#!/usr/bin/python
import sys
import json
from rich.console import Console
from rich.table import Table
import argparse

# Setup Table
console = Console()

table = Table(show_header=True, header_style="bold blue")
table.add_column("Host", style="dim")
table.add_column("IP", style="dim")
table.add_column("Port")
table.add_column("Service")
table.add_column("Version", justify="left")

def pprint_table(hostlist, iponly):

    global table
    for host in hostlist:
        if iponly:
            print(host["ip"])
        else:
            for port in host["ports"]:
                table.add_row(
                    host["hostname"],
                    host["ip"],
                    str(port["port_number"]),
                    port["service"],
                    port["version"]
                )

    if not iponly:
        console.print(table)

def filter_by_port(port):
    global hosts
    out = []
    ports = []
    for host in hosts:
        for p in host["ports"]:
            if port == str(p["port_number"]) and p["state"] == "open":
                ports.append(p)
        host["ports"] = ports
        ports = []
        out.append(host)
    return out


def filter_by_version(version):
    global hosts
    out = []
    ports = []
    for host in hosts:
        for p in host["ports"]:
            if version.lower() in p["version"].lower():
                ports.append(p)
        host["ports"] = ports
        ports = []
        out.append(host)
    return out

# Setup Argument Parser
parser = argparse.ArgumentParser(description='Filtering nmap')
parser.add_argument('file', action='store', nargs='?',
                    help='Input File')
parser.add_argument('--port', dest='port', action='store',
                    help='Filter by port number')
parser.add_argument('--version', dest='version', action='store',
                    help='Filter by version string')
parser.add_argument('--ip', dest='ip', action='store_true',
                    help='Only print the ips')


args = parser.parse_args()

if args.file:
    with open(args.file, "r") as inp_file:
        hosts = json.loads(inp_file.read())
else:
    hosts = json.loads(sys.stdin.read())


if args.port:
    pprint_table(filter_by_port(args.port), args.ip)
elif args.version:
    pprint_table(filter_by_version(args.version), args.ip)
else:
    pprint_table(hosts, args.ip)