#!/usr/bin/env python
import argparse
from termcolor import colored

"""
Script to generate an valid input file consumed by NUCLEI
from a output file coming from NMAP in grepable format (-oG).

Sources:
    https://github.com/nmap/nmap
    https://nmap.org/book/man-output.html
    https://github.com/projectdiscovery/nuclei

Dependencies:
    pip install colorama termcolor
"""


def load_hosts(nmap_file):
    hosts = {}
    with open(nmap_file, mode="r", encoding="utf-8") as f:
        lines = f.read().splitlines()
    if " -oG " not in lines[0] and " -oA " not in lines[0]:
        raise Exception("Input file is not using the NMAP grepable format!")
    for line in lines:
        if "Ports:" in line and "/open/" in line:
            parts = line.split("(")
            # Extract the host IP part
            host = parts[0].split(":")[1].strip(" ()")
            if host not in hosts:
                hosts[host] = []
            # Extract the ports part
            ports = parts[1].split(":")[1].split(",")
            for port in ports:
                port_value = port.split("/")[0].strip(" ")
                if port_value not in hosts[host]:
                    hosts[host].append(port_value)
    return hosts


def generate_nuclei_input_entries(hosts):
    entries = []
    for host in hosts:
        for port in hosts[host]:
            port_as_int = int(port)
            if port_as_int in [80, 8080, 9080]:
                protocol = "http://"
            elif port_as_int in [443, 8443, 9043, 9443]:
                protocol = "https://"
            else:
                protocol = ""
            entries.append(f"{protocol}{host}:{port}")
    return entries


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Script to generate an valid input file, consumed by NUCLEI, from a output file coming from NMAP in grepable format.")
    parser.add_argument("-f", action="store", dest="nmap_file", help="Text file coming from the NMAP scan", required=True)
    parser.add_argument("-o", action="store", dest="nuclei_file", help="Text file consumed by NUCLEI", required=False, default="/tmp/nuclei-input.txt")
    args = parser.parse_args()
    hosts = load_hosts(args.nmap_file)
    print(colored(f"[+] Generate the nuclei file to '{args.nuclei_file}' ...", "yellow"))
    entries = generate_nuclei_input_entries(hosts)
    with open(args.nuclei_file, mode="w", encoding="utf-8") as f:
        f.write("\n".join(entries))
    print(colored(f"[V] Done ({len(entries)} entries).", "green"))
