import copy
import glob
import logging
import os
import re
import socket
import socks
import importlib
import prettytable
from termcolor import colored
from queue import Queue
from urllib.parse import urlsplit

from pocsuite3.lib.core.clear import remove_extra_log_message
from pocsuite3.lib.core.common import boldify_message, check_file, get_file_items, parse_target, \
    get_public_type_members, data_to_stdout
from pocsuite3.lib.core.common import check_path, extract_cookies
from pocsuite3.lib.core.common import get_local_ip, desensitization
from pocsuite3.lib.core.common import single_time_warn_message
from pocsuite3.lib.core.common import OrderedSet
from pocsuite3.lib.core.convert import stdout_encode
from pocsuite3.lib.core.data import conf, cmd_line_options
from pocsuite3.lib.core.data import kb
from pocsuite3.lib.core.data import logger
from pocsuite3.lib.core.data import merged_options
from pocsuite3.lib.core.data import paths
from pocsuite3.lib.core.datatype import AttribDict
from pocsuite3.lib.core.enums import HTTP_HEADER, CUSTOM_LOGGING, PROXY_TYPE
from pocsuite3.lib.core.exception import PocsuiteSyntaxException, PocsuiteSystemException, PocsuiteHeaderTypeException
from pocsuite3.lib.core.log import FORMATTER
from pocsuite3.lib.core.register import load_file_to_module
from pocsuite3.lib.core.settings import DEFAULT_USER_AGENT, DEFAULT_LISTENER_PORT, CMD_PARSE_WHITELIST
from pocsuite3.lib.core.statistics_comparison import StatisticsComparison
from pocsuite3.lib.core.update import update
from pocsuite3.lib.parse.cmd import DIY_OPTIONS
from pocsuite3.lib.parse.configfile import config_file_parser
from pocsuite3.lib.parse.rules import regex_rule
from pocsuite3.lib.request.patch import patch_all
from pocsuite3.modules.listener import start_listener


def _resolve_cross_references():
    import pocsuite3
    pocsuite3.lib.core.revision.stdout_encode = stdout_encode
    pocsuite3.lib.core.convert.single_time_warn_message = single_time_warn_message


def set_verbosity():
    """
    This function set the verbosity of pocsuite output messages.
    """

    if conf.verbose is None:
        conf.verbose = 1

    conf.verbose = int(conf.verbose)

    if conf.verbose == 0:
        logger.setLevel(logging.ERROR)
    elif conf.verbose == 1:
        logger.setLevel(logging.INFO)
    elif conf.verbose == 2:
        logger.setLevel(logging.DEBUG)
    elif conf.verbose == 3:
        logger.setLevel(logging.DEBUG)
        logger.setLevel(CUSTOM_LOGGING.SYSINFO)
    elif conf.verbose == 4:
        logger.setLevel(logging.DEBUG)
        logger.setLevel(CUSTOM_LOGGING.WARNING)
    elif conf.verbose >= 5:
        logger.setLevel(logging.DEBUG)
        logger.setLevel(CUSTOM_LOGGING.ERROR)


def _set_http_user_agent():
    '''
    set user-agent
    :return:
    '''

    conf.http_headers[HTTP_HEADER.USER_AGENT] = DEFAULT_USER_AGENT

    if conf.random_agent:
        uapath = os.path.join(paths.POCSUITE_DATA_PATH, 'user-agents.txt')
        if os.path.exists(uapath):
            with open(uapath) as f:
                agents = f.read().split("\n")
                if len(agents) == 1 and "" in agents:
                    logger.error("user-agents file is empty will use default")
                else:
                    conf.agents = agents
        else:
            logger.error("user-agents file not fond will use default")

    if conf.agent:
        conf.http_headers[HTTP_HEADER.USER_AGENT] = conf.agent


def _set_http_referer():
    if conf.referer:
        conf.http_headers[HTTP_HEADER.REFERER] = conf.referer


