#!/usr/bin/env python3

import argparse
import datetime
import os
import subprocess
import sys
import re


def increment_serial(line):
    current_serial = re.search(R"\b\d\d*\b", line).group(0)

    current = int(current_serial)
    current_num = current % 100
    current_date = (current - current_num) / 100
    new_date = int(datetime.datetime.now().strftime("%Y%m%d"))
    if current_date == new_date:
        next_num = current_num + 1
    else:
        next_num = 0

    if next_num >= 100:
        raise ValueError("Too many serial changes today!")
    new_serial = str(new_date) + str(next_num).zfill(2)
    line = line.replace(current_serial, new_serial)

    return line

def replace_ip(line):
    source_ip, source_port, dest_ip, dest_port = os.environ["SSH_CONNECTION"].split()
    line = re.sub(R"\b\d\d?\d?\.\d\d?\d?\.\d\d?\d?\.\d\d?\d?\b", source_ip, line)
    return line

def update_dyndns(zonefile, dnslabel):
    out = []
    with open(zonefile, encoding="utf-8") as f:
        for line in f:
            if line.find("Serial") >= 0:
                line = increment_serial(line)
            elif line.find("DYNDNS") >= 0 and line.find(dnslabel) >= 0:
                line = replace_ip(line)
            out.append(line)

    with open(zonefile, "w", encoding="utf-8") as f:
        f.writelines(out)

    reload_bind()


def read_zonefile(zonefile):
    with open(zonefile, encoding="utf-8") as f:
        sys.stdout.write(f.read())


def write_zonefile(zonefile):
    data = sys.stdin.read()
    if not data.strip().endswith("; END"):
        print("Missing end of file marker -- ; END")
        sys.exit(1)

    with open(zonefile, "w", encoding="utf-8") as f:
        f.write(data)

    reload_bind()

def tail_logs():
    subprocess.check_call(["sudo", "journalctl", "-u", "nsd", "-f"])

def reload_bind():
    subprocess.check_call(["sudo", "systemctl", "reload", "nsd"])
    print("nsd reloaded")


def parse_ssh_args():
    parser = argparse.ArgumentParser(description="Edit zone files")

    parser.add_argument("--zonefile",
        help="the zone file to operate on")


    action_group = parser.add_mutually_exclusive_group(required=True)

    action_group.add_argument("--logs", action="store_true",
        help="show bind logs")

    action_group.add_argument("--dyndns",
        help="update the specified dnslabel with the SSH origin IP")

    action_group.add_argument("--read", action="store_true",
        help="print the zone file to stdout")

    action_group.add_argument("--write", action="store_true",
        help="save the contents of stdin to the zone file")

    ssh_args = os.environ.get("SSH_ORIGINAL_COMMAND", "--help").split()[1:]
    args = parser.parse_args(ssh_args)

    if not args.logs and not args.zonefile:
        print("Required parameter: --zonefile\n")
        parser.print_help()
        sys.exit(0)

    return args

def parse_cmdline_args():
    parser = argparse.ArgumentParser(description="Edit zone files")

    parser.add_argument("--allow-zonefile", required=True, action="append",
        help="specify allowed zone files")

    parser.add_argument("--allow-write", action="store_true",
        help="allow --write option")

    parser.add_argument("--allow-dyndns",
        help="allow --dyndns option")

    return parser.parse_args()

def parse_args():
    cmdline_args = parse_cmdline_args()
    ssh_args = parse_ssh_args()

    if ssh_args.zonefile and ssh_args.zonefile not in cmdline_args.allow_zonefile:
        print("The specified zonefile is not on the allowed list:", cmdline_args.allow_zonefile)
        sys.exit(1)

    if ssh_args.dyndns and ssh_args.dyndns != cmdline_args.allow_dyndns:
        print("Dynamic DNS is only allowed for:", cmdline_args.allow_dyndns)
        sys.exit(1)

    if ssh_args.write and not cmdline_args.allow_write:
        print("Write to zonefile is not allowed")
        sys.exit(1)

    return ssh_args


def main():
    args = parse_args()

    if args.logs:
        tail_logs()
    elif args.dyndns:
        update_dyndns(args.zonefile, args.dyndns)
    elif args.read:
        read_zonefile(args.zonefile)
    elif args.write:
        write_zonefile(args.zonefile)
    else:
        assert False, "Bad action"

if __name__ == "__main__":
    main()