#!/usr/bin/python3
from argparse import ArgumentParser
from AnmapThread import UDPAnmapThread, ThoroughAnmapThread, BaseAnmapThread, MasscanAnmapThread


if __name__ == "__main__":
    # Argument parsing
    ap = 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")
    ap.add_argument("-u", "--udp", default=1000, type=int, help="The number of UDP ports to scan (Default 1000)")
    ap.add_argument("-v", "--verbose", action="store_true", help="This enables verbose output")
    # Not functional yet
    # ap.add_argument("-m", "--masscan", action="store_true", help="This enables masscan for first scan")
    ap.add_argument("-d", "--debug", action="store_true",
                    help="Sets flags -v and -u 100 and scans only the first 1000 tcp ports")
    ap.add_argument("-o", "--output", action="store_true", help="Enables saving of output files")
    ap.add_argument("HOST", type=str, help="The hosts to scan (Same notations as in nmap possible)")
    args = ap.parse_args()
    if args.debug:
        args.verbose = True
        args.udp = 100
    try:
        c = host_dict = ""
        try:
            ms = args.masscan
        except AttributeError:
            ms = False
        if ms:
            tm = MasscanAnmapThread(args.HOST, "1-1000" if args.debug else "-", args.verbose, args.output)
            tm.start()
            host_dict = tm.rjoin()
            c = input("Do you want to continue without a full nmap scan? (y/N)")
            if c != "y":
                t0 = BaseAnmapThread(args.HOST, "1-1000" if args.debug else "-", args.verbose, args.output)
                t0.start()
                host_dict = t0.rjoin()
        else:
            t0 = BaseAnmapThread(args.HOST, "1-1000" if args.debug else "-", args.verbose, args.output)
            t0.start()
            host_dict = t0.rjoin()
        # Starting thorough and udp scan for each host in separate threads
        thread_list = list()
        for host, open_port_list in host_dict.items():
            thread_list.append(ThoroughAnmapThread(host, open_port_list, args.verbose, args.output))
            thread_list.append(UDPAnmapThread(host, args.udp, args.verbose, args.output))
        for t in thread_list:
            t.start()
        # Waiting for the threads to finish
        for t in thread_list:
            t.join()
    except KeyboardInterrupt:
        print("User Interrupt")