def _set_http_cookie():
    if conf.cookie:
        if isinstance(conf.cookie, dict):
            conf.http_headers[HTTP_HEADER.COOKIE] = '; '.join(map(lambda x: '='.join(x), conf.cookie.items()))
        else:
            conf.http_headers[HTTP_HEADER.COOKIE] = conf.cookie


def _set_http_host():
    if conf.host:
        conf.http_headers[HTTP_HEADER.HOST] = conf.host


def _set_http_extra_headers():
    if conf.headers:
        conf.headers = conf.headers.split("\n") if "\n" in conf.headers else conf.headers.split("\\n")
        for header_value in conf.headers:
            if not header_value.strip():
                continue

            if header_value.count(':') >= 1:
                header, value = (_.lstrip() for _ in header_value.split(":", 1))
                if header and value:
                    if header not in conf.http_headers:
                        conf.http_headers[header] = value


def _set_network_timeout():
    if conf.timeout:
        conf.timeout = float(conf.timeout)
        if conf.timeout < 3.0:
            warn_msg = "the minimum HTTP timeout is 3 seconds, pocsuite "
            warn_msg += "will going to reset it"
            logger.warn(warn_msg)

            conf.timeout = 3.0
    else:
        conf.timeout = 30.0

    socket.setdefaulttimeout(conf.timeout)


def _set_network_proxy():
    if conf.proxy:
        debug_msg = "setting the HTTP/SOCKS proxy for all network requests"
        logger.debug(debug_msg)

        try:
            _ = urlsplit(conf.proxy)
        except Exception as ex:
            err_msg = "invalid proxy address '{0}' ('{1}')".format(conf.proxy, str(ex))
            raise PocsuiteSyntaxException(err_msg)

        hostname_port = _.netloc.split(":")
        scheme = _.scheme.upper()
        hostname = hostname_port[0]
        port = None
        username = None
        password = None

        if len(hostname_port) == 2:
            try:
                port = int(hostname_port[1])
            except Exception:
                pass

        if not all((scheme, hasattr(PROXY_TYPE, scheme), hostname, port)):
            err_msg = "proxy value must be in format '({0})://address:port'".format("|".join(
                _[0].lower() for _ in get_public_type_members(PROXY_TYPE)))
            raise PocsuiteSyntaxException(err_msg)

        if conf.proxy_cred:
            _ = re.search(r"\A(.*?):(.*?)\Z", conf.proxy_cred)
            if not _:
                err_msg = "proxy authentication credentials "
                err_msg += "value must be in format username:password"
                raise PocsuiteSyntaxException(err_msg)
            else:
                username = _.group(1)
                password = _.group(2)

        if scheme in (PROXY_TYPE.SOCKS4, PROXY_TYPE.SOCKS5, PROXY_TYPE.SOCKS5H):
            socks.set_default_proxy(
                socks.PROXY_TYPE_SOCKS4 if scheme == PROXY_TYPE.SOCKS4 else socks.PROXY_TYPE_SOCKS5,
                hostname,
                port,
                username=username,
                password=password,
                rdns=True if scheme == PROXY_TYPE.SOCKS5H else False,
            )
            conf.origin_socks = copy.deepcopy(socket.socket)  # Convenient behind recovery
            socket.socket = socks.socksocket

        if conf.proxy_cred:
            proxy_string = "{0}@".format(conf.proxy_cred)
        else:
            proxy_string = ""

        proxy_string = "{scheme}://{proxy_string}{hostname}:{port}".format(scheme=scheme.lower(),
                                                                           proxy_string=proxy_string,
                                                                           hostname=hostname, port=port)
        conf.proxies = {
            "http": proxy_string,
            "https": proxy_string
        }


