#!/usr/bin/env python3
import socket
import random
import sys
import yaml
import base64
from OpenSSL.crypto import *

if len(sys.argv) != 3:
    print(f'Usage: {sys.argv[0]} <username> <output file>')
    sys.exit(1)

username = sys.argv[1]
output = sys.argv[2]
etcdir = '/etc/ajenti'
config_path = f'{etcdir}/aj.yml'

with open(config_path, 'r') as c:
    config = yaml.load(c.read(), Loader=yaml.Loader)

if not config['ssl']['enable']:
    print('SSL is not enabled in config.yml')
    sys.exit(2)

cn = f'{username}@{socket.gethostname()}'
key = PKey()
key.generate_key(TYPE_RSA, 4096)

with open(config['ssl']['certificate'], 'r') as certificate:
    content = certificate.read()
    ca_key = load_privatekey(FILETYPE_PEM, content)
    ca_cert = load_certificate(FILETYPE_PEM, content)

cert = X509()
cert.get_subject().countryName = 'NA'
cert.get_subject().organizationName = socket.gethostname()
cert.get_subject().commonName = cn
cert.set_pubkey(key)
cert.set_serial_number(random.getrandbits(8 * 20))
cert.gmtime_adj_notBefore(0)
cert.gmtime_adj_notAfter(10 * 365 * 24 * 60 * 60)
cert.set_issuer(ca_cert.get_subject())
cert.sign(ca_key, 'sha256')

pkcs = PKCS12()
#pkcs.set_ca_certificates([ca_cert])
pkcs.set_certificate(cert)
pkcs.set_privatekey(key)
pkcs.set_friendlyname(bytes(cn, encoding="utf-8"))

cert_info = {
    'digest': cert.digest('sha256').decode('utf-8'),
    'name': ','.join(b'='.join(x).decode('utf-8')
                        for x in cert.get_subject().get_components()
                    ),
    'serial': str(cert.get_serial_number()),
    'user': username,
}
config['ssl']['client_auth']['certificates'].append(cert_info)
print(cert_info)
with open(config_path, 'w') as f:
    f.write(yaml.dump(config, default_flow_style=False))

with open(output, 'w') as out:
    out.write(base64.b64encode(pkcs.export()).decode('utf-8'))
