"""
Provides exploit related classes

"""

from collections import defaultdict
import csv
from enum import Enum

from tools.nmap.protocols_list import NMAP_PROTOCOLS
from tools.nmap.services_list import NMAP_SERVICES
from utils.exceptions import ServiceUnsupportedByNmapException, PortUnsupportedException, \
    ProtocolUnsupportedByNmapException, PortRangeUnsupported


class RiskLevel(Enum):
    """
    Risk level object

    """

    def __init__(self, txt, number):
        self.txt = txt
        self.number = number

    HIGH = ('High', 3)
    MEDIUM = ('Medium', 2)
    LOW = ('Low', 1)
    NONE = ('None', 0)

    @classmethod
    def from_name(cls, name):
        """
        Create RiskLevel object basing on string name

        Args:
            name: string representation of risk level, eg. "medium"

        Returns:
            RiskLevel object

        Raises:
            ValueError if not: High, Medium, Low or None

        """
        for val in cls:
            if val.txt == name:
                return val
        raise ValueError('Unsupported risk level name: %s' % name)


class ExploitCategory(Enum):
    """
    Exploits categories

    """
    VULN = 'vuln'
    BACKDOOR = 'backdoor'
    ENUM = 'enum'
    CREDENTIAL = 'credential'


class ExploitScanType(Enum):
    BASIC = 'basic'
    ADVANCED = 'advanced'


class ExploitTag(Enum):
    def __init__(self, mask, text):
        self.mask = mask
        self.text = text

    OTHER = (1 << 0, 'other')
    SSL = (1 << 1, 'ssl')
    TLS = (1 << 2, 'tls')
    HTTP = (1 << 3, 'http')
    HTTPS = (1 << 4, 'https')
    SAMBA = (1 << 5, 'samba')
    SSH = (1 << 6, 'ssh')
    SQL = (1 << 7, 'sql')
    LDAP = (1 << 8, 'ldap')
    VNC = (1 << 9, 'vnc')
    RPC = (1 << 10, 'rpc')
    NTP = (1 << 11, 'ntp')
    SMTP = (1 << 12, 'smtp')
    AFP = (1 << 13, 'afp')
    AMQP = (1 << 14, 'amqp')
    PRINTER = (1 << 15, 'printer')
    IMAP = (1 << 16, 'imap')
    POP3 = (1 << 17, 'pop3')
    APMI = (1 << 18, 'ipmi')
    TELNET = (1 << 19, 'telnet')
    SNMP = (1 << 20, 'snmp')
    MONGODB = (1 << 21, 'mongodb')
    DNS = (1 << 22, 'dns')
    PROXY = (1 << 23, 'proxy')
    XMPP = (1 << 24, 'xmpp')
    FTP = (1 << 25, 'ftp')
    IPMI = (1 << 26, 'ipmi')


