#!/usr/bin/env python27
# -*- coding: utf-8 -*-
'''dnsfire

Usage:
  dnsfire [--rate <RATE>] [FILE...]

Options:
  inputfile             File to take as input, or read from stdin.
  --rate <rate>         DNS requests per second. Otherwise unlimited.
  -h --help             Show this screen.
  --version             Show version.
'''

from __future__ import unicode_literals, print_function
from docopt import docopt
import fileinput
import re
from time import time
import adns

__version__ = "0.1.0"
__author__ = "Justin Hoppensteadt"
__license__ = "MIT"


class AsyncResolver(object):

    def __init__(self, hosts, intensity=100):
        """
        hosts: a list of hosts to resolve
        intensity: how many hosts to resolve at once
        """
        self.hosts = hosts
        self.intensity = intensity
        self.adns = adns.init()

    def resolve(self):
        """
        Resolves hosts and returns a dictionary of { 'host': 'ip' }.
        """
        resolved_hosts = {}
        active_queries = {}
        host_queue = self.hosts[:]

        def collect_results():
            for query in self.adns.completed():
                answer = query.check()
                host = active_queries[query]
                del active_queries[query]
                if answer[0] == 0:
                    ip = answer[3][0]
                    resolved_hosts[host] = ip
                elif answer[0] == 101:  # CNAME
                    query = self.adns.submit(answer[1], adns.rr.A)
                    active_queries[query] = host
                else:
                    resolved_hosts[host] = None

        def finished_resolving():
            return len(resolved_hosts) == len(self.hosts)

        while not finished_resolving():
            while host_queue and len(active_queries) < self.intensity:
                host = host_queue.pop()
                query = self.adns.submit(host, adns.rr.A)
                active_queries[query] = host
            collect_results()

        return resolved_hosts


class dnsfire:

    def __init__(self):
        self.ValidIp = re.compile(
            r'((25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)'
            r'(\.(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)){3})')
        self.CommonHost = re.compile(
            r'((([a-zA-Z0-9]|[a-zA-Z0-9][a-zA-Z0-9\-]*[a-zA-Z0-9])\.)'
            r'*(com|org|net|info|tmcs|pl|tv)\.?)', re.IGNORECASE)


def main():
    '''Main entry point for the dnsfire CLI.'''
    args = docopt(__doc__, version=__version__)
    print(args)
    host_format = "cli%d.dart.ads2.websys.tmcs"
    number_of_hosts = 48

    hosts = [host_format % i for i in range(number_of_hosts)]

    ar = AsyncResolver(hosts, intensity=500)
    start = time()
    resolved_hosts = ar.resolve()
    print resolved_hosts
    end = time()


    print ('It took ', end - start, 'seconds to resolve',
        number_of_hosts)


if __name__ == '__main__':
    main()
