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

""" Cado Response API Integration for the Cortex XSOAR Platform """

import time
import traceback
from typing import Any

import urllib3

from CommonServerUserPython import *

""" Module Level Declarations """

urllib3.disable_warnings()
CadoResponseCombinedOutput = Union[dict[str, Any], List[dict[str, Any]]]

DATE_FORMAT: str = "%Y-%m-%dT%H:%M:%SZ"

""" Cado Response API Client Code """


class Client(BaseClient):
    """Client that makes HTTP requests to the Cado Response API"""

    def heartbeat(self) -> dict[str, Any]:
        """Calls the GET /api/v2/system/status endpoint to verify
        everything is working

        :return JSON response from /system/status endpoint
        :rtype Dict[str, Any]
        """

        return self._http_request(method="GET", url_suffix="/system/status")

    def create_project(self, project_name: str, project_description: str | None) -> dict[str, Any]:
        """Calls the POST /api/v2/projects endpoint to create a new
        project with given parameters

        :param str project_name: Name of the project
        :param Optional[str] project_description: Description for the project

        :return JSON response from /projects endpoint
        :rtype Dict[str, Any]
        """

        if not project_name.endswith("_XSOAR"):
            project_name += "_XSOAR"

        if not project_description:
            project_description = "This is a project in Cado Response created through Cortex XSOAR!"

        payload: dict[str, Any] = {"caseName": project_name, "description": project_description}

        return self._http_request(method="POST", url_suffix="/projects", json_data=payload)

    def get_project(self, project_id: int | None) -> dict[str, Any]:
        """Calls the GET /api/v2/projects endpoint to retrieve a
        project with given parameters

        :param Optional[int] project_id: ID of the project to retrieve

        :return JSON response from /projects endpoint
        :rtype Dict[str, Any]
        """

        if not project_id:
            project_id = demisto.params().get("CadoResponse_DefaultProject", 1)

        return self._http_request(method="GET", url_suffix=f"/projects/{project_id}")

    def list_projects(self, limit: int) -> List[dict[str, Any]]:
        """Calls the GET /api/v2/projects endpoint to retrieve a list
        of created projects

        :return JSON response from /projects endpoint
        :rtype Dict[str, Any]
        """

        data: List[dict[str, Any]] = self._http_request(method="GET", url_suffix="/projects")

        return data[:limit]

    def get_pipeline(self, pipeline_id: int | None, project_id: int | None) -> dict[str, Any]:
        """Calls the GET /api/v2/tasks/pipelines endpoint to
        retrieve details about a given pipeline

        :param Optional[int] pipeline_id: The id of the pipeline to retrieve
        :param Optional[int] project_id: The id of the project the pipeline belongs to

        :return JSON response from /tasks/pipelines endpoint
        :rtype Dict[str, Any]
        """

        if not pipeline_id:
            return {}

        if not project_id:
            project_id = demisto.params().get("CadoResponse_DefaultProject", 1)

        return self._http_request(
            method="GET", url_suffix="/tasks/pipelines", params={"project_id": project_id, "pipeline_id": pipeline_id}
        )

    def list_pipelines(self, project_id: int | None, limit: int) -> List[dict[str, Any]]:
        """Calls the GET /api/v2/tasks/pipelines endpoint to
        retrieve details about all of a projects pipelines

        :param Optional[int] project_id: The id of the project the pipeline belongs to

        :return JSON response from /tasks/pipelines endpoint
        :rtype Dict[str, Any]
        """

        if not project_id:
            project_id = demisto.params().get("CadoResponse_DefaultProject", 1)

        data: dict[str, Any] = self._http_request(
            method="GET",
            url_suffix="/tasks/pipelines",
            params={
                "project_id": project_id,
            },
        )

        pipelines: List[dict[str, Any]] = data["pipelines"]

        return pipelines[:limit]

    def list_instances(self, project_id: int | None, region: str | None, limit: int) -> List[dict[str, Any]]:
        """Calls the GET /api/v2/projects/{id}/imports/ec2 endpoint to
        retrieve details about a regions EC2 instances

        :param Optional[int] project_id: The id of the project to query available instances in
        :param Optional[str] region: The AWS region to search instances in

        :return JSON response from /projects/{id}/imports/ec2 endpoint
        :rtype Dict[str, Any]
        """

        if not project_id:
            project_id = demisto.params().get("CadoResponse_DefaultProject", 1)

        if not region:
            region = demisto.params().get("CadoResponse_DefaultRegion", "us-east-1")

        data: dict[str, Any] = self._http_request(
            method="GET", url_suffix=f"/projects/{project_id}/imports/ec2", params={"region": region}
        )

        instances: List[dict[str, Any]] = data["instances"]

        return instances[:limit]

    def list_buckets(self, project_id: int | None, limit: int) -> dict[str, Any]:
        """Calls the GET /api/v2/projects/{id}/imports/s3 endpoint to
        retrieve details about all the available S3 buckets

        :param Optional[int] project_id: The id of the project to query available buckets in

        :return JSON response from /projects/{id}/imports/s3 endpoint
        :rtype Dict[str, Any]
        """

        if not project_id:
            project_id = demisto.params().get("CadoResponse_DefaultProject", 1)

        data: dict[str, Any] = self._http_request(method="GET", url_suffix=f"/projects/{project_id}/imports/s3")

        data["buckets"] = data["buckets"][:limit]

        return data

    def trigger_instance_acquisition(
        self,
        project_id: int | None,
        instance_id: str | None,
        region: str | None,
        bucket: str | None,
        compress: bool = True,
        include_disks: bool = True,
        include_hash: bool = False,
        include_logs: bool = True,
        include_screenshot: bool = True,
    ) -> dict[str, Any]:
        """Calls the POST /api/v2/projects/{id}/imports/ec2 endpoint to
        trigger an acquisition of a given instance

        :param Optional[int] project_id: The ID of the project you wish to attach the acquisition to
        :param str instance_id: ID of the EC2 instance to acquire
        :param Optional[str] region: AWS region in which the EC2 instance is located
        :param Optional[str] bucket: S3 bucket where the uploaded disk image resides
        :param bool compress: Flag indicating if disk compression is enabled
        :param bool include_disks: Flag indicating if we include disk image in the acquisition
        :param bool include_hash: Flag indicating if we calculate the hash of the disk
        :param bool include_logs: Flag indicating if we include system logs in the acquisition
        :param bool include_screenshot: Flag indicating if we include a screenshot of the system in the acquisition

        :return JSON response from /projects/{id}/imports/ec2 endpoint
        :rtype Dict[str, Any]
        """

        if not project_id:
            project_id = demisto.params().get("CadoResponse_DefaultProject", 1)

        if not region:
            region = demisto.params().get("CadoResponse_DefaultRegion", "us-east-1")

        if not bucket:
            bucket = demisto.params().get("CadoResponse_DefaultBucket", "cado-default-bucket")

        payload: dict[str, Any] = {
            "bucket": bucket,
            "compress": compress,
            "include_disks": include_disks,
            "include_hash": include_hash,
            "include_logs": include_logs,
            "include_screenshot": include_screenshot,
            "instance_id": instance_id,
            "region": region,
        }

        return self._http_request(method="POST", url_suffix=f"/projects/{project_id}/imports/ec2", json_data=payload)

    def trigger_bucket_acquisition(self, project_id: int | None, bucket: str | None, file_name: str | None) -> dict[str, Any]:
        """Calls the POST /api/v2/projects/{id}/imports/s3 endpoint to
        trigger an acquisition of a given bucket or file

        :param Optional[int] project_id: The ID of the project you wish to attach the acquisition to
        :param Optional[str] bucket: The S3 bucket name containing the file
        :param str file_name: The name of the file to process

        :return JSON response from /projects/{id}/imports/ec2 endpoint
        :rtype Dict[str, Any]
        """

        if not project_id:
            project_id = demisto.params().get("CadoResponse_DefaultProject", 1)

        if not bucket:
            bucket = demisto.params().get("CadoResponse_DefaultBucket", "cado-default-bucket")

        payload: dict[str, Any] = {"bucket": bucket, "file_name": file_name}

        return self._http_request(method="POST", url_suffix=f"/projects/{project_id}/imports/s3", json_data=payload)


