import queue
import socket
import threading


TIMEOUT = 1
MAX_THREADS = 64


class Scanner(object):
    def __init__(self, host):
        self.host = host

        # The job queue contains tuples of the form (protocol, port). We might
        # be able to remove this queue by using a lock in _get_next_port(), but
        # that might not be as optimized as accessing a Queue. Since we are too
        # lazy to write and run a benchmark, let's just do this.
        self.job_queue = queue.Queue()
        for job in self._get_next_port():
            self.job_queue.put(job)
        self.opened_ports_queue = queue.Queue()

    def _get_next_port(self):
        """Return the next (protocol, port) tuple to test."""
        # TODO: Improve this function :)
        # TODO: There are probably more "common" ports used for VPNs.
        self.common_ports = [80, 443, 8080]

        for port in self.common_ports:
            yield ('udp', port)

        for port in self.common_ports:
            yield ('tcp', port)

        # TODO: if necessary, check all possible ports

    def _compute_n_threads(self):
        """Compute the number of threads that can be used to scan ports.

        This is not as trivial as it seems: we need as much parallelism as
        possible, but we may not be able to open too many connections at once
        on some networks. Also, we do not want to overload the machine.
        """
        # TODO: Replace this stub by a real function.
        return 4

    def scan(self, n_threads=0):
        if n_threads <= 0 or n_threads > MAX_THREADS:
            n_threads = self._compute_n_threads()

        threads = []
        for i in range(n_threads):
            t = threading.Thread(target=self._scan)
            threads.append(t)
            t.start()

        self.job_queue.join()
        for t in threads:
            t.join()

        # TODO: We could probably start yielding the results while the workers
        # are scanning the network, but that would make the code a bit more
        # complex and would be more error-prone. Being a bit slow seems OK for
        # now :)
        while True:
            try:
                yield self.opened_ports_queue.get(block=False)
            except queue.Empty:
                break
            else:
                self.opened_ports_queue.task_done()

    def _scan(self):
        while True:
            try:
                protocol, port = self.job_queue.get(timeout=1)
                if protocol == 'udp':
                    opened = self._scan_udp(port)
                elif protocol == 'tcp':
                    opened = self._scan_tcp(port)
                else:
                    raise ValueError('Unknown protocol %s' % protocol)

                if opened:
                    self.opened_ports_queue.put((protocol, port))

                self.job_queue.task_done()
            except queue.Empty:
                break

    def _scan_tcp(self, port):
        """Return True if PORT is opened in TCP."""
        try:
            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            s.settimeout(TIMEOUT)
            s.connect((self.host, port))
            s.send(b'x')
            s.recvfrom(1)
        except (ConnectionRefusedError, socket.timeout):
            return False
        finally:
            s.close()

        return True

    def _scan_udp(self, port):
        """Return True if PORT is opened in UDP."""
        try:
            s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            s.settimeout(TIMEOUT)
            # The message is a valid payload sent to OpenVPN in order to
            # trigger an answer from the server, and is computed like this:
            # msg = '38 36 d9 fc 5f e3 fe f7 09 00 00 00 00'
            # msg = bytes.fromhex(msg.replace(' ', ''))
            msg = b'86\xd9\xfc_\xe3\xfe\xf7\t\x00\x00\x00\x00'
            s.sendto(msg, (self.host, port))
            s.recvfrom(1)
        except socket.timeout:
            return False

        return True