import ipaddress
import overseer
from threading import Thread
from io import open
import socket
import os
import re
import time


class ScanManager:
    """Used to manage any ongoing scans

    Methods
    -------
    start()
        Start the scan monitor thread
    stop()
        Stop and cleanup the scan monitor thread
    get_status()
        Get a normalized list of dicts detailing outstanding scans
    perform_scan(target)
        Initiate a scan on target (Can be a hostname, or ip_addr)
    is_ip(target)
        Determines if a target is a hostname of ip_addr
    """

    def __init__(self):
        """Create instance and start thread"""
        self.__pending_shutdown = False
        self.__active_scans = []

    def __scan_monitor(self):
        """Monitors active and completed scans

        Responsible for publishing status to the websockets, cleaning up
        completed scans, and updating the database accordingly.
        """
        while not self.__pending_shutdown:
            time.sleep(1)
            if len(self.__active_scans) == 0:
                continue

            for scan in self.__active_scans:
                # WebSocket progress
                total_progress = (scan.tcp_progress + scan.udp_progress) / 2
                results = scan.get_results()
                overseer.api.send_websocket_event(
                    {
                        "id": scan.scan_history.id,
                        "target": scan.target,
                        "status": "IN_PROGRESS",
                        "results": results,
                        "created_at": scan.scan_history.created_at.isoformat(),
                        "error": scan.scan_history.error,
                        "tcp_progress": scan.tcp_progress,
                        "udp_progress": scan.udp_progress,
                        "total_progress": round(total_progress),
                    }
                )

                if scan.tcp_progress + scan.udp_progress != 200:
                    continue

                # Update database
                scan_history = overseer.database.update_scan_result(
                    scan.scan_history.id, "COMPLETE", results=results
                )

                # WebSocket completion
                overseer.api.send_websocket_event(
                    {
                        "id": scan_history.id,
                        "target": scan.target,
                        "status": scan_history.status,
                        "results": results,
                        "created_at": scan_history.created_at.isoformat(),
                        "error": scan_history.error,
                    }
                )

                # Cleanup active scan
                scan.join()
                self.__active_scans.remove(scan)

    def start(self):
        self.__monitor_thread = Thread(target=self.__scan_monitor)
        self.__monitor_thread.start()

    def stop(self):
        """Shutdown and cleanup the scan monitor thread"""
        self.__pending_shutdown = True
        self.__monitor_thread.join()

    def get_status(self):
        """Get a normalized list of dicts detailing outstanding scans

        Returns
        -------
        list
            List of normalized details on current active scans
        """
        return list(
            map(
                lambda x: {
                    "id": x.scan_history.id,
                    "target": x.target,
                    "status": "IN_PROGRESS",
                    "results": x.get_results(),
                    "created_at": x.scan_history.created_at.isoformat(),
                    "error": x.scan_history.error,
                    "tcp_progress": x.tcp_progress,
                    "udp_progress": x.udp_progress,
                    "total_progress": round(
                        (x.tcp_progress + x.udp_progress) / 2
                    ),  # noqa: E501
                },
                self.__active_scans,
            )
        )

    def perform_scan(self, target):
        """Initiate a scan on target (Can be a hostname, or ip_addr)

        Parameters
        ----------
        target : str
            Either a hostname or IP address of the endpoint to scan
        """
        try:
            socket.gethostbyname(target)
        except socket.error:
            return None

        if self.is_ip(target):
            scan_history = overseer.database.create_scan_result(
                "IN_PROGRESS", ip_addr=target
            )
        else:
            scan_history = overseer.database.create_scan_result(
                "IN_PROGRESS", hostname=target
            )

        new_scan = Scanner(target, scan_history)
        new_scan.start()
        self.__active_scans.append(new_scan)
        return scan_history

    def is_ip(self, target):
        """Determines if a target is a hostname of ip_addr

        Parameters
        ----------
        target : str
            Either a hostname or IP address of the endpoint to scan

        Returns
        -------
        bool
            Whether the target is an IP or not
        """
        try:
            ipaddress.ip_address(target)
            return True
        except ValueError:
            return False