""" Command Line Handlers """


def test_module(client: Client) -> str:
    """Command handler for !test-module"""

    result: dict[str, Any] = client.heartbeat()
    status: str | None = result["status"]

    if status is not None and status == "Running":
        return "ok"

    return "Cado Response is not running"


def create_project_command(client: Client, args: dict[str, Any]) -> CommandResults:
    """Command handler for cado-create-project"""

    unix_timestamp: str = str(int(time.time()))
    project_name: str = args.get("project_name", unix_timestamp)
    project_description: str | None = args.get("project_description", None)
    result: dict[str, Any] = client.create_project(project_name, project_description)

    return CommandResults(outputs_prefix="CadoResponse.Project", outputs_key_field="id", outputs=result)


def list_project_command(client: Client, args: dict[str, Any]) -> CommandResults:
    """Command handler for cado-list-project"""

    project_id: int | None = args.get("project_id", None)
    limit: int = int(args.get("limit", 50))

    if project_id:
        result: Any = client.get_project(project_id)
    else:
        result = client.list_projects(limit)

    return CommandResults(outputs_prefix="CadoResponse.Projects", outputs_key_field="id", outputs=result)


def get_pipeline_command(client: Client, args: dict[str, Any]) -> CommandResults:
    """Command handler for cado-get-pipeline"""

    project_id: int | None = args.get("project_id", None)
    limit: int = int(args.get("limit", 50))
    pipeline_id: int | None = args.get("pipeline_id", None)

    if pipeline_id:
        result: CadoResponseCombinedOutput = client.get_pipeline(pipeline_id, project_id)
    else:
        result = client.list_pipelines(project_id, limit)

    return CommandResults(outputs_prefix="CadoResponse.Pipelines", outputs_key_field="pipeline_id", outputs=result)


