'Certificate generation utilities'

__author__ = 'phoenix'
__version__ = '0.1'

import os
import time
import OpenSSL.crypto as crypto
import argparse
import sys


# Temp list for generating certs
workingList = set()


def create_CA(file):
    '''
    Generate the CA root certificate.
    '''
    key = crypto.PKey()
    key.generate_key(crypto.TYPE_RSA, 2048)
    ca = crypto.X509()
    ca.set_serial_number(0)
    # Value 2 means v3
    ca.set_version(2)
    subj = ca.get_subject()
    subj.countryName = 'CN'
    subj.organizationName = 'Privoxy'
    subj.organizationalUnitName = 'pyOpenSSL'
    subj.commonName = 'Privoxy Fake CA'
    ca.gmtime_adj_notBefore(0)
    ca.gmtime_adj_notAfter(24 * 60 * 60 * 3652)
    ca.set_issuer(ca.get_subject())
    ca.set_pubkey(key)
    ca.add_extensions([
        crypto.X509Extension(b'basicConstraints', True, b'CA:TRUE'),
        crypto.X509Extension(b'nsCertType', False, b'sslCA'),
        crypto.X509Extension(b'extendedKeyUsage', True,
                             b'serverAuth,clientAuth,emailProtection,'
                             b'timeStamping,msCodeInd,msCodeCom,msCTLSign'
                             b',msSGC,msEFS,nsSGC'),
        crypto.X509Extension(b'keyUsage', False, b'keyCertSign, cRLSign'),
        crypto.X509Extension(b'subjectKeyIdentifier', False, b'hash',
                             subject=ca)])
    ca.sign(key, 'sha256')
    file.write(crypto.dump_certificate(crypto.FILETYPE_PEM, ca).decode())
    file.write(crypto.dump_privatekey(crypto.FILETYPE_PEM, key).decode())


def get_cert(name, config):
    '''
    Return cert file path. Create it if it doesn't exist.

    Reads from config:
    cafile: the CA file to create dummpy cert files
    certdir: the path where cert files are looked for or created
    '''
    certfile = os.path.join(config.certdir, name + '.crt')
    if not os.path.exists(certfile):
        dummy_cert(config.ca, certfile, name)
    return certfile


def dummy_cert(cafile, certfile, commonname):
    '''
    Generates and writes a certificate to certfile
    commonname: Common name for the generated certificate
    Ref: https://github.com/mitmproxy/netlib/blob/master/netlib/certutils.py
    '''
    if certfile in workingList:
        # Another thread is working on it, wait until it finish
        while True:
            time.sleep(0.2)
            if certfile not in workingList:
                break
    else:
        workingList.add(certfile)
        with open(cafile, 'rb') as file:
            content = file.read()
            ca = crypto.load_certificate(crypto.FILETYPE_PEM, content)
            key = crypto.load_privatekey(crypto.FILETYPE_PEM, content)
        cert = crypto.X509()
        # Value 2 means v3
        cert.set_version(2)
        cert.gmtime_adj_notBefore(0)
        cert.gmtime_adj_notAfter(60 * 60 * 24 * 3652)
        cert.set_issuer(ca.get_subject())
        if commonname.startswith('.'):
            domain = '*' + commonname
        else:
            domain = commonname
        cert.get_subject().CN = domain
        cert.set_serial_number(int(time.time()*10000))
        cert.set_pubkey(ca.get_pubkey())
        cert.add_extensions([
            crypto.X509Extension(b'subjectAltName', False,
                                 str.encode('DNS:'+domain))])
        cert.sign(key, 'sha256')
        with open(certfile, 'wb') as fp:
            fp.write(crypto.dump_certificate(crypto.FILETYPE_PEM, cert))
            fp.write(crypto.dump_privatekey(crypto.FILETYPE_PEM, key))
        workingList.remove(certfile)


def main():
    parser = argparse.ArgumentParser(description='Generates CA Certificates.')
    parser.add_argument('-f', '--file', type=argparse.FileType('w'),
                        default=sys.stdout,
                        help='CA certificate output file')
    args = parser.parse_args()
    create_CA(args.file)


if __name__ == '__main__':
    main()