#!/usr/bin/env python3

import subprocess
import sys

ENV={"PASSWD": "YOUR_PASSWORD"}  # samba-tool env
USR=["-U", r"REALM\Administrator"]
DOMAIN="ZONE_NAME"


class Record:
    """
    Holds information about a DNS record

    Attributes:
        name (str): name of the record
        type (str): type of the record (A|CNAME|TXT|MX...)
        value (str): value of the record
        attributes (dict): additional attributes (e.g. ttl)
    """
    def __init__(self, name, line):
        """
        Args:
            name (str): name of the record
            line (str): a line of output from "samba-tool dns query"
                        e.g. "A: 10.99.2.22 (flags=f0, serial=8, ttl=900)"
        """
        line = line.strip()
        self.name = name
        self.type = line.split(":")[0]
        self.value = line.split()[1]
        attrib_str = line.split("(")[1].split(")")[0]
        self.attributes = {}
        if "=" in self.value:
            self.value = "."
            attrib_str += " ".join(line.split()[1:]).split("(")[0]
        for pair in [x.split("=") for x in attrib_str.split()]:
            self.attributes[pair[0]] = pair[1]


def get_records():
    """
    Get a list of all records for the domain / zone

    Returns:
        (list(Record)): list of record objects
    """
    records = []
    result = subprocess.run(["samba-tool", "dns", "query", "127.0.0.1", DOMAIN, "@", "ALL"] + USR, env=ENV, capture_output=True, text=True)
    for line in result.stdout.splitlines():
        if line.strip().startswith("Name="):
            name=line.split("Name=")[1].split(",")[0]
        else:
            records.append(Record(name, line))
    return records


def print_dns():
    """
    Print a all records to stdout
    """
    records = get_records()
    for r in records:
        if r.name:
            print(f"{r.type:^8} {r.name:24} {r.value}")


def update_record(rname, rtype, rvalue):
    """
    Update a given record

    Args:
        rname (str): record name
        rtype (str): record type (A|CNAME|TXT...)
        rvalue (str): record value
    """
    records = get_records()
    exrec = next(r for r in records if r.name == rname)
    if not exrec:
        raise ValueError(f"Record {rname} not found")
    result = subprocess.run(["samba-tool", "dns", "update", "127.0.0.1", DOMAIN, rname, rtype, exrec.value, rvalue] + USR, env=ENV)
    result.check_returncode()


def add_record(rname, rtype, rvalue):
    """
    Add a given record

    Args:
        rname (str): record name
        rtype (str): record type (A|CNAME|TXT...)
        rvalue (str): record value
    """
    result = subprocess.run(["samba-tool", "dns", "add", "127.0.0.1", DOMAIN, rname, rtype, rvalue] + USR, env=ENV)
    result.check_returncode()


def del_record(rname, rtype, rvalue):
    """
    Delete a given record

    Args:
        rname (str): record name
        rtype (str): record type (A|CNAME|TXT...)
        rvalue (str): record value
    """
    result = subprocess.run(["samba-tool", "dns", "delete", "127.0.0.1", DOMAIN, rname, rtype, rvalue] + USR, env=ENV)
    result.check_returncode()


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print_dns()
        exit()

    if sys.argv[1] == "update":
        rname = sys.argv[2]
        rtype = sys.argv[3]
        rvalue = sys.argv[4]
        update_record(rname, rtype, rvalue)
        exit()

    if sys.argv[1] == "add":
        rname = sys.argv[2]
        rtype = sys.argv[3]
        rvalue = sys.argv[4]
        add_record(rname, rtype, rvalue)
        exit()

    if sys.argv[1] == "del" or sys.argv[1] == "delete":
        rname = sys.argv[2]
        rtype = sys.argv[3]
        rvalue = sys.argv[4]
        del_record(rname, rtype, rvalue)
        exit()