from utils import log, Col


class AnmapResult(dict):
    def add_host(self, host):
        try:
            if host.services_tcp:
                self[host.hostname].append(host.services_tcp)
            if host.services_udp:
                self[host.hostname].append(host.services_udp, True)
            if host.osmatch:
                self[host.hostname].osmatch.extend(host.osmatch)
        except KeyError:
            self[host.hostname] = host
        return host.open_ports

    def print(self):
        for host in self.values():
            host.print()


class AnmapHost:
    def __init__(self, host, verbose=False):
        self.ip = host["addresses"]["ipv4"]
        self.hostname = host.hostname()
        self.osmatch = list()
        try:
            self.osmatch.extend(host["osmatch"])
        except KeyError:
            self.information = ""
        self.state = host.state()
        self.services_tcp = dict()
        self.services_udp = dict()
        self.open_ports = 0
        log(self.get_info(), verbose)
        for t in host.all_tcp():
            if "open" not in host['tcp'][t]['state']:
                continue
            self.open_ports += 1
            log("Port {}/tcp: {}".format(t, host['tcp'][t]), verbose)
            self.services_tcp[t] = host["tcp"][t]
        for u in host.all_udp():
            if "open" not in host['udp'][u]['state']:
                continue
            self.open_ports += 1
            log("Port {}/udp: {}".format(u, host['udp'][u]), verbose)
            self.services_udp[u] = host["udp"][u]

    def print(self):
        print(self.get_info())
        for guess in self.osmatch:
            print("\tOS-Guess: {} ({}%)".format(guess["name"], guess["accuracy"]))
            for osclass in guess["osclass"]:
                for key, value in osclass.items():
                    print("\t\t{}: {}".format(key, value))
        for port, info in self.services_tcp.items():
            print("\tPort {}/tcp:".format(port))
            for key, value in info.items():
                if key == "script":
                    print("\t\tScript:")
                    for k, v in value.items():
                        print("\t\t\t{}: {}".format(k, v))
                    continue
                print("\t\t{}: {}".format(key, value))
        for port, info in self.services_udp.items():
            print("\tPort {}/udp:".format(port))
            for key, value in info.items():
                if key == "script":
                    print("\t\tScript:")
                    for k, v in value.items():
                        print("\t\t\t{}: {}".format(k, v.replace("\n", "\n\t\t\t\t")))
                    continue
                print("\t\t{}: {}".format(key, value))

    def get_info(self):
        return "{}/{} is {} with {} open ports".format(self.hostname, self.ip, self.state, self.open_ports)

    def append(self, services, udp=False):
        for port, service in services.items():
            if udp:
                if port in self.services_udp:
                    for key, value in service.items():
                        self.services_udp[port][key] = value
                    continue
                self.services_udp[port] = service
                self.open_ports += 1
            else:
                if port in self.services_tcp:
                    for key, value in service.items():
                        self.services_tcp[port][key] = value
                    continue
                self.services_tcp[port] = service
                self.open_ports += 1