class ExploitMetric(Enum):
    """
    Exploits metrics

    """

    AFP = 'AFP'
    AJP = 'AJP'
    AMQP = 'AMQP'
    BROADCAST = 'Broadcast'
    CASSANDRA = 'Cassandra'
    CITRIX = 'Citrix'
    COUCHDB = 'CouchDB'
    CUPS = 'CUPS service and queue info'
    CVE_SEARCH = 'CVE-Search'
    DISTCC = 'distcc'
    DNS = 'DNS SRV enumeration and zone transfer'
    DOCKER = 'Docker'
    EAP = 'EAP'
    FTP = 'FTP'
    HADOOP = 'Hadoop'
    HTTP_CISCO_ASA = 'Cisco ASA HTTP management application'
    HTTP_CREDS = 'Using default login credentials'
    HTTP_FINGERPRINT = 'Web application name and version detection'
    HTTP_HEADERS = 'HTTP headers - if they are properly configured or not'
    HTTP_INFO = 'HTTP vulnerabilities'
    HTTP_PATH_ENUM = 'Path enumeration'
    HTTP_PROXY = 'HTTP Proxy'
    HTTP_SLOWLORIS = 'Slowloris Dos attack'
    HTTP_VULNS = 'Specific HTTP vulnerabilities (for PHP Ruby server IIS Apache JBoss Websphere and so on)'
    HYDRA = 'Hydra'
    IKE = 'IKE'
    IMAP = 'IMAP info'
    IPMI_BRUTE = 'Bruteforce'
    IPMI_DUMP = 'Dump password hashes'
    IPMI_OTHER = 'IPMI'
    IPMI_ZERO = 'Cipher Zero Authentication Bypass'
    IRC = 'IRC'
    ISCSI = 'ISCSI'
    LDAP = 'LDAP searching and retrieving DSE'
    MAIL_HEAP_OFLOW = 'Heap overflow'
    MAIL_MEM_CORRUPT = 'Memory corruption'
    MAIL_POP3 = 'POP3'
    MAIL_POSTFIX = 'Postfix closing'
    MEMCACHEDB = 'Memcache db'
    MONGODB = 'Mongodb information and databases'
    MSSQL_DUMP = 'Dump passwords and determine configuration from MS SQL'
    MYSQL = 'MySQL'
    NFS = 'NFS exports information'
    NTP = 'NTP information'
    PPTP = 'PPTP'
    QUAKE = 'Quake'
    RDP = 'RDP encryption and MS12-02 vulnerability'
    REDIS = 'Redis information'
    RMI = 'RMI'
    SAMBA_BUFF_OFLOW = 'Buffer overflow (ms06-025 ms07-029)'
    SAMBA_ENUM = 'Domains shares and session enumeration'
    SAMBA_OTHER = 'Samba'
    SAMBA_RCE = 'Remote code execution (ms08-067 WannaCry)'
    SKIPFISH = 'Skipfish'
    SNMP_CISCO = 'Obtaining Cisco router configuration'
    SNMP_ENUM = 'Enumerating Windows shares users services and installed software'
    SNMP_STATS = 'Getting netstat and processes from device'
    SSH_CCS = 'CCS injection'
    SSH_CERT = 'Certificate validation'
    SSH_ENUM = 'Algorithms enumeration'
    SSH_SUPPORTED = 'Checking for supported SSL versions and Diffie-Hellman groups'
    SSL_HEARTBLEED = 'Heartbleed'
    SSL_POODLE = 'Poodle'
    TNS = 'TNS'
    UPNP = 'UPNP'
    VNC_BYPASS = 'Bypassing RealVNC'
    VNC_INFO = 'RPC and VNC information'
    WINDOWS = 'Windows exploitation'
    XMPP = 'XMPP information'
    OTHER = ''


class Exploit(object):
    """
    Represents Exploit

    """
    def __init__(self, exploit_id, app=None, name=None, title=None, description=None, risk_level=None, services=None,
                 ports=None, category=None, cve=None, metric=None, scan_type=None, tags=None, cvss=None):
        """

        Args:
            exploit_id (int): exploit id
            app (str): app name, e.g. nmap
            name (str): exploit name, e.g. http-enum
            title (str): title of exploit
            description(str): short description of exploit
            risk_level (RiskLevel): risk level
            services (list): supported services
            ports (list): supported ports
            category (set): exploit's category
            tags (set): exploit's tags

        """
        self.id = exploit_id
        self.app = app
        self.name = name
        self.title = title
        self.description = description
        self.risk_level = risk_level
        self.services = services or set()
        self.ports = ports or {}
        self.category = category or set()
        self.metric = metric if metric is not None else ExploitMetric.OTHER
        self._cve = []
        self.cve = cve
        self.scan_type = scan_type
        self.tags = tags or set()
        self._cvss = []
        self.cvss = cvss

    @property
    def cvss(self):
        return max([0, *self._cvss])

    @cvss.setter
    def cvss(self, value):
        if not value:
            self._cvss = []
            return

        if isinstance(value, (float, int)):
            self._cvss = [float(value)]
            return
        self._cvss = list(map(float, value.split(',')))

    @property
    def cve(self):
        return ",".join(self._cve)

    @cve.setter
    def cve(self, value):
        if value is None:
            self._cve = []
            return

        self._cve = value.split(',')

    def __eq__(self, other):
        return isinstance(other, Exploit) and self.id == other.id

    def __ne__(self, other):
        return (not isinstance(other, Exploit)) or self.id != other.id

    def __hash__(self):
        return hash(self.id)

    def __str__(self):
        return "{0}:{1}[{2}]".format(self.app, self.name, self.id)

    @property
    def tags_mask(self):
        """
        Returns an int which is a mask of different tag bits. See ExploitTag possible values

        """
        return_value = 0

        for tag in self.tags:
            return_value |= tag.mask

        return return_value