def _set_multiple_targets():
    # set multi targets to kb
    if conf.url:
        targets = set()
        for url in conf.url:
            parsed = parse_target(url)
            if parsed:
                targets.add(parsed)
        if not targets:
            err_msg = "incorrect target url or ip format!"
            logger.error(err_msg)
        for target in targets:
            kb.targets.add(target)

    if conf.url_file:
        for line in get_file_items(conf.url_file, lowercase=False, unique=True):
            kb.targets.add(line)

    if conf.dork:
        # enable plugin 'target_from_zoomeye' by default
        if 'target_from_shodan' not in conf.plugins and 'target_from_fofa' not in conf.plugins and 'target_from_quake' not in conf.plugins:
            conf.plugins.append('target_from_zoomeye')

    if conf.dork_zoomeye:
        conf.plugins.append('target_from_zoomeye')

    if conf.dork_shodan:
        conf.plugins.append('target_from_shodan')

    if conf.dork_censys:
        conf.plugins.append('target_from_censys')

    if conf.dork_fofa:
        conf.plugins.append('target_from_fofa')

    if conf.dork_quake:
        conf.plugins.append('target_from_quake')

def _set_task_queue():
    if kb.registered_pocs and kb.targets:
        for poc_module in kb.registered_pocs:
            for target in kb.targets:
                kb.task_queue.put((target, poc_module))


def _check_account_login():
    # TODO
    # check telnet404 account login
    pass


def _check_ceye():
    # TODO
    # check ceye and ceye api token is available
    pass


def _check_seebug():
    # TODO
    # check seebug and seebug api token is available
    pass


def _check_zoomeye():
    # TODO
    # check zoomeye and zoomeye api token is available
    pass


def _set_threads():
    if not isinstance(conf.threads, int) or conf.threads <= 0:
        conf.threads = 1


def _set_connect_back():
    ips = get_local_ip(all=True)
    if ips:
        kb.data.local_ips = ips
    if conf.mode == "shell" and conf.connect_back_host is None:
        data_to_stdout("[i] pocsusite is running in shell mode, you need to set connect back host:\n")
        message = '----- Local IP Address -----\n'
        for i, ip in enumerate(kb.data.local_ips):
            message += "{0}    {1}\n".format(i, desensitization(ip) if conf.ppt else ip)
        data_to_stdout(message)
        while True:
            choose = None
            choose = input('Choose>: ').strip()
            if not choose:
                continue
            try:
                if choose.isdigit():
                    choose = int(choose)
                    conf.connect_back_host = kb.data.local_ips[choose]
                    data_to_stdout("you choose {0}\n".format(
                        desensitization(conf.connect_back_host) if conf.ppt else conf.connect_back_host))
                    break
            except Exception:
                data_to_stdout("wrong number, choose again\n")


def _set_listener():
    if conf.mode == "shell" and not conf.api:
        start_listener()


def _set_user_pocs_path():
    if conf.pocs_path:
        if check_path(conf.pocs_path):
            paths.USER_POCS_PATH = conf.pocs_path
            for root, dirs, files in os.walk(paths.USER_POCS_PATH):
                files = list(filter(lambda x: not x.startswith("__") and x.endswith(".py"), files))
                conf.poc = [os.path.join(paths.USER_POCS_PATH, f) for f in files]
        else:
            warm_msg = "User defined pocs path {0} is invalid".format(conf.pocs_path)
            logger.warn(warm_msg)


def _set_pocs_modules():
    # TODO
    # load poc scripts .pyc file support
    if conf.ssvid:
        conf.plugins.append('poc_from_seebug')
    if conf.poc:
        # step1. load system packed poc from pocsuite3/pocs folder
        exists_poc_with_ext = list(
            filter(lambda x: x not in ['__init__.py', '__init__.pyc'], os.listdir(paths.POCSUITE_POCS_PATH)))
        exists_pocs = dict([os.path.splitext(x) for x in exists_poc_with_ext])
        for poc in conf.poc:
            load_poc_sucess = False
            if any([poc in exists_poc_with_ext, poc in exists_pocs]):
                poc_name, poc_ext = os.path.splitext(poc)
                if poc_ext in ['.py', '.pyc']:
                    file_path = os.path.join(paths.POCSUITE_POCS_PATH, poc)
                else:
                    file_path = os.path.join(paths.POCSUITE_POCS_PATH, poc + exists_pocs.get(poc))
                if file_path:
                    info_msg = "loading PoC script '{0}'".format(file_path)
                    logger.info(info_msg)
                    load_poc_sucess = load_file_to_module(file_path)

            # step2. load poc from given file path
            try:
                if not load_poc_sucess:
                    if not poc.startswith('ssvid-') and check_file(poc):
                        info_msg = "loading PoC script '{0}'".format(poc)
                        logger.info(info_msg)
                        load_poc_sucess = load_file_to_module(poc)
            except PocsuiteSystemException:
                logger.error('PoC file "{0}" not found'.format(repr(poc)))
                continue

            # step3. load poc from seebug website using plugin 'poc_from_seebug'
            if not load_poc_sucess:
                if poc.startswith('ssvid-'):
                    info_msg = "loading Poc script 'https://www.seebug.org/vuldb/{0}'".format(poc)
                    logger.info(info_msg)
                    if "poc_from_seebug" not in conf.plugins:
                        conf.plugins.append('poc_from_seebug')

    load_keyword_poc_sucess = False
    if conf.vul_keyword:
        # step4. load poc with vul_keyword search seebug website
        info_msg = "loading PoC script from seebug website using search keyword '{0}' ".format(conf.vul_keyword)
        logger.info(info_msg)

        conf.plugins.append('poc_from_seebug')


