Cyber Triage

from typing import Any

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

# Disable warning for insecure requests when cert validation is disabled
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)


def IS_2XX(x: int) -> bool:
    return int(x / 100) == 2  # Returns true if status code (int) is 2xx


class CyberTriageClient(BaseClient):
    SCAN_OPTIONS = ["pr", "nw", "nc", "st", "sc", "ru", "co", "lo", "ns", "wb", "fs"]

    def __init__(
        self,
        server: str,
        rest_port: str,
        api_key: str,
        user: str,
        password: str,
        verify_server_cert: bool,
        ok_codes: tuple[int, ...],
    ):
        base_url = (
            f"https://{server}:{rest_port}/api/"
            if not (server.startswith(("https://", "http://")))
            else f"{server}:{rest_port}/api/"
        )
        req_headers = {"restApiKey": api_key}
        self._user = user
        self._password = password
        super().__init__(base_url=base_url, verify=verify_server_cert, headers=req_headers, ok_codes=ok_codes)

    def test_connection(self):
        response = self._http_request("GET", url_suffix="correlation/checkcredentials", resp_type="response")
        return response

    def triage_endpoint(
        self, is_hash_upload_on: bool, is_file_upload_on: bool, endpoint: str, scan_options: str, incident_name: str
    ):
        # Validate scan options
        invalid_options = []
        if scan_options:
            invalid_options = [opt for opt in scan_options.split(",") if opt not in self.SCAN_OPTIONS]
        if invalid_options:
            raise DemistoException("The following are not valid scan options: {}".format(",".join(invalid_options)))

        # Make data dict for rest call
        api_data = {
            "incidentName": incident_name,
            "hostName": endpoint,
            "userId": self._user,
            "password": self._password,
            "scanOptions": scan_options,
            "malwareScanRequested": is_hash_upload_on,
            "sendContent": is_file_upload_on,
            "sendIpAddress": False,
        }
        response = self._http_request("POST", url_suffix="livesessions", data=api_data, resp_type="response")
        return response


def test_connection_command(client: CyberTriageClient) -> str:
    response = client.test_connection()
    response.raise_for_status()
    return "ok"


def triage_endpoint_command(client: CyberTriageClient, args: dict[str, Any]) -> CommandResults:
    def is_true(x: str) -> bool:
        return x == "yes"

    is_hash_upload_on = is_true(args.get("malware_hash_upload", ""))  # arg value = 'yes' or 'no'
    is_file_upload_on = is_true(args.get("malware_file_upload", ""))  # arg value = 'yes' or 'no'
    endpoint = args.get("endpoint", "")
    scan_options = args.get("scan_options", "")
    incident_name = args.get("incident_name", "")
    response = client.triage_endpoint(is_hash_upload_on, is_file_upload_on, endpoint, scan_options, incident_name)

    response.raise_for_status()

    if is_ip_valid(endpoint):
        endpoint_context = {"IPAddress": endpoint}
    else:
        endpoint_context = {"Hostname": endpoint}

    data = response.json()

    ec = {"CyberTriage": data, "Endpoint": endpoint_context}

    return CommandResults(readable_output=f"A collection has been scheduled for {endpoint}", outputs=ec, raw_response=data)


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

    server = params.get("server", "")
    rest_port = params.get("rest_port", "")
    api_key = params.get("api_key", {}).get("password", "")
    user = params.get("credentials", {}).get("identifier", "")
    password = params.get("credentials", {}).get("password", "")
    verify_server_cert = False
    handle_proxy(proxy_param_name="use_proxy")

    demisto.debug(f"Command being called is {command}")
    try:
        acceptable_status_codes: tuple[int, ...] = tuple(int(code) for code in requests.status_codes.codes if IS_2XX(code))
        client = CyberTriageClient(server, rest_port, api_key, user, password, verify_server_cert, acceptable_status_codes)
        # This is the call made when running the ct-triage-endpoint command.
        if command == "ct-triage-endpoint":
            return_results(triage_endpoint_command(client, args))

        # This is the call made when pressing the integration test button.
        elif command == "test-module":
            return_results(test_connection_command(client))
        else:
            raise NotImplementedError(f"command={command} not implemented in this integration")
    except Exception as e:
        return_error(f"Failed to execute {command} command.\nError: {e!s}")


if __name__ in ("__main__", "__builtin__", "builtins"):  # pragma: no cover
    main()