# SPDX-License-Identifier: MPL-2.0
# -*- coding: utf-8 -*-

import abc
import glob
import importlib
import ipaddress
import pathlib
import pkgutil
import time
import urllib

from cryptoparser.common.base import Serializable
from cryptoparser.tls.subprotocol import TlsAlertDescription
from cryptoparser.ssh.subprotocol import SshKeyExchangeInit
from cryptoparser.ssh.version import SshProtocolVersion, SshVersion
from cryptoparser.common.utils import get_leaf_classes

from cryptolyzer.common.transfer import L4TransferSocketParams
from cryptolyzer.common.utils import LogSingleton
from cryptolyzer.dnsrec.client import L7ClientDnsBase
from cryptolyzer.ike.client import L7ClientIPsecBase
from cryptolyzer.httpx.client import L7ClientHttpBase
from cryptolyzer.ssh.client import L7ClientSsh
from cryptolyzer.tls.client import L7ClientTlsBase


class ProtocolHandlerBase(metaclass=abc.ABCMeta):
    @classmethod
    def import_plugins(cls):
        plugin_root_dir_parts = pathlib.PurePath(*pathlib.PurePath(__file__).parts[:-2])  # remove common/analyzer.py
        plugin_module_dir_parts = set()
        plugin_paths = filter(
            lambda path: path != __file__,
            glob.iglob(str(plugin_root_dir_parts / '*' / 'analyzer.py'))
        )
        for path in plugin_paths:
            plugin_path_parts = pathlib.PurePath(path).parts[-3:-1]  # split plugin dirs
            for index in range(len(plugin_path_parts)):
                plugin_module_dir_parts.add('.'.join(plugin_path_parts[:index + 1]))

        plugin_module_dir_parts = list(plugin_module_dir_parts)
        plugin_module_dir_parts.sort(key=len)
        for plugins_dir in plugin_module_dir_parts:
            ns_pkg = importlib.import_module(plugins_dir, package=None)
            for _, name, _ in pkgutil.iter_modules(ns_pkg.__path__, ns_pkg.__name__ + "."):
                if name.endswith('.analyzer'):
                    importlib.import_module(name)

    @classmethod
    def from_protocol(cls, protocol):
        cls.import_plugins()

        for handler_class in get_leaf_classes(cls):
            if handler_class.get_protocol() == protocol:
                return handler_class()
        raise KeyError(protocol)

    @classmethod
    def get_protocols(cls):
        cls.import_plugins()

        return sorted([
            handler_class.get_protocol()
            for handler_class in get_leaf_classes(cls)
        ])

    @classmethod
    @abc.abstractmethod
    def get_protocol(cls):
        raise NotImplementedError()

    @classmethod
    @abc.abstractmethod
    def get_analyzers(cls):
        raise NotImplementedError()

    @classmethod
    @abc.abstractmethod
    def _get_analyzer_args(cls):
        raise NotImplementedError()

    @classmethod
    def _l7_client_from_params(cls, uri, socket_params):
        kwargs = {'scheme': uri.scheme, 'address': uri.host, 'l4_socket_params': socket_params}

        if uri.port:
            kwargs['port'] = int(uri.port)
        if uri.fragment:
            try:
                ipaddress.ip_address(uri.fragment)
            except ValueError:
                pass
            else:
                kwargs['ip'] = uri.fragment

        if uri.query is not None:
            kwargs.update(dict(urllib.parse.parse_qsl(uri.query)))

        for analyzer_class in cls.get_analyzers():
            for client_class in analyzer_class.get_clients():
                if client_class.get_scheme() == uri.scheme:
                    return client_class.from_scheme(**kwargs)

        raise NotImplementedError()

    def analyze(self, analyzer, uri, socket_params=L4TransferSocketParams()):
        LogSingleton().log(
            level=60,
            msg=f'Analysis started; protocol="{self.get_protocol()}", analyzer="{analyzer.get_name()}"'
        )

        l7_client = self._l7_client_from_params(uri, socket_params)
        args, kwargs = self._get_analyzer_args()
        return analyzer.analyze(l7_client, *args, **kwargs)

    @classmethod
    def analyzer_from_name(cls, name):
        analyzer_list = [
            analyzer_class
            for analyzer_class in cls.get_analyzers()
            if analyzer_class.get_name() == name
        ]

        if len(analyzer_list) != 1:
            raise ValueError  # pragma: no cover

        return analyzer_list[0]()


class AnalyzerBase():
    def __init__(self):
        self._probe_attempt = 0

    @classmethod
    @abc.abstractmethod
    def get_name(cls):
        raise NotImplementedError()

    def analyze(self, analyzable):  # pylint: disable=unused-argument
        self._reset_probe_throttle()

    def _reset_probe_throttle(self):
        """Reset probe attempt counter at the start of a new analysis."""
        self._probe_attempt = 0

    def _before_probe(self, analyzable):
        """Call before each probe; sleeps (except before the first) when throttle_delay is set."""
        self._probe_attempt += 1
        if self._probe_attempt > 1:
            params = getattr(analyzable, 'l4_socket_params', None)
            delay = (getattr(params, 'throttle_delay', 0) or 0) if params else 0
            if delay > 0:
                time.sleep(delay)


class AnalyzerResultBase(Serializable):
    pass


class AnalyzerTlsBase(AnalyzerBase):
    _ACCEPTABLE_HANDSHAKE_FAILURE_ALERTS = [
        TlsAlertDescription.HANDSHAKE_FAILURE,  # no matching algorithms
        TlsAlertDescription.CLOSE_NOTIFY,  # no matching algorithms
        TlsAlertDescription.INSUFFICIENT_SECURITY,  # not enough secure matching algorithms
        TlsAlertDescription.ILLEGAL_PARAMETER  # unimplemented matching algorithms
    ]

    @classmethod
    @abc.abstractmethod
    def get_name(cls):
        raise NotImplementedError()

    @classmethod
    def get_clients(cls):
        return list(get_leaf_classes(L7ClientTlsBase))

    @classmethod
    def get_default_scheme(cls):
        return 'https'

    def analyze(self, analyzable, protocol_version):  # pylint: disable=arguments-differ,unused-argument
        self._reset_probe_throttle()


