from weakref import WeakKeyDictionary
from itertools import chain
import threading
import time

from routersploit.utils import print_status


class Option(object):
    """ Exploit attribute that is set by the end user. """

    def __init__(self, default, description=""):
        self.default = default
        self.description = description
        self.data = WeakKeyDictionary()

    def __get__(self, instance, owner):
        return self.data.get(instance, self.default)

    def __set__(self, instance, value):
        self.data[instance] = value


class ExploitOptionsAggregator(type):
    """ Metaclass for exploit base class.

    Metaclass is aggregating all possible Attributes that user can set
    for tab completion purposes.
    """
    def __new__(cls, name, bases, attrs):
        try:
            base_exploit_attributes = chain(map(lambda x: x.exploit_attributes, bases))
        except AttributeError:
            attrs['exploit_attributes'] = {}
        else:
            attrs['exploit_attributes'] = {k: v for d in base_exploit_attributes for k, v in d.iteritems()}

        for key, value in attrs.iteritems():
            if isinstance(value, Option):
                attrs['exploit_attributes'].update({key: value.description})
            elif key == "__info__":
                attrs["_{}{}".format(name, key)] = value
                del attrs[key]
            elif key in attrs['exploit_attributes']:  # Removing exploit_attribute that was overwritten
                del attrs['exploit_attributes'][key]  # in the child and is not a Option() instance.
        return super(ExploitOptionsAggregator, cls).__new__(cls, name, bases, attrs)


class Exploit(object):
    """ Base class for exploits. """

    __metaclass__ = ExploitOptionsAggregator
    target = Option(default="", description="Target IP address.")
    # port = Option(default="", description="Target port.")

    @property
    def options(self):
        """ Returns list of options that user can set.

        Returns list of options aggregated by
        ExploitOptionsAggregator metaclass that user can set.

        :return: list of options that user can set
        """
        return self.exploit_attributes.keys()

    def run(self):
        raise NotImplementedError("You have to define your own 'run' method.")

    def check(self):
        raise NotImplementedError("You have to define your own 'check' method.")

    def run_threads(self, threads, target, *args, **kwargs):
        workers = []
        threads_running = threading.Event()
        threads_running.set()
        for worker_id in xrange(int(threads)):
            worker = threading.Thread(
                target=target,
                args=chain((threads_running,), args),
                kwargs=kwargs,
                name='worker-{}'.format(worker_id),
            )
            workers.append(worker)
            worker.start()

        start = time.time()
        try:
            while worker.isAlive():
                worker.join(1)
        except KeyboardInterrupt:
            threads_running.clear()

        for worker in workers:
            worker.join()
        print_status('Elapsed time: ', time.time() - start, 'seconds')
