﻿import re
import os
import importlib
import inspect
import time
import searchengine.solr_tools
from searchengine.database import DatabaseConnector
from urllib.request import Request
from urllib.parse import urlencode
from searchengine.debugtools import log, log_exception

##
# @class    BaseExploit
#
# @brief    An abstract class defining a basic plugin structure for exploits.
#
# @author   Edward Callahan
# @date 8/3/2016
class BaseExploit:

    ##
    # @fn   __init__(self, exploit_unique_identifier, exploit_cve, exploit_version, exploit_human_readable_name)
    #
    # @brief    Class initializer.
    #
    # @author   Edward Callahan
    # @date 8/3/2016
    #
    # @param    self    The class instance that this method operates on.
    def __init__(self, exploit_cve, exploit_version, exploit_human_readable_name):
        super().__init__()
        self.__exploit_cve = exploit_cve                                    #< A string defining the technical name for this exploit.
        self.__exploit_version = exploit_version                            #< A string defining the technical version for this exploit.
        self.__exploit_human_readable_name = exploit_human_readable_name    #< A string defining the human readable vulnerability that this exploits.
        self.__current_endpoint = ""                                        #< A string correlating to the current endpoint being tested.
        self.__solr_vulnerable_hosts = searchengine.solr_tools.get_solr_instance('vulnerable_hosts', url_offset = 0)

    ##
    # @fn   _set_endpoint_is_vulnerable(self)
    #
    # @brief    Execute this when the current endpoint is vulnerable to this exploit.
    #
    # @author   Intricate
    # @date 8/25/2016
    #
    # @param    self    The class instance that this method operates on.
    def _set_endpoint_is_vulnerable(self):
        # At the moment, this will overwrite an existing document. Update coming soon to append to vulnerable_to field.
        docs = []
        docs.append({
            'id'                : self.__current_endpoint,
            'vulnerable_to'     : {"exploit_cve" : self.__exploit_cve, "exploit_human_readable_name" : self.__exploit_human_readable_name},
            'last_pentest_time' : time.time()
        })
        self.__solr_vulnerable_hosts.add(docs, commit = True, overwrite = True)
        log("{} is vulnerable to {}({})".format(self.__current_endpoint, self.__exploit_cve, self.__exploit_human_readable_name))

    ###################################################################
    # Abstract Functions
    ###################################################################

    ##
    # @fn   feed_endpoint(endpoint)
    #
    # @brief    Feed an endpoint to this plugin (to be implemented in child classes).
    #
    # @author   Edward Callahan
    # @date 8/3/2016
    #
    # @param    endpoint    The endpoint.
    def feed_endpoint(self, endpoint):
        pass

##
# @class    ExploitManager
#
# @brief    Handles loading and initializing exploits.
#
# @author   Edward Callahan
# @date 8/5/2016
class ExploitManager:

    ##
    # @fn   __init__(self)
    #
    # @brief    Class initializer.
    #
    # @author   Edward Callahan
    # @date 8/5/2016
    #
    # @param    self    The class instance that this method operates on.
    def __init__(self):
        super().__init__()
        self.lock = None
        self.hosts_to_pentest = []
        self.__solr_vulnerable_hosts = searchengine.solr_tools.get_solr_instance('vulnerable_hosts', url_offset = 0)
        self.__exploit_modules = []
        self.__load_exploit_modules()

    ##
    # @fn   get_host_to_pentest(self)
    #
    # @brief    Retrieves a host that needs to be pentested.
    #
    # @author   Intricate
    # @date 8/24/2016
    #
    # @param    self    The class instance that this method operates on.
    def get_host_to_pentest(self):
        if len(self.hosts_to_pentest) == 0:
            # Temporary
            import threading
            with threading.Lock():
                response = self.__solr_vulnerable_hosts.search(
                    "last_pentest_time:[0 TO " + str(int(time.time() - (60 * 60 * 24 * 7))) + "] AND domain:* AND -path:*", 
                    rows=20
                    )
                if len(response.docs) == 0:
                    return False
                doc_updates = []
                for doc in response.docs:
                    doc_updates.append({
                        "id"                : doc["id"],
                        "vulnerable_to"     : doc["vulnerable_to"],
                        "last_pentest_time" : int(time.time())
                    })
                    self.hosts_to_pentest.append(doc["id"])
                self.__solr_vulnerable_hosts.add(doc_updates)
        next_host = self.hosts_to_pentest.pop(0)
        return next_host

    ##
    # @fn   pentest_host(self, host)
    #
    # @brief    Pentests a given host.
    #
    # @author   Intricate
    # @date 8/24/2016
    #
    # @param    self    The class instance that this method operates on.
    # @param    host    The host to pentest.
    def pentest_host(self, host):
        for exploit_module in self.__exploit_modules:
            exploit_module.feed_endpoint(host)

    ##
    # @fn   __load_exploit_modules(self)
    #
    # @brief    Loads exploit modules at runtime. This allows an application to iterate through all stored exploits.
    #
    # @author   Edward Callahan
    # @date 8/5/2016
    #
    # @param    self    The class instance that this method operates on.
    def __load_exploit_modules(self):
        # Modified plugin loader:
        # https://copyninja.info/blog/dynamic-module-loading.html
        self.__exploit_modules.clear()
        log("Loading exploits...")
        pysearchre = re.compile('.py$', re.IGNORECASE)
        exploitfiles = filter(pysearchre.search,
                              os.listdir(os.path.join(os.path.dirname(__file__),
                              'exploits')))
        form_module = lambda fp: '.' + os.path.splitext(fp)[0]
        exploits = map(form_module, exploitfiles)
        # import parent module / namespace
        importlib.import_module('searchengine.vulnerability_scanner.exploits')
        modules = []
        for exploit in exploits:
            modules.append(importlib.import_module(exploit, package="searchengine.vulnerability_scanner.exploits"))
        for module in modules:
            for name, obj in inspect.getmembers(module, inspect.isclass):
                if name == "Exploit":
                    log("Loaded exploit: " + str(obj))
                    self.__exploit_modules.append(obj())