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

"""
Script to format a grepable (-oG) report of a NMAP scan to different output format.

nmap -Pn --open -p- -oG report.txt [CDIR]

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

Dependencies:
    pip3 install tabulate 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]:
        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 format_entries_as_lines(hosts):
    entries = [["Host", "Ports"]]
    for host in hosts:
        ports = hosts[host]
        entries.append([host, "/".join(ports)])
    return entries


def format_to_table(entries):
    table_headers = entries[0]
    table_rows = entries[1:]
    print(tabulate(table_rows, headers=table_headers, tablefmt="outline"))


def format_to_csv(entries):
    filename = "nmap-output.csv"
    print(colored(f"[+] Generate the CSV to file '{filename}' ...", "yellow"))
    table_headers = entries[0]
    table_rows = entries[1:]
    rows = []
    for entry in table_rows:
        rows.append(",".join(entry))
    rows.sort()
    with open(filename, mode="w", encoding="utf-8") as f:
        f.write(",".join(table_headers))
        f.write("\n")
        f.write("\n".join(rows))
    print(colored(f"[V] Done ({len(table_rows)} entries).", "green"))


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Format a grepable (-oG) report of a NMAP scan to different output formats.")
    required_params = parser.add_argument_group("required named arguments")
    required_params.add_argument("-f", action="store", dest="nmap_file", help="Path to the grepable (-oG) report of a NMAP scan.", required=True)
    required_params.add_argument("-t", action="store", dest="output_type", choices=["table", "csv"], help="Type of the output.", required=True)
    args = parser.parse_args()
    hosts = load_hosts(args.nmap_file)
    entries = format_entries_as_lines(hosts)
    format_type = args.output_type
    if format_type == "table":
        format_to_table(entries)
    elif format_type == "csv":
        format_to_csv(entries)
