NetQuestOMX

import tempfile
from collections.abc import Callable
from enum import Enum
from pathlib import Path

import demistomock as demisto  # noqa: F401
import urllib3
from CommonServerPython import *  # noqa: F401

from CommonServerUserPython import *  # noqa

# Disable insecure warnings
TOKEN_TTL_S = 120 * 60  # tokens are valid for 120 minutes
urllib3.disable_warnings()

""" CONSTANTS """

DATE_FORMAT = "%Y-%m-%dT%H:%M:%SZ"  # ISO8601 format with UTC, default in XSOAR
DATE_FORMAT_FOR_TOKEN = "%m/%d/%Y, %H:%M:%S"
VENDOR = "NetQuest"
PRODUCT = "OMX"

""" CLIENT CLASS """


class Client(BaseClient):
    def __init__(self, base_url: str, credentials: dict, verify: bool, proxy: bool):
        self._headers = {"Accept": "application/json", "X-Auth-Token": ""}
        super().__init__(base_url=base_url, verify=verify, proxy=proxy, headers=self._headers)
        self.credentials = {"UserName": credentials["identifier"], "Password": credentials["password"]}
        self.login()

    def login(self):
        """
        If there's an existing token in context, its validity is checked.
        Otherwise, logs in and stores token in the `Client` instance
        Refreshes the token as needed.
        """
        now = datetime.utcnow()

        if (cache := get_integration_context()) and cache.get("Token"):
            expiration_time = datetime.strptime(cache["expiration_time"], DATE_FORMAT_FOR_TOKEN)

            # if the token is still valid, continue using it. otherwise, generate a new one.
            if (seconds_left := (expiration_time - now).total_seconds()) > 60:  # decreasing 60s from token expiry for safety
                demisto.debug(f"No need to regenerate the token, it is still valid for {seconds_left} more seconds")
                self._headers["X-Auth-Token"] = cache["Token"]  # in this case, the Token key exists
                return

        demisto.debug("IntegrationContext token cache is empty or token has expired, regenerating a new token")

        self._refresh_access_token(now)

    def _refresh_access_token(self, now: datetime):
        """
        Since the validity of the Access Token is 120 minutes, this method refreshes the token
        and keep the token and the expiration time in the integration context.
        """

        try:
            response = self._http_request(
                method="POST", url_suffix="/api/SessionService/Sessions", json_data=self.credentials, resp_type="response"
            )
        except Exception as e:
            raise DemistoException("An error was occurred when creating a new token. Please verify your credentials.") from e

        new_token = response.headers["X-Auth-Token"]
        self._headers["X-Auth-Token"] = new_token
        set_integration_context(
            {
                "Token": new_token,
                "expiration_time": (now + timedelta(seconds=TOKEN_TTL_S)).strftime(DATE_FORMAT_FOR_TOKEN),
            }
        )

    def address_list_upload_request(self, file_name: str, file_path: str):
        """
        Sends a POST request to upload an address list file to the appliance.

        The file is managed as a temporary file to ensure proper cleanup.

        Args:
            file_name (str): The name of the original file to be uploaded.
            file_path (str): The path to the original file to be uploaded.

        Raises:
            DemistoException: If the HTTP request fails or other errors occur.
        """
        try:
            # Create a temporary file
            with tempfile.NamedTemporaryFile() as temp_file:
                temp_file_name = temp_file.name

                # Copy the contents of the original file to the temporary file
                Path(temp_file_name).write_bytes(Path(file_path).read_bytes())

                # Use the temporary file for the upload request
                self._http_request(
                    method="POST",
                    url_suffix="/api/v1/UpdateService/ImportList/Config",
                    files={"UpdateFile": open(temp_file_name, "rb")},
                    ok_codes=(200,),
                )
        except Exception as e:
            raise DemistoException(f"An error occurred while uploading the file {file_name}.") from e
        finally:
            # Ensure the temporary file is deleted
            Path(temp_file_name).unlink(missing_ok=True)

    def address_list_optimize_request(self) -> dict:
        try:
            response_json = self._http_request(method="GET", url_suffix="/api/Systems/Filters/Address/Status/Optimization")
        except Exception as e:
            raise DemistoException("An error was occurred when optimizing the list of IPs.") from e

        return response_json

    def address_list_create_request(self, name: str):
        try:
            self._http_request(
                method="POST",
                url_suffix="/api/Systems/Filters/ListImport/Config/Install",
                json_data={"Name": name},
                ok_codes=(200,),
            )
        except Exception as e:
            raise DemistoException("An error was occurred when creating the list of IPs.") from e

    def address_list_rename_request(self, new_name: str, existing_name: str):
        try:
            self._http_request(
                method="PUT",
                url_suffix=f"/api/Systems/Filters/ListImport/ListName/{existing_name}/Config/Install",
                json_data={"Name": new_name},
                ok_codes=(200,),
            )
        except Exception as e:
            raise DemistoException(f"An error occurred when renaming the {existing_name} IP list to {new_name}.") from e

    def address_list_delete_request(self, list_name_to_delete: str):
        try:
            self._http_request(
                method="DELETE",
                url_suffix=f"/api/Systems/Filters/Address/ListName/{list_name_to_delete}/Config/List",
                ok_codes=(200,),
            )
        except Exception as e:
            raise DemistoException(f"An error was occurred when deleting the {list_name_to_delete} IP list.") from e

    def metering_stats_request(self, slot_number: str, port_number: str) -> dict[str, Any]:
        try:
            metering_stats_event = self._http_request(
                method="GET", url_suffix=f"/api/Systems/Slot/{slot_number}/Ipfix/Status/Metering", ok_codes=(200,)
            )
        except Exception as e:
            raise DemistoException("An error was occurred when requesting for an event of Metering Stats type.") from e

        metering_stats_event["STAT_TYPE"] = "MeteringStats"

        return metering_stats_event

    def export_stats_request(self, slot_number: str, port_number: str) -> dict[str, Any]:
        try:
            export_stats_event = self._http_request(
                method="GET", url_suffix=f"/api/Systems/Slot/{slot_number}/Ipfix/Status/Export", ok_codes=(200,)
            )
        except Exception as e:
            raise DemistoException("An error was occurred when requesting for an event of Export Stats type.") from e

        export_stats_event["STAT_TYPE"] = "ExportStats"

        return export_stats_event

    def export_peaks_FPS_request(self, slot_number: str, port_number: str) -> dict[str, Any]:
        try:
            export_peaks_FPS_event = self._http_request(
                method="GET", url_suffix=f"/api/Systems/Slot/{slot_number}/Ipfix/Status/ExportHwm", ok_codes=(200,)
            )
        except Exception as e:
            raise DemistoException("An error was occurred when requesting for an event of Export Peaks FPS type.") from e

        export_peaks_FPS_event["STAT_TYPE"] = "ExportPeaksFPS"

        return export_peaks_FPS_event

    def optimization_stats_request(self, slot_number: str, port_number: str) -> dict[str, Any]:
        try:
            optimization_stats_event = self._http_request(
                method="GET",
                url_suffix=f"/api/Systems/Slot/{slot_number}/Port/{port_number}/EthernetInterfaces/Status/EthRxTx",
                ok_codes=(200,),
            )
        except Exception as e:
            raise DemistoException("An error was occurred when requesting for an event of Optimization Stats type.") from e

        optimization_stats_event["STAT_TYPE"] = "OptimizationStats"

        return optimization_stats_event


