import ipaddress
from typing import Optional, List, Tuple
from dataclasses import dataclass, field
from dataclasses_json import dataclass_json
from pprint import pprint
import json
from loguru import logger

from helpers.custom_logging import FlashLog
from pwndoc_api import download_and_return_audit, update_scope

TNS_SCOPE_TABLE_OF_SERVICES = 'TNSTableOfServices'


@dataclass_json
@dataclass
class Service:
    port: int
    protocol: str = 'tcp'
    name: str = ''
    product: str = 'n/a'
    version: str = 'n/a'
    _id: Optional[str] = None


@dataclass_json
@dataclass
class Host:
    services: List[Service] = field(default_factory=lambda: [])
    os: str = 'n/a'
    hostname: str = 'n/a'
    ip: str = 'n/a'
    _id: Optional[str] = None

    def sort(self):
        self.services.sort(key=lambda x: int(x.port))


@dataclass_json
@dataclass
class Scope:
    hosts: List[Host] = field(default_factory=lambda: [])
    name: str = 'n/a'

    def add_entry(self, ip: str, port_number: str, l3_protocol: str, app_protocol: str):
        host = list(filter(lambda x: x.ip == ip, self.hosts))
        if len(host) == 0:
            host = Host(ip=ip)
            self.hosts.append(host)
        else:
            host = host[0]

        service = list(filter(lambda x: x.port == int(port_number) and x.protocol == l3_protocol, host.services))
        if len(service) == 0:
            service = Service(port=int(port_number), protocol=l3_protocol, name=app_protocol)
            host.services.append(service)
        else:
            service = service[0]
            logger.debug("Duplicate service in open ports table. Skipping.")

    def sort(self):
        def try_to_get_ip(x: str):
            try:
                return ipaddress.get_mixed_type_key(ipaddress.ip_address(x))
            except:
                return x

        try:
            self.hosts.sort(key=lambda x: try_to_get_ip(x.ip))
        except:
            logger.warning("Sorting by IPs failed.")

        for host in self.hosts:
            try:
                host.sort()
            except:
                logger.warning("Sorting by ports failed")


@dataclass_json
@dataclass
class AllScopes:
    scope: List[Scope] = field(default_factory=lambda: [])


def get_tns_table_of_services(audit_id: str) -> Scope:
    audit_raw = download_and_return_audit(audit_id)
    if audit_raw is None:
        logger.warning(f"Audit {audit_id} doesn't have a valid scope section? Pretending it's empty.")
        return Scope(name=TNS_SCOPE_TABLE_OF_SERVICES)

    all_scopes_raw = audit_raw.get("datas", {}).get("scope", [])
    all_scopes = Scope.schema().load(all_scopes_raw, many=True)

    tns_scope_dict_list = list(filter(lambda x: x.name == TNS_SCOPE_TABLE_OF_SERVICES, all_scopes))
    return tns_scope_dict_list[0] if len(tns_scope_dict_list) > 0 else Scope(name=TNS_SCOPE_TABLE_OF_SERVICES)


def set_tns_table_of_services(audit_id: str, new_scope: Scope):
    audit_raw = download_and_return_audit(audit_id)
    scope_raw = audit_raw.get("scope", [])
    if len(scope_raw) > 1:
        # todo: fix this
        FlashLog.error("Importer does not support PwnDoc audits which contain Scope. It's being used for TNS Table of Services. Overwriting.")
    new_scope.sort()
    new_scope_dict = {
        # '_id': 'TODO',
        'scope': new_scope.to_dict(),
    }
    update_scope(audit_id, new_scope_dict)


if __name__ == "__main__":
    service1 = Service(port=25, name='Service1')
    host1 = Host(ip='192.168.1.1', services=[service1], hostname='Host1')
    scope1 = Scope(name='Scope1', hosts=[host1])
    all_scope = AllScopes([scope1])

    pprint(all_scope.to_dict())

    with open("test1.out.json", "w", encoding="utf8") as f:
        json.dump(all_scope.to_dict(), f, indent=4)

    with open("test1.out.json", "r", encoding="utf8") as f:
        test1 = json.load(f)
        test2 = AllScopes.from_dict(test1)
        pprint(test2.to_dict())
