make_tls_certs.py

"""Make the custom certificate and private key files used by TLS tests.

Code heavily borrowed from Lib/tests/make_ssl_certs.py in CPython.
"""

import os
import shutil
import tempfile
import subprocess

req_template = """
    [req]
    distinguished_name     = req_distinguished_name
    x509_extensions        = req_x509_extensions
    prompt                 = no

    [req_distinguished_name]
    C                      = XY
    L                      = Dask-distributed
    O                      = Dask
    CN                     = {hostname}

    [req_x509_extensions]
    subjectAltName         = @san

    [san]
    DNS.1 = {hostname}

    [ca]
    default_ca      = CA_default

    [CA_default]
    dir = cadir
    database  = $dir/index.txt
    crlnumber = $dir/crl.txt
    default_md = sha256
    default_days = 3600
    default_crl_days = 3600
    certificate = tls-ca-cert.pem
    private_key = tls-ca-key.pem
    serial    = $dir/serial
    RANDFILE  = $dir/.rand

    policy          = policy_match

    [policy_match]
    countryName             = match
    stateOrProvinceName     = optional
    organizationName        = match
    organizationalUnitName  = optional
    commonName              = supplied
    emailAddress            = optional

    [policy_anything]
    countryName   = optional
    stateOrProvinceName = optional
    localityName    = optional
    organizationName  = optional
    organizationalUnitName  = optional
    commonName    = supplied
    emailAddress    = optional

    [v3_ca]
    subjectKeyIdentifier=hash
    authorityKeyIdentifier=keyid:always,issuer
    basicConstraints = CA:true
    """

here = os.path.abspath(os.path.dirname(__file__))


def make_cert_key(hostname, sign=False):
    print("creating cert for " + hostname)
    tempnames = []
    for i in range(3):
        with tempfile.NamedTemporaryFile(delete=False) as f:
            tempnames.append(f.name)
    req_file, cert_file, key_file = tempnames
    try:
        req = req_template.format(hostname=hostname)
        with open(req_file, "w") as f:
            f.write(req)
        args = [
            "req",
            "-new",
            "-days",
            "3650",
            "-nodes",
            "-newkey",
            "rsa:2048",
            "-keyout",
            key_file,
            "-config",
            req_file,
        ]
        if sign:
            with tempfile.NamedTemporaryFile(delete=False) as f:
                tempnames.append(f.name)
                reqfile = f.name
            args += ["-out", reqfile]

        else:
            args += ["-x509", "-out", cert_file]
        subprocess.check_call(["openssl"] + args)

        if sign:
            args = [
                "ca",
                "-config",
                req_file,
                "-out",
                cert_file,
                "-outdir",
                "cadir",
                "-policy",
                "policy_anything",
                "-batch",
                "-infiles",
                reqfile,
            ]
            subprocess.check_call(["openssl"] + args)

        with open(cert_file, "r") as f:
            cert = f.read()
        with open(key_file, "r") as f:
            key = f.read()
        return cert, key
    finally:
        for name in tempnames:
            os.remove(name)


TMP_CADIR = "cadir"


def unmake_ca():
    shutil.rmtree(TMP_CADIR)


def make_ca():
    os.mkdir(TMP_CADIR)
    with open(os.path.join("cadir", "index.txt"), "a+") as f:
        pass  # empty file
    # with open(os.path.join('cadir','crl.txt'),'a+') as f:
    # f.write("00")
    with open(os.path.join("cadir", "index.txt.attr"), "w+") as f:
        f.write("unique_subject = no")

    with tempfile.NamedTemporaryFile("w") as t:
        t.write(req_template.format(hostname="our-ca-server"))
        t.flush()
        with tempfile.NamedTemporaryFile() as f:
            args = [
                "req",
                "-new",
                "-days",
                "3650",
                "-extensions",
                "v3_ca",
                "-nodes",
                "-newkey",
                "rsa:2048",
                "-keyout",
                "tls-ca-key.pem",
                "-out",
                f.name,
                "-subj",
                "/C=XY/L=Dask-distributed/O=Dask CA/CN=our-ca-server",
            ]
            subprocess.check_call(["openssl"] + args)
            args = [
                "ca",
                "-config",
                t.name,
                "-create_serial",
                "-out",
                "tls-ca-cert.pem",
                "-batch",
                "-outdir",
                TMP_CADIR,
                "-keyfile",
                "tls-ca-key.pem",
                "-days",
                "3650",
                "-selfsign",
                "-extensions",
                "v3_ca",
                "-infiles",
                f.name,
            ]
            subprocess.check_call(["openssl"] + args)
            # args = ['ca', '-config', t.name, '-gencrl', '-out', 'revocation.crl']
            # subprocess.check_call(['openssl'] + args)


if __name__ == "__main__":
    os.chdir(here)
    cert, key = make_cert_key("localhost")
    with open("tls-self-signed-cert.pem", "w") as f:
        f.write(cert)
    with open("tls-self-signed-key.pem", "w") as f:
        f.write(key)

    # For certificate matching tests
    make_ca()
    with open("tls-ca-cert.pem", "r") as f:
        ca_cert = f.read()

    cert, key = make_cert_key("localhost", sign=True)
    with open("tls-cert.pem", "w") as f:
        f.write(cert)
    with open("tls-cert-chain.pem", "w") as f:
        f.write(cert)
        f.write(ca_cert)
    with open("tls-key.pem", "w") as f:
        f.write(key)
    with open("tls-key-cert.pem", "w") as f:
        f.write(key)
        f.write(cert)

    unmake_ca()