""" COMMAND FUNCTIONS """


class StatType(Enum):
    METERING_STATS = "Metering Stats"
    EXPORT_STATS = "Export Stats"
    EXPORT_PEAKS_FPS = "Export Peaks FPS"
    OPTIMIZATION_STATS = "Optimization Stats"


def address_list_upload_command(client: Client, args: dict) -> CommandResults:
    """
    This function uploads a .txt file with address list to the appliance.
    The appliance temporarily stores the file until it is saved to the Library and replaces any previously loaded list file.

    Returns:
        A CommandResults containing a success indication or a DemistoException.
    """
    entry_id = args["entry_id"]  # a required argument
    file_info = demisto.getFilePath(entry_id)
    file_path = file_info["path"]
    file_name = file_info["name"]

    try:
        client.address_list_upload_request(file_name, file_path)
    except Exception as e:
        raise DemistoException(f"An error occurred when uploading {file_name}.") from e

    return CommandResults(readable_output="Address list was successfully uploaded")


def address_list_optimize_command(client: Client, args: dict = None) -> CommandResults:
    """
    If the traffic elements are IP addresses,
    the integration should optimize the list by compressing IP addresses into CIDR groups.

    Returns:
        A CommandResults containing full API response.
    """

    response_json = client.address_list_optimize_request()

    return CommandResults(outputs_prefix="NetQuest.AddressList", outputs=response_json)


