#!/usr/bin/env python3
# vireverse -- interactively edit reverse-DNS records
import argparse
import dns.exception
import dns.name
import dns.resolver
import dns.reversename
import dns.zone
import gssapi
import ipaddress
import logging
import math
import os
import socket
import subprocess
import tempfile

from n.dns.gssapi import gss_tsig_negotiate

RPZ_ZONE = dns.name.from_text("rpz.nullroute.lt.")
RPZ_MASTER = dns.name.from_text("star.nullroute.lt.")
DEFAULT_NET = "10.147.0.0/16"
DEFAULT_NET6 = "2a06:e881:108::/48"
NET_GROUPS = [
    (ipaddress.ip_network("10.147.240.0/22"),           30),
    (ipaddress.ip_network("2a06:e881:108:ffff::/64"),   126),
]

def confirm(text):
    return input(f"{text} ").startswith("y")

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 prefix_for_zone(name):
    # Similar to dns.reversename.to_address() but handles prefixes.
    assert isinstance(name, dns.name.Name)
    if name.is_subdomain(dns.reversename.ipv4_reverse_domain):
        name = name.relativize(dns.reversename.ipv4_reverse_domain)
        octets = [x.decode() for x in reversed(name.labels)]
        # Handle sub-/24 CIDR zones
        if len(octets) == 4 and "/" in octets[-1]:
            octets[-1], bits = octets[-1].split("/", 1)
            bits = int(bits)
        else:
            bits = len(octets) * 8
            octets += ["0"] * (4 - len(octets))
        addr = ".".join(octets)
        prefix = ipaddress.ip_network((addr, bits))
        return prefix
    elif name.is_subdomain(dns.reversename.ipv6_reverse_domain):
        name = name.relativize(dns.reversename.ipv6_reverse_domain)
        nibbles = [x.decode() for x in reversed(name.labels)]
        bits = len(nibbles) * 4
        nibbles += ["0"] * (32 - len(nibbles))
        parts = []
        for i in range(0, len(nibbles), 4):
            parts.append("".join(nibbles[i : i+4]))
        addr = ":".join(parts)
        prefix = ipaddress.ip_network((addr, bits))
        return prefix
    else:
        # XXX: For unusual zones, it might be wise to return 'none' here so
        # that prefix checks would be skipped outright.
        raise ValueError(name)

def swap_suffix(name, old_suffix, new_suffix):
    assert name.is_absolute()
    assert old_suffix.is_absolute()
    assert new_suffix.is_absolute()
    return name.relativize(old_suffix) + new_suffix

def address_for_ptr(name, real_suffix, fake_suffix):
    assert isinstance(name, dns.name.Name)
    if real_suffix != fake_suffix:
        name = swap_suffix(name, real_suffix, fake_suffix)
    addr = dns.reversename.to_address(name)
    return ipaddress.ip_address(addr)

def ptr_for_address(addr, real_suffix, fake_suffix):
    assert isinstance(addr, (ipaddress.IPv4Address,
                             ipaddress.IPv6Address))
    name = dns.reversename.from_address(str(addr))
    if real_suffix != fake_suffix:
        name = swap_suffix(name, fake_suffix, real_suffix)
    return name

def master_for_zone(zone):
    ans = dns.resolver.resolve(zone, "SOA")
    return ans.rrset[0].mname

def probe_dns_server(server_name):
    gai = socket.getaddrinfo(str(server_name),
                             "domain",
                             socket.AF_UNSPEC,
                             socket.SOCK_DGRAM)
    for af, sf, pt, cname, sa in gai:
        yield sa[0]

def axfr(zone, server_name, *, use_gss=False):
    for server_addr in probe_dns_server(server_name):
        if use_gss:
            key_ring, key_name = gss_tsig_negotiate(server_name,
                                                    server_addr)
        else:
            key_ring, key_name = None, None
        resp = dns.query.xfr(server_addr, zone,
                             relativize=False,
                             keyring=key_ring,
                             keyname=key_name)
        return dns.zone.from_xfr(resp,
                                 relativize=False)

def run_visual(filename, line=0):
    editor = os.environ.get("VISUAL") \
          or os.environ.get("EDITOR") \
          or "vi"
    if line and os.path.basename(editor) in {"vi", "vim", "nvim"}:
        subprocess.run([editor, filename, "+:setl cursorline", f"+{line}"])
    else:
        subprocess.run([editor, filename])

def case_eq(a, b):
    """
    Compare two sets or rdatasets, paying attention to case changes (as
    dns.name.Name normally compares case-insensitive).
    """
    assert isinstance(a, (set, dns.rdataset.Rdataset))
    if a == b:
        a = {*map(str, a)}
        b = {*map(str, b)}
    return a == b

def addr_diff(a, b):
    if not a or not b:
        return math.inf
    return abs(a._ip - b._ip)
    #a = int.from_bytes(a.packed)
    #b = int.from_bytes(b.packed)
    #return abs(a - b)