def list_ec2_command(client: Client, args: dict[str, Any]) -> CommandResults:
    """Command handler for cado-list-ec2"""

    project_id: int | None = args.get("project_id", None)
    region: str | None = args.get("region", None)
    limit: int = int(args.get("limit", 100))
    result: List[dict[str, Any]] = client.list_instances(project_id, region, limit)

    return CommandResults(outputs_prefix="CadoResponse.EC2Instances", outputs_key_field="id", outputs=result)


def list_s3_command(client: Client, args: dict[str, Any]) -> CommandResults:
    """Command handler for cado-list-s3"""

    project_id: int | None = args.get("project_id", None)
    limit: int = int(args.get("limit", 100))
    result: dict[str, Any] = client.list_buckets(project_id, limit)

    return CommandResults(outputs_prefix="CadoResponse.S3Buckets", outputs=result)


def trigger_ec2_command(client: Client, args: dict[str, Any]) -> CommandResults:
    """Command handler for cado-trigger-ec2"""

    project_id: int | None = args.get("project_id", None)
    instance_id: str | None = args.get("instance_id", None)
    region: str | None = args.get("region", None)
    bucket: str | None = args.get("bucket", None)
    compress: bool = args.get("compress", True)
    include_disks: bool = args.get("include_disks", True)
    include_hash: bool = args.get("include_hash", False)
    include_logs: bool = args.get("include_logs", True)
    include_screenshot: bool = args.get("include_screenshot", True)

    if not instance_id:
        raise DemistoException("region is a required parameter!")

    result: dict[str, Any] = client.trigger_instance_acquisition(
        project_id, instance_id, region, bucket, compress, include_disks, include_hash, include_logs, include_screenshot
    )

    return CommandResults(outputs_prefix="CadoResponse.EC2Acquistion", outputs_key_field="pipeline_id", outputs=result)


def trigger_s3_command(client: Client, args: dict[str, Any]) -> CommandResults:
    """Command handler for cado-trigger-s3"""

    project_id: int | None = args.get("project_id", None)
    bucket: str | None = args.get("bucket", None)
    file_name: str | None = args.get("file_name", None)

    if not bucket:
        raise DemistoException("bucket is a required parameter!")

    if not file_name:
        raise DemistoException("file_name is a required parameter!")

    result: dict[str, Any] = client.trigger_bucket_acquisition(project_id, bucket, file_name)

    return CommandResults(
        outputs_prefix="CadoResponse.S3Acquisition", outputs_key_field="pipeline_id", outputs=result.get("pipelines")
    )


""" Helper Functions """


def enrich_errors(message: str, command: str) -> str:
    """Helper function to return better error messages.

    :param str message: Error message
    :param str command: Calling command

    :return: A better error message
    :rtype str
    """

    if command == "cado-create-project" and "Project name already exists" in message:
        return f'Project name {demisto.args().get("project_name")} already exists!'
    else:
        return f"Failed to execute {demisto.command()} command.\nError:\n{message}"


""" Entrypoint """


def main() -> None:
    api_key: str = demisto.params().get("apikey")
    base_url: str = urljoin(demisto.params()["url"], "/api/v2")
    verify_certificate: bool = not demisto.params().get("insecure", False)
    proxy: bool = demisto.params().get("proxy", False)
    command: str = demisto.command()
    args: dict[str, Any] = demisto.args()
    headers: dict[str, Any] = {"Authorization": f"Bearer {api_key}"}

    try:
        client: Client = Client(base_url=base_url, verify=verify_certificate, headers=headers, proxy=proxy)

        if command == "test-module":
            return_results(test_module(client))
        elif command == "cado-create-project":
            return_results(create_project_command(client, args))
        elif command == "cado-list-project":
            return_results(list_project_command(client, args))
        elif command == "cado-get-pipeline":
            return_results(get_pipeline_command(client, args))
        elif command == "cado-list-ec2":
            return_results(list_ec2_command(client, args))
        elif command == "cado-list-s3":
            return_results(list_s3_command(client, args))
        elif command == "cado-trigger-ec2":
            return_results(trigger_ec2_command(client, args))
        elif command == "cado-trigger-s3":
            return_results(trigger_s3_command(client, args))

    except Exception as e:
        message: str = str(e)

        if "404" in message:
            return_results(f"Nothing found for {command}")
        else:
            demisto.error(traceback.format_exc())
            return_error(enrich_errors(message, command), error=e)


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