def address_list_create_command(client: Client, args: dict) -> CommandResults:
    """
    This function replaces the old list entity and overrides it.

    Returns:
        A CommandResults containing a success indication or a DemistoException.
    """
    name = args["name"]  # a required argument - The name of the address list to create
    try:
        client.address_list_create_request(name)
    except Exception as e:
        raise DemistoException(
            f"An error was occurred when creating an IP list with {name=}."
            f" This may indicate a list with this name already exists."
        ) from e

    return CommandResults(readable_output=f"Successfully created a new instance of {name}")


def address_list_rename_command(client: Client, args: dict) -> CommandResults:
    """
    This function only meant to change the name of the list.
    Nothing else. If we try to give as a new_name, an existing list name, it will fail, and we'll get an error.

    Returns:
        A CommandResults containing a success indication or a DemistoException.
    """
    new_name = args["new_name"]  # a required argument - The new name for an existing address list
    existing_name = args["existing_name"]  # a required argument - Name of the existing address list that we want to modify

    try:
        client.address_list_rename_request(new_name, existing_name)
    except Exception as e:
        raise DemistoException(
            f"An error was occurred when renaming {existing_name} IPs list to {new_name}."
            f" Make sure {existing_name} exists and that {new_name} doesn't."
        ) from e

    return CommandResults(readable_output=f"Successfully renamed {existing_name} to {new_name}")


def address_list_delete_command(client: Client, args: dict) -> CommandResults:
    """
    This function deletes the list by the given list's name.

    Returns:
        A CommandResults containing a success indication or a DemistoException.
    """
    list_name_to_delete = args["name"]  # a required argument - The name of the address list to delete

    try:
        client.address_list_delete_request(list_name_to_delete)
    except Exception as e:
        raise DemistoException(
            f"An error was occurred when deleting the {list_name_to_delete} IPs list. Make sure {list_name_to_delete} exists."
        ) from e

    return CommandResults(readable_output=f"Successfully deleted {list_name_to_delete} list")


def add_time_to_events(events: list[dict]) -> None:
    """
    Adds the _time key to the events.
    Args:
        events: list[dict] - list of events to add the _time key to.
    Returns:
        list: The events with the _time key.
    """
    create_time = datetime.utcnow().strftime(DATE_FORMAT)
    for event in events:
        event["_time"] = create_time  # type: ignore[union-attr]


