#!/usr/bin/env python3
# setreverse -- set up rDNS PTR records for a hostname using DNS UPDATE
import argparse
import dns.name
import dns.resolver
import ipaddress
import logging
import sys
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 lookup_ptr(name):
    try:
        ans = dns.resolver.resolve(name, "PTR")
    except dns.resolver.NoAnswer:
        return None
    except dns.resolver.NXDOMAIN:
        return None
    else:
        return ans.rrset[0].target

def chase_one_cname(name):
    try:
        ans = dns.resolver.resolve(name, "CNAME")
    except dns.resolver.NoAnswer:
        return name
    except dns.resolver.NXDOMAIN:
        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()

# dnspython 2.0 in Debian lacks Answer.chaining_result
def compat_chaining_result_cnames(answer):
    if hasattr(answer, "chaining_result"):
        return answer.chaining_result.cnames

    resp = answer.response
    names = []
    count = 0
    qname = answer.qname
    while count < 16:
        try:
            rrset = resp.find_rrset(resp.answer,
                                    qname,
                                    resp.question[0].rdclass,
                                    resp.question[0].rdtype)
            break
        except KeyError:
            if resp.question[0].rdtype != dns.rdatatype.CNAME:
                try:
                    rrset = resp.find_rrset(resp.answer,
                                            qname,
                                            resp.question[0].rdclass,
                                            dns.rdatatype.CNAME)
                    names.append(rrset)
                    count += 1
                    for rr in rrset:
                        qname = rr.target
                        break
                except KeyError:
                    break
            else:
                break
    return names

def lookup_addrs(name, *, allow_cname=False,
                          only_a=False,
                          only_aaaa=False):
    addrs = []
    for rrtype in ["A", "AAAA"]:
        if (only_a and rrtype != "A") or (only_aaaa and rrtype != "AAAA"):
            continue
        try:
            ans = dns.resolver.resolve(name, rrtype)
        except dns.resolver.NoAnswer:
            continue
        addrs += [rr.address for rr in ans.rrset]
        #if cnames := ans.chaining_result.cnames: # needs v2.1.0+
        if cnames := compat_chaining_result_cnames(ans):
            print(f"warning: \"{name}\" is an alias to \"{ans.canonical_name}\"",
                  file=sys.stderr)
            if not allow_cname:
                exit("error: Refusing to point rDNS to an alias; use --allow-cname if needed")
    return addrs

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("--allow-cname", action="store_true",
                        help="allow the hostname to be an alias (a CNAME)")
parser.add_argument("-r", "--remove", action="store_true",
                        help="unset rDNS pointers")
parser.add_argument("-4", "--only-v4", action="store_true",
                        help="only update IPv4 rDNS")
parser.add_argument("-6", "--only-v6", action="store_true",
                        help="only update IPv6 rDNS")
parser.add_argument("-R", "--rpz", action="store_true",
                        help="update the Response Policy Zone")
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("host", nargs="+",
                        help="domain name to update PTR for")
args = parser.parse_args()

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

if args.only_v4 and args.only_v6:
    exit("error: Options -4 and -6 are mutually exclusive")

failed = False

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

    addresses = lookup_addrs(host, allow_cname=args.allow_cname,
                                   only_a=args.only_v4,
                                   only_aaaa=args.only_v6)
    print(f"Host <{host}> has addresses:", *sorted(addresses))

    for addr in addresses:
        addr = ipaddress.ip_address(addr)
        rname = dns.name.from_text(addr.reverse_pointer)
        if args.rpz:
            rname = rname.relativize(dns.name.root) + RPZ_ZONE
        old = lookup_ptr(rname)

        # Detect IPv4 classless delegation, as well as other uses of CNAME.
        rname = chase_one_cname(rname)

        zone = dns.resolver.zone_for_name(rname)
        logging.debug(f"Updating zone <{zone}>")

        if args.remove:
            if old:
                print(f"Removing PTR for [{addr}] (was \"{old}\")")
            else:
                print(f"Removing PTR for [{addr}]")
            cmds = [f"del {rname} {args.ttl} PTR\n"]
        else:
            if old:
                print(f"Changing PTR for [{addr}] to \"{host}\" (was \"{old}\")")
            else:
                print(f"Changing PTR for [{addr}] to \"{host}\"")
            cmds = [f"del {rname} {args.ttl} PTR\n",
                    f"add {rname} {args.ttl} PTR {host}\n"]

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

        if not args.dry_run:
            try:
                send_nsupdate(zone, cmds,
                              use_gss=(not args.no_gss),
                              debug=(args.debug >= 2))
            except subprocess.CalledProcessError as e:
                print("error: %s" % e, file=sys.stderr)
                failed = True

exit(failed)
