Virtual Partitioning

#!/usr/bin/python3
"""
This example script adds segregation mappings for TLS certs as defined in
segregation_map
"""
from pymongo import MongoClient
from pymongo.collation import Collation
from cryptography import x509
from os import path

config = {
    "MONGO_HOST": "mongo1",
    "MONGO_PORT": 30001,
    "SEGREGATION_DB": "segregation_db",
    "SEGREGATION_COLLECTION": "segregations",
}

# When no label is provided, the client belongs to the common virtual partition
segregation_map = {
    # client certificate : partition label
    "client1.crt": "Label_A",
    "client2.crt": "Label_A",
    "client3.crt": "Label_B",
    "client4.crt": "Label_C",
    "client5.crt": None,
    "client6.crt": None
}

try:
    # Create the Segregation database and collection
    with MongoClient(host=config["MONGO_HOST"], port=config["MONGO_PORT"]) as mongoClient:
        mongoClient[config["SEGREGATION_DB"]
                    ][config["SEGREGATION_COLLECTION"]].drop()
        # The index of Subject and Issuer is case insensitive
        collection = mongoClient[config["SEGREGATION_DB"]].create_collection(
            config["SEGREGATION_COLLECTION"], collation=Collation(locale='en', strength=2))
        collection.create_index([("issuer", 1), ("subject", 1)], unique=True)

        documents_to_add = []

        # Process certificates
        for certificate in segregation_map:
            filename = path.join(path.dirname(
                path.abspath(__file__)), certificate)
            with open(filename, 'rb') as pfile:
                cert = x509.load_pem_x509_certificate(pfile.read())
                
                subject = []

                subject_serial_number = cert.subject.get_attributes_for_oid(
                    x509.NameOID.SERIAL_NUMBER)
                if subject_serial_number:
                    subject.append(subject_serial_number[0].rfc4514_string())

                subject_common_name = cert.subject.get_attributes_for_oid(
                    x509.NameOID.COMMON_NAME)
                if subject_common_name:
                    subject.append(subject_common_name[0].rfc4514_string())

                subject_organization_unit_name = cert.subject.get_attributes_for_oid(
                    x509.NameOID.ORGANIZATIONAL_UNIT_NAME)
                if subject_organization_unit_name:
                    subject.append(
                        subject_organization_unit_name[0].rfc4514_string())

                subject_organization_name = cert.subject.get_attributes_for_oid(
                    x509.NameOID.ORGANIZATION_NAME)
                if subject_organization_name:
                    subject.append(
                        subject_organization_name[0].rfc4514_string())

                subject_postal_code = cert.subject.get_attributes_for_oid(
                    x509.NameOID.POSTAL_CODE)
                if subject_postal_code:
                    subject.append(subject_postal_code[0].rfc4514_string())

                subject_street_address = cert.subject.get_attributes_for_oid(
                    x509.NameOID.STREET_ADDRESS)
                if subject_street_address:
                    subject.append(subject_street_address[0].rfc4514_string())

                subject_locality_name = cert.subject.get_attributes_for_oid(
                    x509.NameOID.LOCALITY_NAME)
                if subject_locality_name:
                    subject.append(subject_locality_name[0].rfc4514_string())

                subject_province_name = cert.subject.get_attributes_for_oid(
                    x509.NameOID.STATE_OR_PROVINCE_NAME)
                if subject_province_name:
                    subject.append(subject_province_name[0].rfc4514_string())

                subject_country_name = cert.subject.get_attributes_for_oid(
                    x509.NameOID.COUNTRY_NAME)
                if subject_country_name:
                    subject.append(subject_country_name[0].rfc4514_string())

                subject = ','.join(subject)

                issuer = []

                issuer_serial_number = cert.issuer.get_attributes_for_oid(
                    x509.NameOID.SERIAL_NUMBER)
                if issuer_serial_number:
                    issuer.append(issuer_serial_number[0].rfc4514_string())

                issuer_common_name = cert.issuer.get_attributes_for_oid(
                    x509.NameOID.COMMON_NAME)
                if issuer_common_name:
                    issuer.append(issuer_common_name[0].rfc4514_string())

                issuer_organization_unit_name = cert.issuer.get_attributes_for_oid(
                    x509.NameOID.ORGANIZATIONAL_UNIT_NAME)
                if issuer_organization_unit_name:
                    issuer.append(
                        issuer_organization_unit_name[0].rfc4514_string())

                issuer_organization_name = cert.issuer.get_attributes_for_oid(
                    x509.NameOID.ORGANIZATION_NAME)
                if issuer_organization_name:
                    issuer.append(issuer_organization_name[0].rfc4514_string())

                issuer_postal_code = cert.issuer.get_attributes_for_oid(
                    x509.NameOID.POSTAL_CODE)
                if issuer_postal_code:
                    issuer.append(issuer_postal_code[0].rfc4514_string())

                issuer_street_address = cert.issuer.get_attributes_for_oid(
                    x509.NameOID.STREET_ADDRESS)
                if issuer_street_address:
                    issuer.append(issuer_street_address[0].rfc4514_string())

                issuer_locality_name = cert.issuer.get_attributes_for_oid(
                    x509.NameOID.LOCALITY_NAME)
                if issuer_locality_name:
                    issuer.append(issuer_locality_name[0].rfc4514_string())

                issuer_province_name = cert.issuer.get_attributes_for_oid(
                    x509.NameOID.STATE_OR_PROVINCE_NAME)
                if issuer_province_name:
                    issuer.append(issuer_province_name[0].rfc4514_string())

                issuer_country_name = cert.issuer.get_attributes_for_oid(
                    x509.NameOID.COUNTRY_NAME)
                if issuer_country_name:
                    issuer.append(issuer_country_name[0].rfc4514_string())

                issuer = ','.join(issuer)

            record = {}
            record["subject"] = subject
            record["issuer"] = issuer
            if segregation_map[certificate]:
                record["label"] = segregation_map[certificate]
            else:
                # No label value indicates the common partition
                pass

            documents_to_add.append(record)

        collection.insert_many(documents_to_add)

        # Reporting
        for i in collection.find({}):
            # display 'null' when no label
            if 'label' not in i:
                i['label'] = 'null'
            print("Added label: '{0}' subject: '{1}', issuer: '{2}'".format(
                i['label'], i['subject'], i['issuer']))

except Exception as ex:
    print(ex)
    import sys
    sys.exit(1)