def group_for_addr(addr):
    """
    Round the address up to the nearest IPv4 /24 or IPv6 /64 that covers it.
    """
    global NET_GROUPS
    for grp_net, grp_bits in NET_GROUPS:
        if addr in grp_net:
            return ipaddress.ip_interface((addr, grp_bits)).network
    bits = {4: 24, 6: 64}[addr.version]
    return ipaddress.ip_interface((addr, bits)).network

def send_nsupdate(zone, cmds, *, server=None, use_gss=True, debug=False):
    if not server:
        server = master_for_zone(zone)

    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("-l", "--ttl", type=int, default=3600,
                        help="set the TTL for created records")
parser.add_argument("-6", "--ipv6", action="store_true",
                        help="edit the default IPv6 reverse zone")
parser.add_argument("-R", "--rpz", action="store_true",
                        help="edit the Response Policy Zone")
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("network",
                        help="IP address to edit the zone for",
                        nargs="?")
args = parser.parse_args()

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

if (args.ipv6 + args.rpz) > 1:
    exit(f"vireverse: options -6 and -R conflict")

if args.rpz:
    # For RPZ, the prefix is only used to jump to the nearest line.
    # If an interface was given, ignore nbits and treat it as /32 or /128.
    if args.network:
        arg = ipaddress.ip_interface(args.network)
        jump_net = ipaddress.ip_network(arg.ip)
    else:
        jump_net = None
    zone = RPZ_ZONE
    logging.info(f"Editing RPZ zone <{zone}>")
    zone_net = None
    real_suffix = RPZ_ZONE
    fake_suffix = dns.name.root
else:
    if args.network:
        # Deliberately accept interfaces (i.e. prefixes with non-0 host bits)
        # for convenient pasting. Use network/nbits for the initial DNS zone,
        # but address/32 or address/128 as the jump target.
        try:
            arg = ipaddress.ip_interface(args.network)
            jump_net = ipaddress.ip_network(arg.ip)
            rname = dns.name.from_text(arg.network[0].reverse_pointer)
        except ValueError:
            logging.debug(f"Failed to parse {args.network!r} as IP prefix, assuming DNS zone")
            jump_net = None
            rname = dns.name.from_text(args.network)
    else:
        jump_net = ipaddress.ip_network(DEFAULT_NET6 if args.ipv6 else DEFAULT_NET)
        rname = dns.name.from_text(jump_net[0].reverse_pointer)

    logging.debug(f"Target network is [{jump_net}]")
    logging.debug(f"Initial DNS name is <{rname}>")
    cname = chase_one_cname(rname)
    if cname != rname:
        logging.debug(f"Aliases followed to <{cname}>")
    zone = dns.resolver.zone_for_name(cname)
    logging.debug(f"Zone for <{cname}> is <{zone}>")
    zone_net = prefix_for_zone(zone)
    logging.debug(f"IP range for <{zone}> is [{zone_net}]")
    if zone_net != jump_net:
        # Warn that input prefix was widened.
        logging.warning(f"Editing zone <{zone}> for {zone_net}")

    real_suffix = zone
    fake_suffix = zone
    if zone.is_subdomain(dns.reversename.ipv4_reverse_domain):
        if zone_net.prefixlen > 24:
            assert b"/" in zone.labels[0]
            fake_suffix = zone.parent()
            logging.debug(f"Derived classful zone is <{fake_suffix}>")

if args.rpz:
    # XXX: hack (this should be a properly delegated .sym zone)
    master = RPZ_MASTER
    logging.debug(f"Master for zone <{zone}> is <{master}> (hardcoded)")
else:
    master = master_for_zone(zone)
    logging.debug(f"Master for zone <{zone}> is <{master}>")

if args.no_gss:
    logging.debug("Performing unauthenticated zone transfer")
else:
    logging.debug("Using GSS-TSIG authentication for zone transfer")
try:
    zone_data = axfr(zone, master, use_gss=(not args.no_gss))
except dns.xfr.TransferError as ex:
    exit(f"vireverse: zone transfer failed: {ex!r}")
except gssapi.exceptions.GSSError as ex:
    exit(f"vireverse: GSS-TSIG zone transfer failed: {ex!r}")
logging.debug(f"Transferred records for {len(zone_data.keys())} names.")

old = {}
new = {}
is_dnssec = False

for name, node in zone_data.items():
    if rdset := node.get_rdataset(dns.rdataclass.IN, dns.rdatatype.PTR):
        addr = address_for_ptr(name, real_suffix, fake_suffix)
        old[addr] = rdset
        #old[addr] = set(rdset)

old_soa = zone_data[zone].get_rdataset(dns.rdataclass.IN, dns.rdatatype.SOA)
logging.debug(f"Old serial number {old_soa[0].serial}.")

