hurricanedns

#!/usr/bin/env python
"""HurricaneDNS Command Line tools/shell

Inspired by & heavily borrowed from the EveryDNS Command Line tool/shell by
    Scott Yang: http://svn.fucoder.com/fucoder/pyeverydns/everydnscmd.py
"""


try:
    from cStringIO import StringIO
except ImportError:
    from io import StringIO

import cmd
from importlib.metadata import version
from shlex import split as split_args

import HurricaneDNS

try:
    __version__ = version("hurricanedns")
except:
    __version__ = "1.0.3"
__author__ = "Brian Hartvigsen <brian.andrew@brianandjenny.com>"
__copyright__ = "Copyright 2015, Brian Hartvigsen"
__credits__ = ["Scott Yang", "Brian Hartvigsen"]
__license__ = "MIT"


def write_help(func):
    def decorator(self):
        for line in StringIO(func.__doc__):
            line = line.strip()
            if line.startswith("!"):
                print(line[1:].upper())
            else:
                print("        " + line)

    return decorator

class HurricaneDNSShell(cmd.Cmd):
    def __init__(self, username, password, tfa):
        cmd.Cmd.__init__(self)
        self.heDNS = HurricaneDNS.HurricaneDNS(username, password, tfa)
        self.prompt = "[%s@dns.he.net] " % username

    def __print_error(self, errmsg, command=None):
        if not command:
            command = self.lastcmd.split()[0]
        print(f"hdnssh: {command}: {errmsg}")

    @staticmethod
    def filter_down(args, pos, possibiles):
        start = args[pos - 1] if len(args) == pos else None
        final = list(
            filter(lambda x: x.startswith(start) if start else True, possibiles)
        )
        return final

    def default(self, line):
        self.__print_error("command not found", line.split()[0])

    def completedefault(self, text, line, begidx, endidx):
        pass

    def do_cp(self, args):
        """!NAME
        cp - copy records from one domain to another

        !SYNOPSIS
        cp src-domain target-domain

        !DESCRIPTION
        Copy all records from src-domain to target-domain. If target-domain
        does not exist, it will be created.

        """
        args = split_args(args)
        if len(args) < 2 or len(args) > 3:
            self.__print_error("Must specify both a source and target")
            return

        if len(args) == 3 and args[0].lower() != "-f":
            self.__print_error("If passing 3 arguments, first arg shoudl be -f")
            return

        existing = self.heDNS.domain_list

        src = args[0].lower()
        target = args[1].lower()

        if src not in existing:
            self.__print_error("There is no domain '%s', source must exist" % src)
            return

        if target not in existing:
            print("Domain '%s' does not exist, attempting to create" % target)
            try:
                self.heDNS.add_domain(target)
            except HurricaneDNS.HurricaneError as e:
                self.__print_error(e)
                return

        src_records = self.heDNS.get_domain_records(src)

        for record in src_records:
            host = record["host"][::-1].replace(src[::-1], target[::-1], 1)[::-1]
            if record["type"] == "SOA":
                continue
            if "extended" not in record:
                record["extended"] = record["value"]
            self.heDNS.add_record(
                target,
                host,
                record["type"],
                record["extended"],
                mx=record["mx"],
                ttl=record["ttl"],
            )

    def complete_cp(self, text, line, begidx, endidx):
        args = line.split()
        pos = len(args)
        if not text:
            pos += 1

        if pos < 2 or pos > 3:
            return

        domains = self.heDNS.domain_list
        if pos == 3:
            domains = filter(lambda x: x.lower() != args[1].lower(), domains)

        return self.filter_down(args, pos, domains)

    def do_add(self, args):
        """!NAME
        add - Add a domain/record

        !SYNOPSIS
        add domain [domain option]
        add domain host type value [mx] [ttl] [ddns_key]

        !DESCRIPTION
        Add either a domain or a host record to HurricaneDNS. If the number of
        arguments is 1 or 2, it will add the domain in arg 1, and then use arg
        2 as option. If the number of arguments is between 4 and 6, it will
        add the host record with optional MX and TTL. If argument 7 is offered,
        DDNS will be enabled for the record and used the argument as the key.

        """
        args = split_args(args)
        try:
            if len(args) and len(args) < 3:
                domain = args[0]
                extra = {}
                if len(args) == 2:
                    (option, value) = args[1].split("=", 1)
                    extra[option] = value
                self.heDNS.add_domain(domain, **extra)
            elif 7 >= len(args) >= 4:
                self.heDNS.add_record(*args)
            else:
                self.__print_error("Invalid arguments")
        except HurricaneDNS.HurricaneError as e:
            self.__print_error(e)

    def complete_add(self, text, line, begidx, endidx):
        args = line.split()
        pos = len(args)
        if not text:
            pos += 1

        domains = self.heDNS.domain_list
        if pos == 2:
            return self.filter_down(args, pos, domains)
        elif args[1] not in domains:
            if pos == 3:
                return self.filter_down(args, pos, ["method=", "master="])
            else:
                return []
        else:
            if pos == 4:
                types = []
                if args[1].endswith(".in-addr.arpa") or args[1].endswith(".ip6.arpa"):
                    types = ["CNAME", "NS", "PTR", "TXT"]
                else:
                    types = [
                        "A",
                        "AAAA",
                        "CNAME",
                        "ALIAS",
                        "MX",
                        "NS",
                        "TXT",
                        "CAA",
                        "AFSDB",
                        "HINFO",
                        "RP",
                        "LOC",
                        "NAPTR",
                        "PTR",
                        "SSHFP",
                        "SPF",
                        "SRV",
                    ]
                return self.filter_down(args, pos, types)
            else:
                pass

    def do_del(self, args):
        """!NAME
        del - Delete domain or host records

        !SYNOPSIS
        del domain
        del domain host [type] [value] [mx] [ttl]

        !DESCRIPTION
        Delete either domain or host records. If only 1 argument is given, it
        will use that as the domain name ot delete. Otherwise it will delete
        host records that match the arguments.

        """
        args = split_args(args)
        try:
            if len(args) == 1:
                self.heDNS.del_domain(args[0])
            elif len(args) > 1:
                self.heDNS.del_records(*args)
        except HurricaneDNS.HurricaneError as e:
            self.__print_error(e)

    def complete_del(self, text, line, begidx, endidx):
        args = line.split()
        pos = len(args)
        if not text:
            pos += 1

        domains = self.heDNS.domain_list
        if pos <= 1:
            return
        elif pos == 2:
            return self.filter_down(args, pos, domains)
        else:
            if args[1] in domains:
                records = self.heDNS.get_domain_records(args[1])
                records = filter(lambda x: x["status"] != "locked", records)
            else:
                records = None

        if pos == 3:
            records = map(lambda x: x["host"], records)
            return self.filter_down(args, pos, records)
        else:
            records = filter(lambda x: x["host"].lower() == args[2].lower(), records)

        if pos == 4:
            types = map(lambda x: x["type"], records)
            if types:
                pass
            elif args[1].endswith(".in-addr.arpa") or args[1].endswith(".ip6.arpa"):
                types = ["CNAME", "NS", "PTR", "TXT"]
            else:
                types = [
                    "A",
                    "AAAA",
                    "CNAME",
                    "MX",
                    "NS",
                    "TXT",
                    "AFSDB",
                    "HINFO",
                    "RP",
                    "LOC",
                    "NAPTR",
                    "PTR",
                    "SSHFP",
                    "SPF",
                    "SRV",
                ]
            return self.filter_down(args, pos, types)
        else:
            records = filter(lambda x: x["type"].lower() == args[3].lower(), records)

        if pos == 5:
            values = map(
                lambda x: x["extended"] if "extended" in x else x["value"], records
            )
            return self.filter_down(args, pos, values)
        else:
            records = filter(
                lambda x: (x["extended"] if "extended" in x else x["value"]).lower()
                == args[4].lower(),
                records,
            )

        if pos == 6:
            mx = map(lambda x: x["mx"], records)
            return self.filter_down(args, pos, mx)
        else:
            records = filter(lambda x: x["mx"] == args[5], records)

        if pos == 7:
            ttl = map(lambda x: x["ttl"], records)
            return self.filter_down(args, pos, ttl)

    def do_EOF(self, args):
        print()
        return 1

    def do_exit(self, args):
        return 1

    def complete_ls(self, text, line, begidx, endidx):
        domains = list(
            filter(
                lambda x: x.startswith(text) if text else True,
                self.heDNS.domain_list,
            )
        )
        return domains

    def do_ls(self, args):
        """!NAME
        ls - List domains or host records

        !SYNOPSIS
        ls
        ls domain [domain...]

        !DESCRIPTION
        Listing all the domains when there is no argument. Otherwise list all
        host records from the specified domains.

        """
        if args:
            existing = self.heDNS.domain_list
            records = []
            for domain in split_args(args):
                if domain.lower() not in existing:
                    self.__print_error("Invalid domain: " + domain)
                    continue
                records.extend(self.heDNS.get_domain_records(domain))

            if records:
                maxhost = max(len(item["host"]) for item in records)
                maxvalue = max(len(item["value"]) for item in records)
                maxttl = max(len(item["ttl"]) for item in records)
                template = f"%-10s %{maxhost}s %-5s %-{maxvalue}s %{maxttl}s %5s %4s"
                print(template % ("ID", "HOST", "TYPE", "VALUE", "TTL", "MX", "DDNS"))
                for record in records:
                    print(
                        template
                        % (
                            record["id"],
                            record["host"],
                            record["type"],
                            record["value"],
                            record["ttl"],
                            record["mx"],
                            "Y" if record["ddns"] == "1" else "N",
                        )
                    )
        else:
            # ls domains
            domain_dict = self.heDNS.get_domain_info("all")
            template = "%-10s %-9s %s"
            print(template % ("ID", "DOMAIN", "TYPE"))
            for domain in sorted(domain_dict.keys()):
                print(
                    template
                    % (domain_dict[domain]["id"], domain, domain_dict[domain]["type"])
                )

    def do_import(self, args):
        """!NAME
        import - Import a domain from BIND zone file

        !SYNOPSIS
        import domain zone_file

        !DESCRIPTION
        Importing a BIND zone file into a domain. Just normal domain (not reserve) for now.

        """
        try:
            import dns.rdatatype
            from dns import zone
        except ImportError as e:
            self.__print_error("You do not have python-dns installed!")
            return

        args = split_args(args)
        if len(args) and len(args) == 2:
            domain = args[0]
            try:
                zonedata = open(args[1]).read()
            except:
                self.__print_error("Unable to open %s for reading" % args[1])
                return
        else:
            self.__print_error("Invalid arguments")
            return

        try:
            Z = zone.from_text(zonedata, domain)
        except SyntaxError as e:
            if e.message:
                self.__print_error(e.message)
            else:
                self.__print_error("Syntax error found in zone file")
            return
        except Exception as e:
            self.__print_error("Unable to parse zone file")
            return

        # create the domain
        try:
            self.heDNS.add_domain(domain)
        except HurricaneDNS.HurricaneError as e:
            self.__print_error(e)
            return

        for host, data in Z.iteritems():
            if str(host) == "@":
                name = domain
            else:
                name = f"{host}.{domain}"
            for r in data.rdatasets:
                for rdata in r:
                    if (
                        "*" in name
                        or r.rdtype == dns.rdatatype.SOA
                        or r.rdtype == dns.rdatatype.NS
                    ):
                        print(
                            "SKIPPING:",
                            name,
                            dns.rdatatype.to_text(r.rdtype),
                            r.ttl,
                            rdata,
                        )
                        continue
                    elif r.rdtype == dns.rdatatype.SRV:
                        _value = f"{rdata.weight} {rdata.port} {rdata.target.to_text()}"
                        _priority = rdata.priority
                    elif r.rdtype == dns.rdatatype.MX:
                        _value = rdata.exchange.to_text()
                        _priority = rdata.preference
                    else:
                        _priority = None
                        if hasattr(rdata, "address"):
                            _value = rdata.address
                        elif hasattr(rdata, "value"):
                            _value = rdata.value
                        elif hasattr(rdata, "target"):
                            _value = rdata.target
                        else:
                            _value = ""

                    try:
                        self.heDNS.add_record(
                            domain,
                            name,
                            dns.rdatatype.to_text(rdata.rdtype),
                            _value,
                            _priority,
                            r.ttl,
                        )
                    except HurricaneDNS.HurricaneError as e:
                        if "properly delegated" in str(e):
                            pass
                        self.__print_error(e)

    def emptyline(self):
        pass

    help_add = write_help(do_add)
    help_del = write_help(do_del)
    help_ls = write_help(do_ls)
    help_import = write_help(do_import)
    help_cp = write_help(do_cp)

    def cmdloop(self):
        try:
            cmd.Cmd.cmdloop(self)
        except KeyboardInterrupt:
            print()
            self.cmdloop()


if __name__ == "__main__":
    import argparse
    from sys import exit, stdin

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-v",
        "--version",
        action="version",
        version=f"%(prog)s {__version__} ({__author__})",
    )
    parser.add_argument("username", help="Your HE DNS username")
    parser.add_argument("password", help="Your HE DNS password")
    parser.add_argument("--tfa", "-t", type=int, help="Your 2fa code")
    parser.add_argument(
        "command",
        nargs=argparse.REMAINDER,
        help="An optional command, if blank we drop into interactive mode",
    )
    options = parser.parse_args()

    try:
        shell = HurricaneDNSShell(options.username, options.password, options.tfa)
        if not options.command:
            if stdin.isatty():
                shell.cmdloop()
            else:
                for line in stdin.readlines():
                    shell.onecmd(line)
        else:
            from pipes import quote

            shell.onecmd(" ".join(map(lambda x: quote(x), options.command)))
    except HurricaneDNS.HurricaneAuthenticationError as e:
        print(f"{parser.prog}: HE sent an error ({e})")
        exit(1)
    except HurricaneDNS.HurricaneError as e:
        print(f"{parser.prog}: {e}")