#!/usr/bin/env python3
# setdns -- change DNS records using DNS UPDATE
import argparse
import dns.name
import dns.rdata
from dns.rdataclass import IN
import dns.rdatatype
import dns.resolver
import ipaddress
import logging
import subprocess

RPZ_ZONE = dns.name.from_text("rpz.nullroute.lt.")

def lookup_suffix(name):
    ans = dns.resolver.resolve(name,
                               search=True,
                               raise_on_no_answer=False)
    return ans.qname

def chase_one_cname(name):
    try:
        ans = dns.resolver.resolve(name, "CNAME")
    except dns.resolver.NoAnswer:
        return name
    else:
        return ans.rrset[0].target

def master_for_zone(zone):
    # Because Bind, the SOA response might have an upper-case tail, e.g.
    # "star.nullroute.LT" for our RPZ zone, and nsupdate blindly gives that to
    # Kerberos which fails. So we manually look up the master server here and
    # casefold the result.
    ans = dns.resolver.resolve(zone, "SOA")
    return ans.rrset[0].mname.canonicalize()

def rtype_for_address(rdata):
    addr = ipaddress.ip_address(rdata)
    return {4: "A", 6: "AAAA"}[addr.version]

def send_nsupdate(zone, cmds, *, server=None, use_gss=True, debug=False):
    if not server:
        server = master_for_zone(zone)
    logging.debug(f"Sending update to server \"{server}\"")

    cmds = [f"server {server}\n",
            f"zone {zone}\n",
            *cmds,
            f"send\n"]

    nsupdate_args = ["nsupdate"]
    if debug:
        nsupdate_args += ["-d"]
    if use_gss:
        nsupdate_args += ["-g"]

    subprocess.run(nsupdate_args,
                   input="".join(cmds).encode(),
                   check=True)

parser = argparse.ArgumentParser()
parser.add_argument("-a", "--add", action="store_true",
                        help="keep existing records of same type")
parser.add_argument("-r", "--remove", action="store_true",
                        help="remove all specified records")
parser.add_argument("-z", "--zone",
                        help="override automatically determined zone")
parser.add_argument("-s", "--server",
                        help="override automatically determined server")
parser.add_argument("-R", "--rpz", action="store_true",
                        help="update the Response Policy Zone")
parser.add_argument("-t", "--type",
                        help="manage records of the specified type")
parser.add_argument("-l", "--ttl", type=int, default=3600,
                        help="set the TTL for created records")
parser.add_argument("-n", "--dry-run", action="store_true",
                        help="only show updates that would be done")
parser.add_argument("-x", "--no-gss", action="store_true",
                        help="disable Kerberos (GSS-TSIG) authentication")
parser.add_argument("-d", "--debug", action="count", default=0,
                        help="enable nsupdate debugging")
parser.add_argument("-v", "--verbose", action="store_true",
                        help="show detailed information")
parser.add_argument("name",
                        help="domain name to update")
parser.add_argument("data", nargs="?",
                        help="record data")
args = parser.parse_args()

logging.basicConfig(level=[logging.INFO, logging.DEBUG][args.verbose],
                    format="%(message)s")

if args.zone and args.rpz:
    exit(f"setptr: --rpz and --zone are mutually exclusive")

if "." in args.name:
    rname = dns.name.from_text(args.name)
else:
    print(f"Looking up \"{args.name}\"...", end=" ", flush=True)
    rname = lookup_suffix(args.name)
    print(f"canonicalized to <{rname}>")

if args.rpz:
    rname = rname.relativize(dns.name.root) + RPZ_ZONE

if args.zone:
    zone = dns.name.from_text(args.zone)
else:
    print(f"Finding zone root...", end=" ", flush=True)
    zone = dns.resolver.zone_for_name(rname)
    print(f"starts at <{zone}>")

if args.remove and args.add:
    exit("error: Options '--add' and '--remove' are mutually exclusive")
elif args.remove and args.data:
    rtype = dns.rdatatype.from_text(args.type or "ANY").name
    if rtype == "ANY":
        exit("error: Cannot specify data with type 'ANY'")
    rdata = dns.rdata.from_text(IN, rtype, args.data)
    print(f"Removing {rtype} record from <{rname}>")
    cmds = [f"del {rname} 0 {rtype} {rdata}\n"]
elif args.remove:
    rtype = dns.rdatatype.from_text(args.type or "ANY").name
    if rtype == "ANY":
        print(f"Removing all records from <{rname}>")
    else:
        print(f"Removing all {rtype} records from <{rname}>")
    cmds = [f"del {rname} 0 {rtype}\n"]
elif args.data:
    if args.type:
        rtype = dns.rdatatype.from_text(args.type).name
    else:
        rtype = rtype_for_address(args.data)
    if rtype == "ANY":
        exit("error: Cannot create record with type 'ANY'")
    rdata = dns.rdata.from_text(IN, rtype, args.data)
    if args.add:
        print(f"Adding {rtype} record at <{rname}>")
        cmds = [f"add {rname} {args.ttl} {rtype} {rdata}\n"]
    else:
        print(f"Changing {rtype} record at <{rname}>")
        cmds = [f"del {rname} 0 {rtype}\n",
                f"add {rname} {args.ttl} {rtype} {rdata}\n"]
else:
    if args.add:
        exit("error: Record data must be specified")
    else:
        exit("error: Either record data or '--remove' must be specified")

if args.debug >= 1:
    print("Commands:")
    print("\t" + "\t".join(cmds), end="")

if not args.dry_run:
    send_nsupdate(zone, cmds,
                  server=args.server,
                  use_gss=(not args.no_gss),
                  debug=(args.debug >= 2))
