#!/usr/bin/python3
from argparse import ArgumentParser
from AnmapThread import AnmapThread, MasscanThread
from AnmapResult import AnmapResult


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 10 and scans only the first 100 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 = 10
    result = AnmapResult()
    try:
        c = ""
        try:
            ms = args.masscan
        except AttributeError:
            ms = False
        ports = "1-100" if args.debug else "-"
        if ms:
            tm = MasscanThread(args.HOST, ports, args.verbose, args.output)
            tm.start()
            tm.join()
            c = input("Do you want to continue without a full nmap scan? (y/N)")
            if c != "y":
                t0 = AnmapThread(args.HOST, ports, args.verbose, args.output, result, "quick")
                t0.start()
                t0.join()
        else:
            t0 = AnmapThread(args.HOST, ports, args.verbose, args.output, result, "quick")
            t0.start()
            t0.join()
        # Starting thorough and udp scan for each host in separate threads
        thread_list = list()
        for name, host in result.items():
            thread_list.append(AnmapThread(name, host.services_tcp.keys(), args.verbose, args.output, result, "tcp"))
            thread_list.append(AnmapThread(name, args.udp, args.verbose, args.output, result, "udp"))
        for t in thread_list:
            t.start()
        # Waiting for the threads to finish
        for t in thread_list:
            t.join()
        result.print()
    except KeyboardInterrupt:
        print("User Interrupt")