#!/usr/bin/env python
import argparse, atexit, os, subprocess, sqlite3, sys, xmlrpclib
from OpenSSL import crypto
from re6st import utils

def create(path, text=None, mode=0666):
    fd = os.open(path, os.O_CREAT | os.O_WRONLY | os.O_TRUNC, mode)
    if text is None:
        return fd
    try:
        os.write(fd, text)
    finally:
        os.close(fd)

def main():
    parser = argparse.ArgumentParser(
        description="Setup script for re6stnet.",
        formatter_class=utils.HelpFormatter)
    _ = parser.add_argument
    _('--registry', required=True, metavar='URL',
        help="HTTP URL of the server delivering certificates.")
    _('--ca-only', action='store_true',
        help='Only fetch CA from registry and exit.')
    _('-d', '--dir',
        help="Directory where the key and certificate will be stored.")
    _('-r', '--req', nargs=2, action='append', metavar=('KEY', 'VALUE'),
        help="The registry only sets the Common Name of your certificate,"
             " which actually encodes your allocated prefix in the network."
             " You can repeat this option to add any field you want to its"
             " subject.")
    _('--email',
        help="Email address where your token is sent. Use -r option if you"
             " want to show an email in your certificate.")
    _('--token', help="The token you received.")
    config = parser.parse_args()
    if config.dir:
        os.chdir(config.dir)
    ca_path = 'ca.crt'
    cert_path = 'cert.crt'
    key_path = 'cert.key'
    dh_path = 'dh2048.pem'

    # Establish connection with server
    s = xmlrpclib.ServerProxy(config.registry)

    # Get CA
    create(ca_path, s.getCa())

    if config.ca_only:
        sys.exit()

    # Generating dh file
    if not os.access(dh_path, os.F_OK):
        r = subprocess.call(('openssl', 'dhparam', '-out', dh_path, '2048'))
        if r:
            sys.exit(r)

    req = crypto.X509Req()
    subj = req.get_subject()
    if config.req:
        for k, v in config.req:
            if k == 'CN':
                sys.exit("CN field is reserved.")
            setattr(subj, k, v)

    cert_fd = token_advice = None
    try:
        if not config.token:
            if not config.email:
                config.email = raw_input('Please enter your email address: ')
            s.requestToken(config.email)
            token_advice = "Use --token to retry without asking a new token\n"
            config.token = raw_input('Please enter your token: ')

        # First make sure we can store private key and open certificate file
        # for writing, to avoid using our token for nothing.
        print "Generating key and cert requests ..."
        cert_fd = create(cert_path)
        pkey = crypto.PKey()
        pkey.generate_key(crypto.TYPE_RSA, 2048)
        key = crypto.dump_privatekey(crypto.FILETYPE_PEM, pkey)

        req.set_pubkey(pkey)
        req.sign(pkey, 'sha1')
        req = crypto.dump_certificate_request(crypto.FILETYPE_PEM, req)

        create(key_path, key, 0600)
        print "Requesting certificate ..."
        cert = s.requestCertificate(config.token, req)
        if not cert:
            token_advice = None
            sys.exit("Error: invalid or expired token")
    except:
        if cert_fd:
            os.remove(cert_path)
        if token_advice:
            atexit.register(sys.stdout.write, token_advice)
        raise
    os.write(cert_fd, cert)
    os.close(cert_fd)

    print "Certificate setup complete."

    cn = utils.subnetFromCert(cert_path)
    subnet = utils.networkFromCa(ca_path) + utils.binFromSubnet(cn)
    print "Your subnet: %s/%u (CN=%s)" \
        % (utils.ipFromBin(subnet), len(subnet), cn)

if __name__ == "__main__":
    main()