import nmap
import masscan
from threading import Thread
from AnmapResult import AnmapHost
from utils import date, log


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

    def run(self):
        try:
            log("Starting {} scan on {}".format(self.proto, self.host), self.verbose)
            if self.proto == "tcp":
                xml = 5
                self.scanner.scan(self.host, arguments='-p 1,{} -sSVC -A -Pn{}'.format(",".join(map(str, self.ports)), output(self.out, self.host, 2)))
            elif self.proto == "udp":
                xml = 6
                self.scanner.scan(self.host, arguments='-sVCU -A -Pn --top-ports {}{}'.format(self.ports, output(self.out, self.host, 3)))
            elif self.proto == "quick":
                xml = 4
                self.scanner.scan(self.host, arguments='-sS -Pn -p{}{}'.format(self.ports, output(self.out, self.host, 1)))
            elif self.proto == "masscan":
                xml = False
                self.scanner.scan(self.host, ports=self.ports, arguments=output(self.out, self.host, 7), sudo=True)
            else:
                return
            log(self.scanner.command_line(), self.verbose)
            if self.out and xml:
                with open(output(True, self.host, xml), "w") as outfile:
                    outfile.write(self.scanner.get_nmap_last_output())
            np = nh = 0
            for hn in self.scanner.all_hosts():
                np += self.result.add_host(AnmapHost(self.scanner[hn], self.verbose))
                nh += 1
            log("Found {} open ports on {} host(s)".format(np, nh), self.verbose)
            log("Finished {} scan on {}".format(self.proto, self.host), self.verbose)
        except KeyboardInterrupt:
            return


class MasscanThread(AnmapThread):
    def __init__(self, hostname, ports, verbose, out):
        AnmapThread.__init__(self, hostname, ports, verbose, out)
        self.scanner = masscan.PortScanner()


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())