from nmap import PortScanner
from threading import Thread
from argparse import ArgumentParser
from datetime import datetime


class AnmapThread(Thread):
    def __init__(self, hostname, ports, logger):
        Thread.__init__(self)
        self.host = hostname
        self.ports = ports
        self.nm = PortScanner()
        self.logger = logger
        self.daemon = True


class ThoroughAnmapThread(AnmapThread):
    def run(self):
        self.logger.log("Starting thorough scan on " + self.host)
        self.nm.scan(self.host, "1," + ",".join(self.ports), arguments="-sSVC -A -Pn")
        host = self.nm[self.host]
        for p in host.all_tcp():
            print("Port {}: {}".format(p, host['tcp'][p]))
        self.logger.log("Finished thorough scan on " + self.host)


class UDPAnmapThread(AnmapThread):
    def run(self):
        self.logger.log("Starting UDP scan on " + self.host)
        self.nm.scan(self.host, arguments="-sVCU -A -Pn --top-ports {}".format(self.ports))
        host = self.nm[self.host]
        for p in host.all_udp():
            print("Port {}: {}".format(p, host['udp'][p]))
        self.logger.log("Finished UDP scan on " + self.host)


class Logger:
    def __init__(self, verbose):
        self.verbose = verbose

    def log(self, message):
        if self.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")


def run(args):
    l = Logger(args.v)

    # Scanning all tcp ports
    nm = PortScanner()
    l.log("Starting quick scan")
    nm.scan(args.HOST, arguments='-sS -Pn -p1-1000')
    l.log("Finished quick scan")
    host_list = dict()
    for hostname in nm.all_hosts():
        host = nm[hostname]
        port_list = list()
        for p in host.all_tcp():
            if nm[hostname]['tcp'][p]['state'] == 'open':
                port_list.append(str(p))
        if port_list is not list():
            host_list[hostname] = port_list

    # Starting thorough and udp scan in separate threads
    thread_list = []
    for host, open_port_list in host_list.items():
        t1 = ThoroughAnmapThread(host, open_port_list, l)
        t1.start()
        thread_list.append(t1)
        t2 = UDPAnmapThread(host, args.u, l)
        t2.start()
        thread_list.append(t2)

    # Waiting for the threads to finish
    for t in thread_list:
        t.join()


if __name__ == "__main__":
    # Argument parsing
    parser = ArgumentParser(description="This script automates nmap scans by quickly scanning all TCP ports first and "
                                        "executing a thorough scan on all ports found open afterwards. "
                                        "Additionally it scans a given number of most used UDP ports.",
                            prog="anmap.py")
    parser.add_argument("-u", default=1000, type=int, help="The number of UDP ports to scan (Default 1000)")
    parser.add_argument("-v", action="store_true", help="This enables verbose output")
    parser.add_argument("HOST", type=str, help="The hosts to scan (Same notations as in nmap possible)")

    try:
        run(parser.parse_args())
    except KeyboardInterrupt:
        print("User Interrupt")
        exit(0)