#!/usr/bin/env python
# Checks a bulk of DNS names for their existence (based on the SOA record).
#
# Copyright 2015 Peter Wu <peter@lekensteyn.nl>
import datetime
import dns.resolver, dns.exception
from itertools import product
import sys

def crange(begin, end):
    """Generates a range of ascii characters."""
    return [chr(i) for i in range(ord(begin), ord(end) + 1)]

def parse_pattern(pattern):
    """Parses a pattern with character classes into a list of character sets."""
    components = []
    class_chars = None
    range_begin = None
    for c in pattern:
        if class_chars is not None:
            if c == ']':
                if range_begin is not None:
                    # Matches '[X-]'
                    class_chars += [range_begin, '-']
                    range_begin = None
                components.append(sorted(set(class_chars)))
                class_chars = None
            elif range_begin is not None:
                # Detected 'X-Y'
                class_chars += crange(range_begin, c)
                range_begin = None
            else:
                if c == '-' and class_chars:
                    # Detected 'X-' (at least one prev. char in char class).
                    range_begin = class_chars.pop()
                else:
                    # Detected 'X' or '-'
                    class_chars.append(c)
        else:
            if c == '[':
                class_chars = []
            else:
                components.append(c)
    assert range_begin is None
    assert class_chars is None, "Character class must end with ]"
    return components

def get_input(patterns):
    """
    Accepts patterns such as:
     - [0-9].tld
     - [0a-z].tld
    """
    for pattern in patterns:
        format = parse_pattern(pattern)
        for item in product(*format):
            yield ''.join(item)

def resolve(domain):
    # May throw dns.exception.DNSException
    answer = dns.resolver.query(domain, 'SOA')

    # May throw if malformed... hopefully it does not
    return answer[0].to_text()

def resolve_inputs(inputs):
    for line in inputs:
        domain = line.strip()
        try:
            name = resolve(domain)
            yield "{}: {}".format(line, repr(name))
        except dns.exception.DNSException as e:
            yield "{}! {}".format(line, type(e).__name__)

def main(argv):
    if len(argv) < 3:
        sys.stderr.write('Usage: {} resolved-names.txt pattern...\n')
        sys.stderr.write('Usage: {} "" pattern... (to print patterns)\n')
        return 1

    if argv[1] == '':
        for key in get_input(argv[2:]):
            print(key)
        return 0

    with open(argv[1], 'a') as outf:
        for result in resolve_inputs(get_input(sys.argv[2:])):
            now = datetime.datetime.now()
            sys.stderr.write('{}  '.format(now.time()))
            sys.stderr.flush()
            print(result)
            outf.write(result + "\n")

if __name__ == '__main__':
    sys.exit(main(sys.argv))