def fetch_events(client: Client, slot_number: str, port_number: str, statistic_types_to_fetch: list[str]) -> list[dict]:
    """
    Args:
        client (Client): NetQuest client to use.
        slot_number (str): A Target Netquest device slot number.
        port_number (str): A port number to use.
        statistic_types_to_fetch (list): List of event types to fetch.
    Returns:
        events (list[dict]): A list of events (number of events equal to the number of statistic types given)
        that will be created in XSIAM.
    """
    demisto.debug(f"Starting Fetch: {slot_number=}, {port_number=}")

    events: list[dict] = []
    statistic_types_mapping: Dict[str, Callable] = {
        "Metering Stats": client.metering_stats_request,
        "Export Stats": client.export_stats_request,
        "Export Peaks FPS": client.export_peaks_FPS_request,
        "Optimization Stats": client.optimization_stats_request,
    }

    events.extend(
        statistic_types_mapping[statistic_type](slot_number, port_number) for statistic_type in statistic_types_to_fetch
    )

    return events


def get_events(client: Client, params: dict, args: dict) -> list[dict]:
    """
    Args:
        client (Client): NetQuest client to use.
        params (dict)
        args (dict)
    Returns:
        events (list[dict]): A list of events (number of events equal to the number of statistic types given)
        that will be created in XSIAM.
    """

    # validate the input
    statistic_types_to_fetch = argToList(
        args.get("statistic_types_to_fetch", []) or params.get("statistic_types_to_fetch", [])
    )  # arg overrides the param

    for statistic_type in statistic_types_to_fetch:
        if statistic_type not in StatType._value2member_map_:
            raise DemistoException(
                f"{statistic_types_to_fetch} is not a valid type."
                f" Valid types are {list(StatType._value2member_map_.keys())}."
                f" Execute the command get-events again with valid input."
            )

    # execute the command
    events = fetch_events(
        client=client,
        slot_number=params["slot"],  # a required param
        port_number=params["port"],  # a required param
        statistic_types_to_fetch=statistic_types_to_fetch,
    )

    return events


""" MAIN FUNCTION """


def main() -> None:  # pragma: no cover
    """
    main function, parses params and runs command functions
    """

    params = demisto.params()
    args = demisto.args()
    command = demisto.command()
    proxy = params.get("proxy", False)
    verify_certificate = not params.get("insecure", False)

    demisto.debug(f"Command being called is {command}")

    try:
        client = Client(
            base_url=params["url"],
            credentials=params["credentials"],
            verify=verify_certificate,
            proxy=proxy,
        )

        commands: Dict[str, Callable] = {
            "netquest-address-list-upload": address_list_upload_command,
            "netquest-address-list-optimize": address_list_optimize_command,
            "netquest-address-list-create": address_list_create_command,
            "netquest-address-list-rename": address_list_rename_command,
            "netquest-address-list-delete": address_list_delete_command,
        }

        if demisto.command() == "test-module":
            # This is the call made when pressing the integration Test button - it returns 'ok' (success) if the client was built.
            return_results("ok")

        elif command in commands:
            return_results(commands[command](client, args))

        elif command == "fetch-events":
            events = fetch_events(
                client=client,
                slot_number=params["slot"],  # a required param (used only for urls, so it will remain a string)
                port_number=params["port"],  # a required param (used only for urls, so it will remain a string)
                statistic_types_to_fetch=argToList(params.get("statistic_types_to_fetch", [])),
            )

            add_time_to_events(events)
            send_events_to_xsiam(events, vendor=VENDOR, product=PRODUCT)

            demisto.debug(f"fetched and sent {len(events)} events to xsiam.")

        elif command == "get-events":
            events = get_events(client, params, args)

            return_results(CommandResults(readable_output=tableToMarkdown(f"{VENDOR} Events:", events), outputs=events))

            if argToBoolean(args["should_push_events"]):
                add_time_to_events(events)
                send_events_to_xsiam(events, vendor=VENDOR, product=PRODUCT)

        else:
            raise NotImplementedError(f"{command} command is not implemented.")

    # Log exceptions and return errors
    except Exception as e:
        demisto.error(traceback.format_exc())  # prints the traceback
        return_error(f"Failed to execute {command} command.\nError:\n{e!s}")


""" ENTRY POINT """

if __name__ in ("__main__", "__builtin__", "builtins"):
    main()