rug-website/rugwebsite/management/commands/init-saml2-settings.py

from django.core.management.base import BaseCommand, CommandError
import datetime
import os.path
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives import hashes, serialization
from cryptography import x509
from cryptography.x509.oid import NameOID

try:
    import urllib  # python 2
except:
    import urllib.request as urllib  # python 3


class Command(BaseCommand):
    help = 'Create SAML2 settings, they are printed to standard out and it is up to de developer or deployer to add' \
           'them to the django settings.'

    def add_arguments(self, parser):
        parser.add_argument('--country',            required=True, nargs=1, type=str, dest='country', help='For example: US, NL, DE, etc')
        parser.add_argument('--city',               required=True, nargs=1, type=str, dest='city', help='For example: London or Groningen')
        parser.add_argument('--state',              required=True, nargs=1, type=str, dest='state', help='For example: State or province, for example California or Groningen.')
        parser.add_argument('--organisation',       required=True, nargs=1, type=str, dest='organisation', help='Typically \'University of Groningen\'')
        parser.add_argument('--organisation-unit',  required=True, nargs=1, type=str, dest='organisation_unit', help='For example: \'Research and Innovation\' Support or \'Faculty of smart people\'')
        parser.add_argument('--common-name',        required=True, nargs=1, type=str, dest='common_name')
        parser.add_argument('--alternative',        nargs='*', type=str, dest='alternatives')
        parser.add_argument('--support-name',       required=True, nargs=1, type=str, dest='support_name')
        parser.add_argument('--support-email',      required=True, nargs=1, type=str, dest='support_email')
        parser.add_argument('--technical-name',     required=True, nargs=1, type=str, dest='technical_name')
        parser.add_argument('--technical-email',    required=True, nargs=1, type=str, dest='technical_email')
        parser.add_argument('--entity-id',          required=True, nargs=1, type=str, dest='entity_id', help='Used as an identifyer of your service, typically the url to your service: www.rug.nl/yourservice')
        parser.add_argument('--base-url',           required=True, nargs=1, type=str, dest='base_url', help='The base url of your service, the route <base_url>sso/saml/ should exist.')
        parser.add_argument('--expires-after-days', nargs='?', type=int, default=10 * 365, dest='expires')

    def handle(self, *args, **options):
        for option in {'country', 'city', 'state', 'organisation', 'organisation_unit', 'common_name', 'support_name',
                       'support_email', 'technical_name', 'technical_email', 'entity_id', 'base_url'}:
            assert option in options and options[option] is not None and len(options[option]) == 1, "Expected one " \
                                                                                                    "value for option" \
                                                                                                    ": " + option
            options[option] = options[option][0]

        key = rsa.generate_private_key(public_exponent=65537, key_size=2048, backend = default_backend())
        subject = issuer = x509.Name([
            x509.NameAttribute(NameOID.COUNTRY_NAME, options['country']),
            x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, options['state']),
            x509.NameAttribute(NameOID.LOCALITY_NAME, options['city']),
            x509.NameAttribute(NameOID.ORGANIZATION_NAME, options['organisation']),
            x509.NameAttribute(NameOID.ORGANIZATIONAL_UNIT_NAME, options['organisation_unit']),
            x509.NameAttribute(NameOID.COMMON_NAME, options['common_name']),
        ])
        cert = (
            x509.CertificateBuilder()
                .subject_name(subject)
                .issuer_name(issuer)
                .public_key(key.public_key())
                .serial_number(x509.random_serial_number())
                .not_valid_before(datetime.datetime.utcnow())
                .not_valid_after(datetime.datetime.utcnow() + datetime.timedelta(days=options['expires']))
                .add_extension(x509.SubjectAlternativeName([
                    x509.DNSName(alternative)
                    for alternative in options['alternatives']
                ] if options['alternatives'] is not None else []), critical = False)
                .sign(key, hashes.SHA256(), default_backend())
        )
        settings_template = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'settings_template.py')

        key_text = key.private_bytes(encoding=serialization.Encoding.PEM,
                                     format = serialization.PrivateFormat.TraditionalOpenSSL,
                                     encryption_algorithm = serialization.NoEncryption()).decode('utf-8')
        x509_text = cert.public_bytes(encoding=serialization.Encoding.PEM).decode('utf-8')

        assert (
            key_text.startswith('-----BEGIN RSA PRIVATE KEY-----\n') and
            key_text.endswith('\n-----END RSA PRIVATE KEY-----\n')), "Private key starts and ends improperly, should " \
                                                                     "be -----BEGIN RSA PRIVATE KEY----- and " \
                                                                     "-----END RSA PRIVATE KEY-----"
        assert (
            x509_text.startswith('-----BEGIN CERTIFICATE-----\n') and
            x509_text.endswith('\n-----END CERTIFICATE-----\n')), "Certificate starts and ends improperly, should " \
                                                                     "be -----BEGIN CERTIFICATE----- and " \
                                                                     "-----END CERTIFICATE-----"
        key_text = key_text[len('-----BEGIN RSA PRIVATE KEY-----\n'):-len('\n-----END RSA PRIVATE KEY-----\n')]
        x509_text = x509_text[len('-----BEGIN CERTIFICATE-----\n'):-len('\n-----END CERTIFICATE-----\n')]


        with open(settings_template, 'r') as f:
            print(f.read().format(private_key=key_text , x509=x509_text, **options))