#! /usr/bin/env python3

import argparse
import docker
import sys
import re

def get_zones(container):
    out = container.exec_run(['pdnsutil', 'list-all-zones'])
    if out is not None:
        out = out.decode('UTF-8') # bytres to string
        zones = out.split('\n') # convert to list, one zone per line
        zones = list(filter(None, zones)) # remove any empty lines
        zones = zones[0:len(zones) - 1] # remove last line 'All zonecount: xx'
        return zones
    return None

def add_zone(container, zone):
    out = container.exec_run(['pdnsutil', 'create-zone', zone])
    if out is not None:
        out = out.decode('UTF-8')
        if out != 'Creating empty zone \'%s\'\n' % (zone): # might to strict, but as we can not get return code
            return False
        out = container.exec_run(['pdnsutil', 'set-kind', zone, 'master'])
        out = out.decode('UTF-8')
        return out == '';
    return False

def get_tsig_keys(container):
    out = container.exec_run(['pdnsutil', 'list-tsig-keys'])
    if out is not None:
        out = out.decode('UTF-8') # bytres to string
        keys = out.split('\n') # convert to list, one key per line
        keys = list(filter(None, keys)) # remove any empty lines
        keys = {"%s|%s" % (key[0], key[1]): key for key in [key.split(' ') for key in keys]}
        return keys
    return None

def add_tsig_key(container,  key_name, key_algo):
    out = container.exec_run(['pdnsutil', 'generate-tsig-key', key_name, key_algo])
    if out is not None:
        out = out.decode('UTF-8')
        match = re.search('Create new TSIG key %s %s (.+)' % (re.escape(key_name), re.escape(key_algo)), out)
        if match:
            key = match.group(1)
            return key
        else:
            return None
    return None

def add_tsig_key_to_zone(container, zone, key_name):
    out = container.exec_run(['pdnsutil', 'set-meta', zone, 'TSIG-ALLOW-DNSUPDATE', key_name])
    if out is not None:
        out = out.decode('UTF-8')
        return out == 'Enabled TSIG key %s for %s\n' % (key_name, zone) # might to strict, but as we can not get return code
    return False

def main():
    parser = argparse.ArgumentParser(description='Init Powerdns server')
    parser.add_argument('--dhcp-zone', dest='dhcp_zone', default=None, help='DHCP zone to create (eg: dhcp.city-a.example.com)')
    parser.add_argument('--site-zone', dest='site_zone', default=None, help='Site zone to create (eg: city-a.example.com)')
    parser.add_argument('--rev-zone', dest='rev_zone', default=None, help='Reverse zone to create (eg: 100.15.10), arpa suffix will be added automatically')
    parser.add_argument('--tsig-key-name', dest='tsig_key_name', default='dhcpupdate', help='TSIG key name to add (eg: dnsupdate)')
    parser.add_argument('--tsig-key-algo', dest='tsig_key_algo', default='hmac-md5', help='TSIG key algorithm to use (eg: hmac-md5)')
    parser.add_argument('--container-pdns', dest='container_pdns', default='pdns-pdns', help='Name of pdns auth docker container (eg: pdns-pdns)')

    args = parser.parse_args()

    if args.dhcp_zone is None and args.site_zone is None:
        print('At least one of dhcp-zone or site-zone is required')
        return 64

    if args.rev_zone is not None and not args.rev_zone.endswith('.in-addr.arpa'):
        args.rev_zone = args.rev_zone + '.in-addr.arpa'

    print('Getting docker client instance...')
    try:
        client = docker.DockerClient(version='auto')
        client.ping()
    except Exception as e:
        print('Failed to ping docker server: %s' % (e))
        return 1

    print('Getting docker container %s...' % (args.container_pdns))
    try:
        container = client.containers.get(args.container_pdns)
    except Exception as e:
        print('Can not find container %s' % (e))
        return 2

    if container.status != 'running':
        print('Container is not running: %s' % (container.status))
        return 3

    print('Checking existing zones...')
    zones = get_zones(container)
    if zones is None:
        print('Failes to get existing zones: Unknown error')
        return 4

    print('Checking existing TSIG keys...')
    keys = get_tsig_keys(container)
    if keys is None:
        print('Failes to get existing TSIG keys: Unknown error')
        return 5

    if '%s.|%s.' % (args.tsig_key_name, args.tsig_key_algo) in keys:
        print('Not adding TSIG key %s %s: already exists' % (args.tsig_key_name, args.tsig_key_algo))
        tsig_key_secret = keys['%s.|%s.' % (args.tsig_key_name, args.tsig_key_algo)][2]
    else:
        print('Adding TSIG key %s %s...' % (args.tsig_key_name, args.tsig_key_algo))
        tsig_key_secret = add_tsig_key(container, args.tsig_key_name, args.tsig_key_algo)
        if tsig_key_secret is None:
            print('Failed to add TSIG key: Unknown error')
            return 6


    for zone in [args.dhcp_zone, args.site_zone, args.rev_zone]:
        if zone is not None:
            if zone in zones:
                print('Not adding zone %s: already exists' % (zone))
            else:
                print('Adding zone...')
                if not add_zone(container, zone):
                    print('Failed to add zone: Unknown error')
                    return 7
            print('Adding TSIG %s to zone %s...' % (args.tsig_key_name, zone))
            if not add_tsig_key_to_zone(container, zone, args.tsig_key_name):
                print('Failed to add TSIG key to zone: Unknown error')

    print('DHCP_TSIG_KEY_NAME=%s' % (args.tsig_key_name))
    print('DHCP_TSIG_KEY_ALGO=%s' % (args.tsig_key_algo))
    print('DHCP_TSIG_KEY_SECRET=%s' % (tsig_key_secret))

    return 0

if __name__ == '__main__':
    sys.exit(main())