def _set_plugins():
    # TODO
    # load plugin scripts .pyc file support
    if conf.plugins:
        founds = []
        cache_plugins = copy.deepcopy(conf.plugins)

        for found in glob.glob(os.path.join(paths.POCSUITE_PLUGINS_PATH, "*.py*")):
            dirname, filename = os.path.split(found)
            plugin_name = os.path.splitext(filename)[0]
            if found.endswith(('__init__.py', '__init__.pyc')):
                continue
            if plugin_name not in conf.plugins:
                continue
            cache_plugins.remove(plugin_name)
            founds.append(found)
        if len(cache_plugins) > 0:
            for file in cache_plugins:
                if os.path.exists(file):
                    founds.append(file)

        for file in founds:
            debug_msg = "loading plugin script '{0}'".format(file)
            logger.debug(debug_msg)
            load_file_to_module(file)


def _cleanup_options():
    """
    Cleanup configuration attributes.
    """
    if conf.agent:
        conf.agent = re.sub(r"[\r\n]", "", conf.agent)

    if conf.cookie:
        if isinstance(conf.cookie, str):
            conf.cookie = re.sub(r"[\r\n]", "", conf.cookie)
            conf.cookie = extract_cookies(conf.cookie)
        elif not isinstance(conf.cookie, dict):
            raise PocsuiteHeaderTypeException('Does not support type for cookie')
    if conf.delay:
        conf.delay = float(conf.delay)

    if conf.retry:
        conf.retry = min(conf.retry, 10)

    if conf.url:
        if isinstance(conf.url, str):
            conf.url = [conf.url]
        conf.url = [x.strip() for x in conf.url]

    if conf.poc:
        if isinstance(conf.poc, str):
            conf.poc = [conf.poc]
        conf.poc = [poc.lower() if poc.lower().startswith('ssvid-') else poc for poc in conf.poc]

    if conf.url_file:
        conf.url_file = os.path.expanduser(conf.url_file)
        check_file(conf.url_file)

    if conf.plugins:
        conf.plugins = conf.plugins.split(',')
        conf.plugins = [i.strip() for i in conf.plugins]
        conf.plugins = list(set(conf.plugins))

    if conf.connect_back_port:
        conf.connect_back_port = int(conf.connect_back_port)

    conf.origin_socks = None


def _basic_option_validation():
    _check_account_login()
    _check_seebug()
    _check_zoomeye()
    _check_ceye()


def _adjust_logging_formatter():
    """
    Solves problem of line deletition caused by overlapping logging messages
    and retrieved data info in inference mode
    """
    if hasattr(FORMATTER, '_format'):
        return

    def new_format(record):
        message = FORMATTER._format(record)
        message = boldify_message(message)
        return message

    FORMATTER._format = FORMATTER.format
    FORMATTER.format = new_format


