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, out):
        Thread.__init__(self)
        self.host = hostname
        self.ports = ports
        self.nm = PortScanner()
        self.logger = logger
        self.daemon = True
        self.out = out


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{}'.format(output_argument(self.out, self.host, 2)))
        host = self.nm[self.host]
        for p in host.all_tcp():
            if p == 1:
                continue
            print("Port {}/tcp: {}".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, output_argument(self.out, self.host, 3)))
        host = self.nm[self.host]
        for p in host.all_udp():
            print("Port {}/udp: {}".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 output_argument(o, host, st):
    host = host.replace("/", "x")
    host = host.replace(" ", "")
    if not o:
        return ""
    if st == 1:
        return " -oG nmap_{}_S_{}".format(host, date())
    if st == 2:
        return " -oG nmap_{}_SVCA_{}".format(host, date())
    if st == 3:
        return " -oG nmap_{}_VCUA_{}".format(host, date())


def run(args):
    if args.d:
        args.v = True
        args.u = 100
    l = Logger(args.v)

    # Scanning all tcp ports
    nm = PortScanner()
    l.log("Starting quick scan")
    if args.d:
        nm.scan(args.HOST, arguments='-sS -Pn -p1-1000{}'.format(output_argument(args.o, args.HOST, 1)))
    else:
        nm.scan(args.HOST, arguments='-sS -Pn -p-{}'.format(output_argument(args.o, args.HOST, 1)))
    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, args.o)
        t1.start()
        thread_list.append(t1)
        t2 = UDPAnmapThread(host, args.u, l, args.o)
        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("-d", action="store_true",
                        help="Sets flags -v and -u 100 and scans only the first 1000 tcp ports")
    parser.add_argument("-o", action="store_true", help="Enables saving of output files")
    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)