class Scanner(Thread):
    """Subclass of Thread - used to perform a TCP and UDP scan

    Attributes
    ----------
    target : str
        The hostname or ip_addr to scan
    scan_history : ScanHistory
        The ScanHistory DB model reference for this scan
    tcp_progress : int
        The current progress percentage of the threaded TCP scan
    udp_progress : int
        The current progress percentage of the threaded UDP scan
    tcp_results : list
        The current list if ints of open TCP ports
    udp_results : list
        The current list if ints of open UDP ports

    Methods
    -------
    run()
        Overridden run method from the Thread superclass. Starts the scan.
    get_results()
        Returns a normalized list of string of open ports
        (E.g. ["53 UDP", "53 TCP"])
    """

    def __init__(self, target, scan_history):
        """
        Parameters
        ----------
        target : str
            The hostname or ip_addr to scan
        scan_history : ScanHistory
            The ScanHistory DB model reference for this scan
        """
        Thread.__init__(self)
        self.target = target
        self.scan_history = scan_history

        self.__port_count = 1000

        self.tcp_progress = 0
        self.udp_progress = 0

        self.tcp_results = []
        self.udp_results = []

        self.__udp_payloads = {}
        self.__load_nmap_payloads()

    def __load_nmap_payloads(self):
        """Load and parse nmap UDP payloads

        Because of how UDP is designed, we have to test with specific payloads.
        This parses nmaps protocol specific payloads [0] in preperation.

        TODO: This should be cached. No need to parse it on every scan.

        [0] https://nmap.org/book/nmap-payloads.html
        """

        # Open file & remove comments
        nmap_payloads = os.path.join(
            os.path.dirname(__file__), "./resources/nmap-payloads"
        )
        f = open(nmap_payloads, "r", encoding="unicode_escape")
        raw_file = re.sub(
            r"^(#|\s*#).*$", "", f.read(), flags=re.MULTILINE | re.UNICODE
        )
        f.close()

        # Find all matches
        results = re.findall(
            r"(^udp.*?)(?=udp|\Z)",
            raw_file,
            flags=re.MULTILINE | re.DOTALL | re.UNICODE,
        )

        for raw_match in results:
            match = raw_match.strip()
            match_payloads = list(
                map(lambda x: x.strip(), re.findall(r'\"(.*)"', match))
            )
            raw_ports = re.match(r"^udp\s([\d,-]*)", match).group(1)

            for raw_port in raw_ports.split(","):
                if "-" in raw_port:
                    port_range = raw_port.split("-")
                    start_port = int(port_range[0])
                    end_port = int(port_range[1])
                    for port_match in range(start_port, end_port + 1):
                        if port_match not in self.__udp_payloads:
                            self.__udp_payloads[port_match] = []
                        self.__udp_payloads[port_match].extend(match_payloads)
                else:
                    port_match = int(raw_port)
                    if port_match not in self.__udp_payloads:
                        self.__udp_payloads[port_match] = []
                    self.__udp_payloads[port_match].extend(match_payloads)

    def run(self):
        """Overridden run method from the Thread superclass. Starts the scan"""
        tcp_thread = Thread(target=self.__scan_tcp)
        udp_thread = Thread(target=self.__scan_udp)
        tcp_thread.start()
        udp_thread.start()
        tcp_thread.join()
        udp_thread.join()

    def get_results(self):
        """Returns a normalized list of string of open ports

        Returns
        -------
        list
            List of open ports (E.g. ["53 UDP", "53 TCP"])
        """

        results = list(map(lambda x: "%s UDP" % x, self.udp_results))
        results.extend(
            list(map(lambda x: "%s TCP" % x, self.tcp_results))
        )  # noqa: E501
        results.sort(key=lambda x: int(x.split(" ")[0]))
        return results

    def __scan_tcp(self):
        """Threaded TCP Scanner"""
        for port in range(1, self.__port_count):
            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            s.settimeout(0.1)
            result = s.connect_ex((self.target, port))
            if result == 0:
                self.tcp_results.append(port)
            s.close()
            self.tcp_progress = round(port / self.__port_count * 100)

    def __scan_udp(self):
        """Threaded UDP Scanner"""
        for port in range(1, self.__port_count):
            s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            s.settimeout(0.01)
            payloads = ["\x00"]
            if port in self.__udp_payloads:
                payloads.extend(self.__udp_payloads[port])
            for payload in payloads:
                s.sendto(payload.encode("utf-8"), (self.target, port))
                try:
                    s.recvfrom(1)
                    self.udp_results.append(port)
                    break
                except ConnectionRefusedError:
                    pass
                except socket.timeout:
                    pass
            s.close()
            self.udp_progress = round(port / self.__port_count * 100)
