#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from core.constants import EXTERNAL, IG
from core.regressiontest import RegressionTest
from core.module import Module
import xmltodict
import json
import os


class NmapModule(Module):
    _type = [IG]
    _services = []
    _target_type = EXTERNAL
    _defaults = {}
    _extension = ".xml"

    def run(self):
        command = [
            "sudo",
            "nmap",
            "-sS",
            "-p-",
            "-oX",
            (self._tmpfile_loc()),
        ]

        self.run_cmd(" ".join(command + [self._target]))

    def parse(self):
        # Get current dir and check that path exists, else create path.
        target_loc = f"{os.getcwd()}/targets/{self._target}/results"

        # Open xml scan and convert to json
        if os.path.exists(f"{target_loc}/{self._res_name}.xml"):
            with open(f"{target_loc}/{self._res_name}.xml", "r") as f:
                xml = json.dumps(xmltodict.parse(f.read()), indent=4)

            # Overwrite scan results json
            with open(f"{target_loc}/{self._res_name}.json", "w+") as f:
                f.write(xml)

    def init(self):
        pass


"""
Define regression tests here
"""


class NmapPortRegression(RegressionTest):
    _extension = ".xml" # Specifies what extension the result file has
    _description = "Checks if a given set of Nmap ports are still open"
    _defaults = {
        
    }

    def run(self):
        command = [
            "sudo",
            "nmap",
            "-sS",
            "-p",
            ",".join(self._payload["ports"]),
            "-oX",
            (self._tmpfile_loc()),
        ]

        self.run_cmd(" ".join(command + [self._target]))

    def parse(self):
        # Get current dir and check that path exists, else create path.
        target_loc = f"{os.getcwd()}/targets/{self._target}/regtests"

        # Open xml scan and convert to json
        if os.path.exists(f"{target_loc}/{self._res_name}.xml"):
            data = {}
            with open(f"{target_loc}/{self._res_name}.xml", "r") as f:
                data = xmltodict.parse(f.read())
            
            result = {port: "closed" for port in self._payload["ports"]}

            # Make into list
            if type(data["nmaprun"]["host"]["ports"]["port"]) is not list:
                data["nmaprun"]["host"]["ports"]["port"] = [data["nmaprun"]["host"]["ports"]["port"]]

            for port in data["nmaprun"]["host"]["ports"]["port"]:
                if (
                    port["state"]["@state"] == "open"
                ):
                    result[port["@portid"]] = "open"
            
            # Overwrite scan results json
            with open(f"{target_loc}/{self._res_name}.json", "w+") as f:
                f.write(json.dumps(result, indent=4))

    def check(target):
        ports = [] # Specifies which ports that are of interest (Leave empty if all)
        open_ports = []

        # Check nmap file for open ports
        with open(
            f"{os.getcwd()}/targets/{target}/results/nmap_result.json",
            "r"
        ) as f:
            data = json.load(f)

        # Make into list
        if type(data["nmaprun"]["host"]["ports"]["port"]) is not list:
            data["nmaprun"]["host"]["ports"]["port"] = [data["nmaprun"]["host"]["ports"]["port"]]
        
        # Check all ports
        if not ports:
            for port in data["nmaprun"]["host"]["ports"]["port"]:
                if port["state"]["@state"] == "open":
                    open_ports.append(
                        port["@portid"]
                    )
        # add ports with specific services
        else:
            for port in data["nmaprun"]["host"]["ports"]["port"]:
                if (
                    port["@portid"] in ports
                    and port["state"]["@state"] == "open"
                ):
                    open_ports.append(port["@portid"])

        # Return None if empty
        if len(open_ports) == 0:
            return None

        return {"ports": open_ports}