class Exploits(object):
    """
    Collection of exploits

    """

    def __init__(self):
        self._exploits = defaultdict(dict)
        self._all_exploits = set()

    def __iter__(self):
        """
        Implements iteration over exploits in collection

        Returns:
            Exploit

        """
        return iter(self._all_exploits)

    def __len__(self):
        return len(self._all_exploits)

    def add(self, exploit):
        """
        Add exploit to collection

        Args:
            exploit (Exploit):

        Returns:
            None
        """
        self._exploits[exploit.app][exploit.name] = exploit
        self._all_exploits.add(exploit)

    def find(self, app, name):
        """
        Find all exploits of provided name and app name

        Args:
            app (str):
            name (str):

        Returns:
            Exploit

        """
        return self._exploits[app].get(name)

    def find_all_matching(self, port=None):
        """
        Find all exploits meet the criteria

        Args:
            port (Port): Port, which we want to exploit

        Returns:
            list: all exploits for which:
                - app is equal
                - there is no port or service is equal or transport protocol and port is equal
                - there is no supported ports and services (exploit support all ports and services)

        """
        return_value = {}
        for exploit in self._all_exploits:
            if not port or port.protocol in exploit.services or \
                    (port.transport_protocol.db_val in exploit.ports.keys() and
                     port.number in exploit.ports[port.transport_protocol.db_val])or \
                    (not exploit.ports and not exploit.services):

                return_value.setdefault(exploit.app, []).append(exploit)

        return return_value

    def find_all_phy(self):
        """
        Finds all exploits for physical interface scanning

        Returns:
            list

        """
        return_value = {}
        for exploit in self._all_exploits:
            if "PHY" in exploit.ports.keys():
                return_value.setdefault(exploit.app, []).append(exploit)
        return return_value

    def remove(self, exploit):
        """
        Removes exploit from collection

        Args:
            exploit (Exploit):

        Returns:
            None

        """
        self._all_exploits.remove(exploit)
        self._exploits.get(exploit.app, {}).pop(exploit.name, None)

    @classmethod
    def read(cls, file_name):
        """
        Reads csv exploits database, and return collection of exploits

        Args:
            file_name:

        Returns:
            Exploits: collection of exploits

        """
        result = cls()

        with open(file_name, newline='', encoding='utf-8') as csvfile:
            reader = csv.reader(csvfile, delimiter=';', quotechar='"')
            for num, row in enumerate(reader):
                if num == 0:
                    continue  # header
                exploit = Exploit(exploit_id=int(row[0]))
                exploit.app = row[1]
                exploit.name = row[2]
                exploit.title = row[3]
                exploit.description = row[4]
                exploit.risk_level = RiskLevel.from_name(row[5])
                exploit.services = cls.parse_services(row[6])
                exploit.ports = cls.parse_ports(row[7])
                exploit.category = ExploitCategory[row[8].upper()]
                exploit.cve = row[9]
                exploit.cvss = row[10]
                exploit.metric = ExploitMetric[row[11].upper()]
                exploit.scan_type = ExploitScanType[row[12].upper()]
                exploit.tags = cls.parse_tags(row[13])
                result.add(exploit)
        return result

    @classmethod
    def parse_services(cls, text):
        """
        This function parses service string and returns list of services

        Args:
            text (str):

        Returns:
            list

        """
        if not text:
            return set()

        return_value = set()
        for service in text.split(','):
            service_name = service.strip()
            if service_name not in NMAP_SERVICES:
                raise ServiceUnsupportedByNmapException(service_name)
            return_value.add(service_name)

        return return_value

    @classmethod
    def parse_ports(cls, text):
        """
        This function parses port string and returns list of Ports

        Args:
            text (str):

        Returns:
            list

        """

        if not text:
            return {}

        return_value = {}
        protocols = text.split(",")
        ports = [protocol.split(":") for protocol in protocols]

        try:
            for protocol, port_set in ports:
                protocol = protocol.strip()
                if protocol not in NMAP_PROTOCOLS:
                    raise ProtocolUnsupportedByNmapException(protocol)
                return_value[protocol] = set()

                all_ports = port_set.split(".")
                for port in all_ports:
                    if "-" in port:
                        port_range = [int(port_number) for port_number in port.split("-")]
                        if port_range[0] >= port_range[1]:
                            raise PortRangeUnsupported(port)

                        return_value[protocol].update(range(int(port_range[0]), int(port_range[1]) + 1))
                        continue
                    return_value[protocol].add(int(port))
        except ValueError:
            raise PortUnsupportedException(text)

        return return_value

    @classmethod
    def parse_tags(cls, text):
        """

        Args:
            text:

        Returns:
            set

        """
        if not text:
            return set()

        return {ExploitTag[cat.upper()] for cat in text.split(",")}

    def find_by_apps(self, apps):
        """
        Find all exploits for given apps list

        Args:
            apps (list):

        Returns:
            dict: {app: exploits, ...)

        """
        return {app: list(self._exploits.get(app, {}).values()) for app in apps}