def _create_directory():
    if not os.path.isdir(paths.POCSUITE_OUTPUT_PATH):
        try:
            os.makedirs(paths.POCSUITE_OUTPUT_PATH)
        except FileExistsError:
            pass

    if not os.path.isdir(paths.POCSUITE_TMP_PATH):
        try:
            os.makedirs(paths.POCSUITE_TMP_PATH)
        except FileExistsError:
            pass

    if not os.path.isfile(paths.POCSUITE_RC_PATH):
        open(paths.POCSUITE_RC_PATH, 'a').close()


def _set_conf_attributes():
    """
    This function set some needed attributes into the configuration
    singleton.
    """

    debug_msg = "initializing the configuration"
    logger.debug(debug_msg)

    conf.url = None
    conf.url_file = None
    conf.mode = 'verify'
    conf.poc = None
    conf.cookie = None
    conf.host = None
    conf.referer = None
    conf.agent = None
    conf.headers = None
    conf.random_agent = None
    conf.proxy = None
    conf.proxy_cred = None
    conf.proxies = {}
    conf.timeout = 30
    conf.retry = 0
    conf.delay = 0
    conf.http_headers = {}
    conf.agents = [DEFAULT_USER_AGENT]  # When loading from the plug-in, if the data source has no default value, it needs to be processed
    conf.login_user = None
    conf.login_pass = None
    conf.shodan_token = None
    conf.fofa_user = None
    conf.fofa_token = None
    conf.quake_token = None
    conf.censys_uid = None
    conf.censys_secret = None
    conf.dork = None
    conf.dork_zoomeye = None
    conf.dork_shodan = None
    conf.dork_fofa = None
    conf.dork_quake = None
    conf.dork_censys = None
    conf.dork_b64 = False
    conf.max_page = 1
    conf.search_type = 'host'
    conf.comparison = False
    conf.vul_keyword = None
    conf.ssvid = None
    conf.plugins = []
    conf.threads = 1
    conf.batch = False
    conf.check_requires = False
    conf.quiet = False
    conf.update_all = False
    conf.verbose = 1

    conf.ipv6 = False
    conf.multiple_targets = False
    conf.pocs_path = None
    conf.output_path = None
    conf.plugin_name = None
    conf.plugin_code = None
    conf.connect_back_host = None
    conf.connect_back_port = DEFAULT_LISTENER_PORT
    conf.console_mode = False
    conf.show_version = False
    conf.api = False  # api for zipoc
    conf.ppt = False
    conf.pcap = False
    conf.rule = False
    conf.rule_req = False
    conf.rule_filename = None
    conf.show_options = False


def _set_kb_attributes(flush_all=True):
    """
    This function set some needed attributes into the knowledge base
    singleton.
    """

    debug_msg = "initializing the knowledge base"
    logger.debug(debug_msg)

    kb.abs_file_paths = set()
    kb.os = None
    kb.os_version = None
    kb.arch = None
    kb.dbms = None
    kb.auth_header = None
    kb.counters = {}
    kb.multi_thread_mode = False
    kb.thread_continue = True
    kb.thread_exception = False
    kb.word_lists = None
    kb.single_log_flags = set()

    kb.cache = AttribDict()
    kb.cache.addrinfo = {}
    kb.cache.content = {}
    kb.cache.regex = {}

    kb.data = AttribDict()
    kb.data.local_ips = []
    kb.data.connect_back_ip = None
    kb.data.connect_back_port = DEFAULT_LISTENER_PORT
    kb.data.clients = []
    kb.targets = OrderedSet()
    kb.plugins = AttribDict()
    kb.plugins.targets = AttribDict()
    kb.plugins.pocs = AttribDict()
    kb.plugins.results = AttribDict()
    kb.results = []
    kb.current_poc = None
    kb.registered_pocs = AttribDict()
    kb.task_queue = Queue()
    kb.cmd_line = DIY_OPTIONS or []

    kb.comparison = None


def _merge_options(input_options, override_options):
    """
    Merge command line options with configuration file and default options.
    """
    if hasattr(input_options, "items"):
        input_options_items = input_options.items()
    else:
        input_options_items = input_options.__dict__.items()

    for key, value in input_options_items:
        if key not in conf or value not in (None, False) or override_options:
            conf[key] = value

    if input_options.get("configFile"):
        config_file_parser(input_options["configFile"])

    merged_options.update(conf)