pref_line = 0
near_line = 0
near_addr = None

with tempfile.NamedTemporaryFile("w", delete=False) as temp:
    n_lines = 0
    last_net = None

    print(f"# Editing {zone} (serial {old_soa[0].serial})", file=temp.file)
    n_lines += 1

    # A mixed list of IPv4Address and IPv6Address (such as in a RPZ zone)
    # refuses to be sorted
    old4 = [a for a in old if a.version == 4]
    old6 = [a for a in old if a.version == 6]
    assert len(old) == len(old4 + old6)

    # Maximum width of the address column...
    addr_max_width = max(len(str(addr)) for addr in old) if old else 0
    # ...rounded up to the effective width (whole number of tabs), with an
    # increment to guarantee that if it is already an exact multiple of 8, then
    # a tab is inserted regardless.
    addr_max_width = math.ceil((addr_max_width + 1) / 8) * 8

    for addr in sorted(old4) + sorted(old6):
        # Output a blank line between each /24 or /64
        this_net = group_for_addr(addr)
        if last_net and this_net != last_net:
            print("", file=temp.file)
            n_lines += 1
        last_net = this_net

        # Let `vireverse <addr>` open editor with the cursor at this location
        # (or the nearest line if the address is not in the zone yet).
        if jump_net:
            if addr in jump_net and not pref_line:
                pref_line = n_lines + 1
            elif addr_diff(jump_net[0], addr) < addr_diff(jump_net[0], near_addr):
                near_addr = addr
                near_line = n_lines + 1

        for rdata in old[addr]:
            pad = addr_max_width - len(str(addr))
            tabs = "\t" * (math.ceil(pad / 8) + 0)
            target = rdata.target.relativize(dns.name.root)
            print(f"{addr}{tabs}{target}", file=temp.file)
            n_lines += 1

    print(f"# End of file", file=temp.file)
    edit_path = temp.name

run_visual(edit_path, pref_line or near_line)

parse_err = False
with open(edit_path, "r") as temp:
    for i, line in enumerate(temp, start=1):
        line = line.rstrip()
        if not line:
            continue
        if line[0] in {"#", ";"}:
            continue

        try:
            addr, target = line.split()
            addr = ipaddress.ip_address(addr)
            target = dns.name.from_text(target)
        except (ValueError, dns.exception.SyntaxError) as err:
            logging.error(f"error: {err}")
            logging.error(f"{edit_path}:{i}: {line}")
            parse_err = True
            continue

        if zone_net and addr not in zone_net:
            logging.error(f"error: {addr} is out of range for {zone_net}")
            logging.error(f"{edit_path}:{i}: {line}")
            parse_err = True
            continue

        rdata = dns.rdtypes.ANY.PTR.PTR(dns.rdataclass.IN,
                                        dns.rdatatype.PTR,
                                        target)
        if addr not in new:
            new[addr] = dns.rdataset.Rdataset(dns.rdataclass.IN,
                                              dns.rdatatype.PTR)
            #new[addr] = set()
        new[addr].add(rdata)

if parse_err:
    exit(f"error: Could not parse input (retained in {edit_path!r})")

addrs_missing = {*old} - {*new}
addrs_created = {*new} - {*old}
addrs_updated = {addr
                 for addr in {*old} & {*new}
                 if not case_eq(old[addr], new[addr])}

if not new:
    if not old:
        print("No changes were made (zone remains empty).")
        exit()
    else:
        exit(f"error: Zone became empty.")

if len(addrs_missing) + len(addrs_created) + len(addrs_updated) == 0:
    print("No changes were made.")
    exit()
else:
    fmt = lambda rs: ", ".join(sorted([str(r.target) for r in rs]))
    if addrs_missing:
        print("Deletions:")
        for addr in addrs_missing:
            print(f"\t{addr} -= (was {fmt(old[addr])})")
    if addrs_created:
        print("Additions:")
        for addr in addrs_created:
            print(f"\t{addr} += {fmt(new[addr])}")
    if addrs_updated:
        print("Changes:")
        for addr in addrs_updated:
            print(f"\t{addr} := {fmt(new[addr])} (was {fmt(old[addr])})")
    if not confirm("Continue?"):
        exit("Changes discarded.")

cmds = []

for addr in addrs_missing | addrs_updated:
    rname = ptr_for_address(addr, real_suffix, fake_suffix)
    cmds += [f"update del {rname} {args.ttl} PTR\n"]
    for rdata in old[addr]:
        print(f"Deleting: {addr} -= (was {rdata.target})")

for addr in addrs_created | addrs_updated:
    rname = ptr_for_address(addr, real_suffix, fake_suffix)
    for rdata in new[addr]:
        cmds += [f"update add {rname} {args.ttl} PTR {rdata.target}\n"]
        print(f"Adding:   {addr} += {rdata.target}")

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

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

os.unlink(edit_path)
