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
"""

from __future__ import print_function

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


import cmd
import pkg_resources
from shlex import split as split_args
import HurricaneDNS

__version__ = pkg_resources.require("hurricanedns")[0].version
__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.__username = username
        self.__password = password
        self.__tfa = tfa
        self.__hdns = None
        self._make_prompt()
        self._get_hdns()

    @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._do_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._do_error("Must specify both a source and target")
            return

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

        existing = self._get_hdns().cache_domains()
        existing = list(item['domain'] for item in existing)

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

        if src not in existing:
            self._do_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._get_hdns().add_domain(target)
            except HurricaneDNS.HurricaneError as e:
                self._do_error(e)
                return

        src_records = self._get_hdns().cache_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._get_hdns().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 = map(lambda x: x['domain'], self._get_hdns().cache_domains())
        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]

        !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.

        """
        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._get_hdns().add_domain(domain, **extra)
            elif 6 >= len(args) >= 4:
                    self._get_hdns().add_record(*args)
            else:
                self._do_error('Invalid arguments')
        except HurricaneDNS.HurricaneError as e:
            self._do_error(e)

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

        domains = map(lambda x: x['domain'], self._get_hdns().cache_domains())
        if pos == 2:
            return self.filter_down(args, pos, domains)
        elif not args[1] 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._get_hdns().del_domain(args[0])
            elif len(args) > 1:
                self._get_hdns().del_records(*args)
        except HurricaneDNS.HurricaneError as e:
            self._do_error(e)

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

        domains = map(lambda x: x['domain'], self._get_hdns().cache_domains())
        if pos <= 1:
            return
        elif pos == 2:
            return self.filter_down(args, pos, domains)
        else:
            if args[1] in domains:
                records = self._get_hdns().cache_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,
                         map(lambda x: x['domain'], self._get_hdns().cache_domains())))
        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._get_hdns().cache_domains()
            existing = set([item['domain'] for item in existing])
            records = []
            for domain in split_args(args):
                if domain.lower() not in existing:
                    self._do_error('Invalid domain: ' + domain)
                    continue
                records.extend(self._get_hdns().cache_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 = '%%%ds %%-5s %%-%ds %%%ds %%6s' % (maxhost, maxvalue, maxttl)
                print(template % ('HOST', 'TYPE', 'VALUE', 'TTL', 'MX'))
                for record in records:
                    print(template % (record['host'], record['type'], record['value'], record['ttl'], record['mx']))
        else:
            domains = self._get_hdns().cache_domains()
            domains = sorted(domains, key=lambda item: item['domain'])
            print('TYPE       DOMAIN')
            for domain in domains:
                print('%-9s %s' % (domain['type'], domain['domain']))

    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:
            from dns import zone
            import dns.rdatatype
        except ImportError as e:
            self._do_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], 'r').read()
            except:
                self._do_error('Unable to open %s for reading' % args[1])
                return
        else:
            self._do_error('Invalid arguments')
            return

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

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

        for host, data in Z.iteritems():
            if str(host) == '@':
                name = domain
            else:
                name = '%s.%s' % (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 = '%s %s %s' % (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._get_hdns().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._do_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 _do_error(self, errmsg, command=None):
        if not command:
            command = self.lastcmd.split()[0]
        print('hdnssh: %s: %s' % (command, errmsg))

    def _get_hdns(self):
        if not self.__hdns:
            self.__hdns = HurricaneDNS.HurricaneDNS(self.__username, self.__password, self.__tfa)
        return self.__hdns

    def _make_prompt(self):
        self.prompt = '[%s@dns.he.net] ' % self.__username

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


def main():
    import argparse
    from sys import exit, stdin
    parser = argparse.ArgumentParser()
    parser.add_argument('-v', '--version', action='version', version="%%(prog)s %s (%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('%s: HE sent an error (%s)' % (parser.prog, e))
        exit(1)
    except HurricaneDNS.HurricaneError as e:
        print('%s: %s' % (parser.prog, e))


if __name__ == '__main__':
    main()