def _set_poc_options(input_options):
    for line in input_options.keys():
        if line not in CMD_PARSE_WHITELIST:
            DIY_OPTIONS.append(line)


def init_options(input_options=AttribDict(), override_options=False):
    cmd_line_options.update(input_options)
    _set_conf_attributes()
    _set_poc_options(input_options)
    _set_kb_attributes()
    _merge_options(input_options, override_options)
    # export rules, dont run the poc in the default status
    if conf.rule or conf.rule_req:
        logger.info("The rule export function is in use. The POC is not executed at this point")
        if conf.pocs_path:
            if check_path(conf.pocs_path):
                paths.USER_POCS_PATH = conf.pocs_path
                for root, dirs, files in os.walk(paths.USER_POCS_PATH):
                    files = list(filter(lambda x: not x.startswith("__") and x.endswith(".py"), files))
                regex_rule(list(paths.USER_POCS_PATH + i for i in files))

        if conf.poc:
            regex_rule(conf.poc)
        exit()
    # if check version
    if conf.show_version:
        exit()


def _init_targets_plugins():
    for _, plugin in kb.plugins.targets.items():
        plugin.init()


def _init_pocs_plugins():
    for _, plugin in kb.plugins.pocs.items():
        plugin.init()


def _init_results_plugins():
    for _, plugin in kb.plugins.results.items():
        plugin.init()


def _init_kb_comparison():
    if conf.comparison:
        kb.comparison = StatisticsComparison()


def _init_target_from_poc_dork():
    if len(kb.targets) > 0:
        return

    for poc_module, poc_class in kb.registered_pocs.items():
        if not hasattr(poc_class, 'dork'):
            continue
        # find a available target source
        target_source = ''
        for i in ["zoomeye", "fofa", "shodan", "quake", "censys"]:
            if i in poc_class.dork.keys():
                target_source = i
                break
        # fetch target from target source, add it to kb.targets
        conf.dork = poc_class.dork[target_source]
        plugin_name = f'target_from_{target_source}'
        importlib.import_module(f'pocsuite3.plugins.{plugin_name}')
        for _, plugin in kb.plugins.targets.items():
            if target_source in plugin.__class__.__name__.lower():
                plugin.init()


def _show_pocs_modules_options():
    if not conf.show_options:
        return
    for module_name, poc_class in kb.registered_pocs.items():
        module_options = poc_class.options
        tb = prettytable.PrettyTable(
            ['Name', 'Current settings', 'Type', 'Description'])
        # add target option
        for name, opt in module_options.items():
            value = opt.value
            if opt.require and value == '':
                value = colored('*require*', 'red')
            tb.add_row([name, value, opt.type, opt.description])
        data_to_stdout(f'\nModule ({module_name}) options:\n')
        data_to_stdout(tb.get_string())
        data_to_stdout('\n')
    exit()


def init():
    """
    Set attributes into both configuration and knowledge base singletons
    based upon command line and configuration file options.
    """
    set_verbosity()
    patch_all()
    _adjust_logging_formatter()
    _cleanup_options()
    _basic_option_validation()
    _create_directory()
    _init_kb_comparison()
    update()
    _set_multiple_targets()
    _set_user_pocs_path()
    _set_pocs_modules()  # The poc module module must be in front of the plug-in module, and some parameters in the poc option call the plug-in
    _set_plugins()
    _init_targets_plugins()
    _init_pocs_plugins()
    _show_pocs_modules_options()
    _init_target_from_poc_dork()
    _set_task_queue()
    _init_results_plugins()

    if any((conf.url, conf.url_file, conf.plugins)):
        _set_http_cookie()
        _set_http_host()
        _set_http_referer()
        _set_http_user_agent()
        _set_http_extra_headers()

    _set_connect_back()
    _set_network_proxy()
    _set_network_timeout()
    _set_threads()
    _set_listener()
    remove_extra_log_message()