class ProtocolHandlerTlsBase(ProtocolHandlerBase):
    @classmethod
    @abc.abstractmethod
    def get_protocol_version(cls):
        raise NotImplementedError()

    @classmethod
    def get_protocol(cls):
        return cls.get_protocol_version().identifier

    @classmethod
    def _get_analyzer_args(cls):
        return ([], {'protocol_version': cls.get_protocol_version()})


class ProtocolHandlerTlsExactVersion(ProtocolHandlerTlsBase):
    @classmethod
    @abc.abstractmethod
    def get_protocol_version(cls):
        raise NotImplementedError()


class ProtocolHandlerSshBase(ProtocolHandlerBase):
    @classmethod
    def get_protocol(cls):
        return SshProtocolVersion(SshVersion.SSH2).identifier

    @classmethod
    def _get_analyzer_args(cls):
        return ([], {})

    @classmethod
    @abc.abstractmethod
    def get_analyzers(cls):
        raise NotImplementedError()


class ProtocolHandlerSshExactVersion(ProtocolHandlerSshBase):
    @classmethod
    @abc.abstractmethod
    def get_protocol_version(cls):
        raise NotImplementedError()


class ProtocolHandlerIKEBase(ProtocolHandlerBase):
    @classmethod
    @abc.abstractmethod
    def get_protocol_version(cls):
        raise NotImplementedError()

    @classmethod
    def get_protocol(cls):
        return cls.get_protocol_version().identifier  # pragma: no cover

    @classmethod
    def _get_analyzer_args(cls):
        return ([], {'protocol_version': cls.get_protocol_version()})


class ProtocolHandlerIKEExactVersion(ProtocolHandlerIKEBase):
    """Base class for IKE protocol handlers with exact version."""
    @classmethod
    def get_protocol_version(cls):
        raise NotImplementedError()


class AnalyzerSshBase(AnalyzerBase):
    _KEXINIT_CIPHER_ATTRIBUTES = (
        'kex_algorithms',
        'host_key_algorithms',
        'encryption_algorithms_client_to_server',
        'encryption_algorithms_server_to_client',
        'mac_algorithms_client_to_server',
        'mac_algorithms_server_to_client',
        'compression_algorithms_client_to_server',
        'compression_algorithms_server_to_client',
    )

    @classmethod
    @abc.abstractmethod
    def get_name(cls):
        raise NotImplementedError()

    @classmethod
    def get_clients(cls):
        return list(get_leaf_classes(L7ClientSsh))

    @classmethod
    def get_default_scheme(cls):
        return 'ssh'

    @staticmethod
    def _get_known_algorithms(algorithms):
        return filter(lambda algorithm: not isinstance(algorithm, str), algorithms)

    @classmethod
    def _build_limited_key_exchange_init_message(
            cls,
            key_exchange_init,
            analyzer_result,
            keep_from_template=(),
            server_overrides=None,
    ):
        if isinstance(key_exchange_init, type):
            key_exchange_init_message = key_exchange_init()
        else:
            key_exchange_init_message = key_exchange_init

        keep_from_template = set(keep_from_template)
        server_overrides = {} if server_overrides is None else dict(server_overrides)

        init_kwargs = {}
        for attribute_name in cls._KEXINIT_CIPHER_ATTRIBUTES:
            if attribute_name in server_overrides:
                init_kwargs[attribute_name] = server_overrides[attribute_name]
            elif attribute_name in keep_from_template:
                init_kwargs[attribute_name] = getattr(key_exchange_init_message, attribute_name)
            else:
                init_kwargs[attribute_name] = list(cls._get_known_algorithms(getattr(analyzer_result, attribute_name)))

        SshKeyExchangeInit.__init__(key_exchange_init_message, **init_kwargs)
        return key_exchange_init_message

    @abc.abstractmethod
    def analyze(self, analyzable):
        self._reset_probe_throttle()


class AnalyzerHttpBase(AnalyzerBase):
    @classmethod
    @abc.abstractmethod
    def get_name(cls):
        raise NotImplementedError()

    @classmethod
    def get_clients(cls):
        return list(get_leaf_classes(L7ClientHttpBase))

    @classmethod
    def get_default_scheme(cls):
        return 'https'

    def analyze(self, analyzable, protocol_version):  # pylint: disable=arguments-differ,unused-argument
        self._reset_probe_throttle()


class AnalyzerDnsRecordBase(AnalyzerBase):
    @classmethod
    @abc.abstractmethod
    def get_name(cls):
        raise NotImplementedError()

    @classmethod
    def get_clients(cls):
        return list(get_leaf_classes(L7ClientDnsBase))

    @classmethod
    def get_default_scheme(cls):
        return 'dns'

    def analyze(self, analyzable):
        self._reset_probe_throttle()


class AnalyzerIKEBase(AnalyzerBase):
    _MAX_PROPOSALS_PER_INIT_MESSAGE = 128

    @classmethod
    @abc.abstractmethod
    def get_name(cls):
        raise NotImplementedError()

    @classmethod
    def get_clients(cls):
        return [
            L7ClientIPsecBase
        ]

    @classmethod
    def get_default_scheme(cls):
        return 'ipsec'

    def analyze(self, analyzable, protocol_version):  # pylint: disable=arguments-differ,unused-argument
        self._reset_probe_throttle()
