#!python

import argparse
import subprocess
import os
import time

def get_args():
    parser = argparse.ArgumentParser(description='a2utils certbot wrapper to generate certificates with remote validation')
    parser.add_argument('--ssh', nargs='+', help='SSH arguments, e.g.  --ssh root@example.com')
    parser.add_argument('-d','--domain', nargs='+', help='Domain name to generate certificate')
    parser.add_argument('--aliases', default=False, action='store_true')
    parser.add_argument('--test-cert', default=False, action='store_true')
    
    return parser.parse_args()



def printenv():
    print(f"PID: {os.getpid()}")
    for k, v in os.environ.items():
        if k.startswith('CERTBOT') or k.startswith('A2UTILS'):
            print(k,v)

def process_aliases(ssh, domain):
    cmd = ['ssh']
    cmd.extend(ssh)
    cmd.extend(['a2vhost', '-d', *domain, '--aliases'])
    out = subprocess.run(cmd, capture_output=True)
    aliases = out.stdout.decode().strip().split(' ')
    return aliases

def certbot_wrapper(ssh, domain, test_cert):
    test_options = ['--test-cert'] if test_cert else list()
    my_env = os.environ
    my_env['A2UTILS_SSH'] = ' '.join(ssh)

    cmd = [
        'certbot', 'certonly', '--manual', *test_options,
        '--manual-auth-hook', __file__,
        '--manual-cleanup-hook', __file__,
    ]

    for d in domain:
        cmd.extend(['-d', d])
    
    subprocess.run(cmd, env=my_env)

def get_role():
    
    if 'CERTBOT_AUTH_OUTPUT' in os.environ:
        return 'cleanup'
    
    if 'CERTBOT_DOMAIN' in os.environ:
        return 'place'
    
    return 'wrapper'

def place():
    domain = os.getenv('CERTBOT_DOMAIN')
    token = os.getenv('CERTBOT_TOKEN')
    validation = os.getenv('CERTBOT_VALIDATION')
    sshcmd = os.getenv('A2UTILS_SSH').split(' ')

    cmd = [ 'ssh', *sshcmd ] + ['a2certbot', '-d', domain, '--place', token, validation ]
    # print("place cmd:", cmd)
    subprocess.run(cmd)
    # we need any output to have

def cleanup():
    domain = os.getenv('CERTBOT_DOMAIN')
    token = os.getenv('CERTBOT_TOKEN')
    sshcmd = os.getenv('A2UTILS_SSH').split(' ')

    cmd = [ 'ssh', *sshcmd ] + ['a2certbot', '-d', domain, '--cleanup', token ]
    # print("cleanup cmd:", cmd)
    # time.sleep(20)
    subprocess.run(cmd)

def main():
    args = get_args()

    if not args.domain:
        print("Need domain(s) (-d)")
        return

    if not args.ssh:
        args.ssh = [ 'ssh', f'root@{args.domain[0]}' ]
        print(f"--ssh not set, using {' '.join(args.ssh)}")

    if args.aliases:
        args.domain = process_aliases(args.ssh, args.domain)

    # print(f"args: {args}")
    # printenv()

    role = get_role()
    if role == 'wrapper':
        certbot_wrapper(args.ssh, args.domain, args.test_cert)
    elif role == 'place':
        place()
    elif role == 'cleanup':
        cleanup()
    

if __name__ == '__main__':
    main()