import nmap
import masscan
from threading import Thread
from datetime import datetime
from pprint import pprint


class AnmapThread(Thread):
    def __init__(self, hostname, ports, verbose, out):
        Thread.__init__(self)
        self.host = hostname
        self.ports = ports
        self.scanner = nmap.PortScanner()
        self.verbose = verbose
        self.daemon = True
        self.out = out


class ThoroughAnmapThread(AnmapThread):
    def run(self):
        log("Starting thorough scan on " + self.host, self.verbose)
        self.scanner.scan(self.host, "1," + ",".join(self.ports),
                          arguments='-sSVC -A -Pn{}'.format(output(self.out, self.host, 2)))
        log(self.scanner.command_line(), self.verbose)
        if self.out:
            with open(output(True, self.host, 5), "w") as outfile:
                outfile.write(self.scanner.get_nmap_last_output())
        host = self.scanner[self.host]
        log("{}/{} is {}".format(host.hostname(), host["addresses"]["ipv4"], host["osmatch"][0]["name"]), self.verbose)
        for p in host.all_tcp():
            if p == 1:
                continue
            log("Port {}/tcp: {}".format(p, host['tcp'][p]), self.verbose)
        log("Finished thorough scan on " + self.host, self.verbose)


class UDPAnmapThread(AnmapThread):
    def run(self):
        log("Starting UDP scan on " + self.host, self.verbose)
        self.scanner.scan(self.host, arguments='-sVCU -A -Pn --top-ports {}{}'.
                          format(self.ports, output(self.out, self.host, 3)))
        log(self.scanner.command_line(), self.verbose)
        if self.out:
            with open(output(True, self.host, 6), "w") as outfile:
                outfile.write(self.scanner.get_nmap_last_output())
        host = self.scanner[self.host]
        log("{}/{} is {}".format(host.hostname(), host["addresses"]["ipv4"], host["osmatch"][0]["name"]), self.verbose)
        for p in host.all_udp():
            log("Port {}/udp: {}".format(p, host['udp'][p]), self.verbose)
        log("Finished UDP scan on " + self.host, self.verbose)


class BaseAnmapThread(AnmapThread):
    def __init__(self, hostname, ports, verbose, out):
        AnmapThread.__init__(self, hostname, ports, verbose, out)
        self.host_dict = dict()

    def run(self):
        log("Starting quick scan", self.verbose)
        self.scanner.scan(self.host, arguments='-sS -Pn -p{}{}'.format(self.ports, output(self.out, self.host, 1)))
        log(self.scanner.command_line(), self.verbose)
        if self.out:
            with open(output(True, self.host, 4), "w") as outfile:
                outfile.write(self.scanner.get_nmap_last_output())
        log("Finished quick scan", self.verbose)
        np = 0
        for hostname in self.scanner.all_hosts():
            host = self.scanner[hostname]
            port_list = list()
            for p in host.all_tcp():
                if self.scanner[hostname]['tcp'][p]['state'] == 'open':
                    port_list.append(str(p))
            if port_list is not list():
                self.host_dict[hostname] = port_list
                np += len(port_list)
        log("Found {} open ports on {} host(s) with {}".format(np, len(self.host_dict), "nmap"), self.verbose)

    def rjoin(self):
        Thread.join(self)
        return self.host_dict


class MasscanAnmapThread(BaseAnmapThread):
    def __init__(self, hostname, ports, verbose, out):
        AnmapThread.__init__(self, hostname, ports, verbose, out)
        self.host_dict = dict()
        self.scanner = masscan.PortScanner()

    def run(self):
        log("Starting masscan scan", self.verbose)
        self.scanner.scan(self.host, ports=self.ports, arguments=output(self.out, self.host, 7), sudo=True)
        log(self.scanner.command_line(), self.verbose)
        log("Finished quick scan", self.verbose)
        np = 0
        for hostname in self.scanner.all_hosts():
            host = self.scanner[hostname]
            port_list = list()
            for p in host.all_tcp():
                if self.scanner[hostname]['tcp'][p]['state'] == 'open':
                    port_list.append(str(p))
            if port_list is not list():
                self.host_dict[hostname] = port_list
                np += len(port_list)
        log("Found {} open ports on {} host(s) with {}".format(np, len(self.host_dict), "masscan"), self.verbose)


def output(o, host, st):
    host = host.replace("/", "x")
    host = host.replace(" ", "")
    if not o:
        return ""
    if st == 1:
        return " -oG nmap_{}_S_{}.gnmap".format(host, date())
    if st == 2:
        return " -oG nmap_{}_SVCA_{}.gnmap".format(host, date())
    if st == 3:
        return " -oG nmap_{}_VCUA_{}.gnmap".format(host, date())
    if st == 4:
        return "nmap_{}_S_{}.xml".format(host, date())
    if st == 5:
        return "nmap_{}_SVCA_{}.xml".format(host, date())
    if st == 6:
        return "nmap_{}_VCUA_{}.xml".format(host, date())
    if st == 7:
        return " -oG masscan_{}_S_{}.gnmap -oX masscan_{}_S_{}.xml".format(host, date(), host, date())


def log(message, verbose):
    if verbose: print("{}: {}".format(date(True), message))


def date(long=False):
    if long:
        return datetime.now().strftime("%Y-%m-%d_%H%M%S")
    return datetime.now().strftime("%Y-%m-%d_%H%M")