Pentera

import csv
import io
import json
from enum import Enum

import demistomock as demisto
import jwt
import requests

# Disable insecure warnings
import urllib3
from CommonServerPython import *

from CommonServerUserPython import *

urllib3.disable_warnings()

HEADERS = {"Accept": "application/json"}
DATE_FORMAT = "%Y-%m-%dT%H:%M:%SZ"

HEALTH_URL_SUFFIX = "/health/dbChecks"
AUTH_URL_SUFFIX = "/auth/token"
LIST_TEMPLATES_URL_SUFFIX = "/api/v1/templates"
CANCEL_TASK_URL_SUFFIX = "/api/v1/taskRun/{taskRunId}/cancel"
GET_TASK_RUN_STATUS_URL_SUFFIX = "/api/v1/taskRun/{taskRunId}"
RUN_BULK_URL_SUFFIX = "/api/v1/template/runBulk"
EXPORT_CSV_URL_SUFFIX = "/api/v1/taskRun/{taskRunId}/fullActionReportCSV"


class Request(Enum):
    POST = "POST"
    GET = "GET"


class AuthorizationError(Exception):
    pass


class Client(BaseClient):
    def __init__(self, base_url: str, tgt: str, client_id: str, verify: bool, proxy: bool, headers):
        super().__init__(base_url=f"{base_url}", headers=headers, verify=verify, proxy=proxy)
        self.session = requests.Session()
        self.session.headers = headers
        self.client_id = client_id
        self.tgt = tgt
        self.access_token = ""
        self.expiry = 0
        self.load_session_parameters()

    def load_session_parameters(self):
        context: dict = get_integration_context()
        if context and context["base_url"] == self._base_url:
            self.tgt = context["tgt"]
            self.access_token = context["accessToken"]
            self.expiry = context["expiry"]

    def generic_request(
        self,
        method: str,
        url_suffix: str = None,
        full_url: str = None,
        headers: dict = None,
        params: dict = None,
        data: dict = None,
        response_type: str = "json",
    ):
        full_url = full_url if full_url else f"{self._base_url}{url_suffix}"
        headers = headers if headers else self._headers
        try:
            res = self.session.request(method, full_url, headers=headers, verify=self._verify, data=data, params=params)
            demisto.debug(f"Got response: {res}")
            if not res.ok:
                status_code = res.status_code
                if status_code == requests.status_codes.codes.UNAUTHORIZED:  # pylint: disable=no-member
                    info = "Check that your system clock is set to the correct date and time before you try again."
                    raise AuthorizationError(f"Status code: {status_code}, reason: {res.text}. {info}")
                raise ValueError(f"Error in API call to Pentera. Status code: {status_code}, reason: {res.text}")

            try:
                if response_type == "json":
                    demisto.debug("result is JSON")
                    return res.json()
                demisto.debug("result is TEXT")
                return res.text
            except Exception:
                raise ValueError(f"Failed to parse http response to JSON format. Original response body: \n{res.text}")

        except requests.exceptions.ConnectTimeout as exception:
            err_msg = (
                "Connection Timeout Error - potential reasons might be that the Server URL parameter"
                " is incorrect or that the Server is not accessible from your host."
            )
            raise DemistoException(err_msg, exception)

        except requests.exceptions.SSLError as exception:
            err_msg = (
                "SSL Certificate Verification Failed - try selecting 'Trust any certificate' checkbox in"
                " the integration configuration."
            )
            raise DemistoException(err_msg, exception)

        except requests.exceptions.ProxyError as exception:
            err_msg = (
                "Proxy Error - if the 'Use system proxy' checkbox in the integration configuration is"
                " selected, try clearing the checkbox."
            )
            raise DemistoException(err_msg, exception)

        except requests.exceptions.ConnectionError as exception:
            err_msg = getattr(exception, "message", str(exception))
            raise DemistoException(err_msg, exception)

        except Exception as request_error:
            message = getattr(request_error, "message", str(request_error))
            raise DemistoException(f"Could not send request to Pentera, reason: {message}", exception=request_error)

    def authenticate(self):
        data = {"client_id": self.client_id, "tgt": self.tgt}
        res = self.generic_request(method=Request.POST.value, url_suffix=AUTH_URL_SUFFIX, data=data)
        self.tgt = res.get("tgt")
        self.access_token = res.get("token")
        jwt_decode_dict = jwt.get_unverified_header(self.access_token)
        self.expiry = jwt_decode_dict.get("exp", 0) if jwt_decode_dict else 0
        self.save_session_parameters()

    def save_session_parameters(self):
        context = {"base_url": self._base_url, "tgt": self.tgt, "accessToken": self.access_token, "expiry": self.expiry}
        set_integration_context(context)

    def is_access_token_valid(self):
        if not self.access_token or not self.expiry or self.expiry < int(datetime.utcnow().timestamp()):  # noqa: SIM103
            return False
        return True

    def create_basic_authentication_header(self):
        authentication_headers = HEADERS.copy()
        token = self.access_token + ":"
        encoded_bytes = base64.b64encode(token.encode("utf-8"))
        encoded_str = str(encoded_bytes, "utf-8")
        authentication_headers["Authorization"] = "Basic " + encoded_str
        return authentication_headers

    def run_health_checks(self):
        res = self.generic_request(method=Request.GET.value, url_suffix=HEALTH_URL_SUFFIX)
        return res

    def run_template_by_name(self, template_name):
        headers = self.create_basic_authentication_header()
        data = {"templateNames": [template_name]}
        res = self.generic_request(method=Request.POST.value, url_suffix=RUN_BULK_URL_SUFFIX, headers=headers, data=data)
        return res

    def get_task_run_status_by_task_run_id(self, task_run_id: str):
        headers = self.create_basic_authentication_header()
        url_suffix = GET_TASK_RUN_STATUS_URL_SUFFIX.format(taskRunId=task_run_id)
        res = self.generic_request(method=Request.GET.value, url_suffix=url_suffix, headers=headers, data={})
        task_status = res.get("taskRuns")[0]
        return task_status

    def get_task_run_full_action_report_by_task_run_id(self, task_run_id: str):
        headers = self.create_basic_authentication_header()
        url_suffix = EXPORT_CSV_URL_SUFFIX.format(taskRunId=task_run_id)
        res = self.generic_request(method=Request.GET.value, url_suffix=url_suffix, headers=headers, response_type="csv")
        return res


def pentera_test_module_command(client: Client):
    try:
        response = client.run_health_checks()
    except Exception as test_error:
        message = getattr(test_error, "message", str(test_error))
        raise DemistoException(message)
    exceptions: list = response.get("exceptions")
    if exceptions:
        raise DemistoException(", ".join(exceptions))
    return "ok"


def pentera_run_template_command(client: Client, args):
    template_name = args.get("template_name")
    try:
        response = client.run_template_by_name(template_name)
        task_run_json = response.get("taskRuns")[0]
        parsed_response = parse_task_run_status(task_run_json)
        readable_output = tableToMarkdown(template_name, parse_task_run_status(task_run_json), removeNull=True)
        return (
            readable_output,
            {"Pentera.TaskRun(val.ID == obj.ID)": parsed_response},
            response,  # raw response - the original response
        )

    except Exception as run_template_error:
        message = getattr(run_template_error, "message", str(run_template_error))
        raise DemistoException(
            f"Could not run template with template_name: '{template_name}', reason: {message}", exception=run_template_error
        )


def pentera_get_task_run_status_command(client: Client, args):
    task_run_id = args.get("task_run_id")
    try:
        task_run_status = client.get_task_run_status_by_task_run_id(task_run_id)
        parsed_response = parse_task_run_status(task_run_status)
        title = parsed_response["TemplateName"] + ": " + parsed_response["Status"]
        readable_output = tableToMarkdown(title, parsed_response, removeNull=True)
        return (
            readable_output,
            {"Pentera.TaskRun(val.ID == obj.ID)": parsed_response},
            task_run_status,  # raw response - the original response
        )
    except Exception as status_error:
        message = getattr(status_error, "message", str(status_error))
        raise DemistoException(
            f"Could not get task run status for task_run_id: '{task_run_id}', reason: {message}", exception=status_error
        )


def pentera_get_task_run_full_action_report_command(client: Client, args):
    def _convert_csv_file_to_dict(csv_file):
        def _map_parameters_string_to_object(str_parameters: str = None):
            if str_parameters:
                return json.loads(str_parameters)
            return None

        csv_reader = csv.DictReader(io.StringIO(csv_file))
        data = []
        for row in csv_reader:
            row_copy = row.copy()
            converted_params = _map_parameters_string_to_object(row_copy.get("Parameters"))
            if converted_params:
                row_copy["Parameters"] = converted_params
            data.append(row_copy)
        return data

    def _convert_full_action_report_time(full_action_report_list: list[dict]):
        def _parse_date(full_date, separator):
            if isinstance(full_date, str) and isinstance(separator, str):
                date = full_date.split(separator)
                if len(date) > 2:
                    first_arg = date[0]
                    second_arg = date[1]
                    third_arg = date[2]
                    return first_arg, second_arg, third_arg
                return None
            return None

        res_list: list[dict] = []
        for ordered_dict in full_action_report_list:
            full_date_to_convert = ordered_dict["Time"]
            full_date_list = full_date_to_convert.split(" ")
            year, month, day = _parse_date(full_date_list[0], "-")
            hours, minutes, seconds = _parse_date(full_date_list[1], ":")
            converted_date = year + "-" + month + "-" + day + "T" + hours + ":" + minutes + ":" + seconds + "Z"
            new_ordered_dict = ordered_dict.copy()
            new_ordered_dict["Time"] = converted_date
            res_list.append(new_ordered_dict)
        return res_list

    entries = []
    task_run_id = args.get("task_run_id")
    try:
        response_csv = client.get_task_run_full_action_report_by_task_run_id(task_run_id)
        readable_output = f"# Pentera Report for TaskRun ID {task_run_id}"
        entry = fileResult(f"penterascan-{task_run_id}.csv", response_csv, entryTypes["entryInfoFile"])
        entry["HumanReadable"] = readable_output
        entry["ContentsFormat"] = formats["markdown"]
        entries.append(entry)
        csv_dict = _convert_csv_file_to_dict(response_csv)
        date_converted_csv_dict = _convert_full_action_report_time(csv_dict)
        human_readable = tableToMarkdown(readable_output, date_converted_csv_dict)
        entries.append(
            {
                "Type": entryTypes["note"],
                "ContentsFormat": formats["json"],
                "ReadableContentsFormat": formats["markdown"],
                "Contents": date_converted_csv_dict,
                "EntryContext": {
                    "Pentera.TaskRun(val.ID == obj.ID)": {"FullActionReport": date_converted_csv_dict, "ID": task_run_id}
                },
                "HumanReadable": human_readable,
            }
        )
        return entries
    except Exception as report_error:
        message = getattr(report_error, "message", str(report_error))
        raise DemistoException(
            f"Could not get full action report for task_run_id: '{task_run_id}', reason: {message}", exception=report_error
        )


def parse_task_run_status(json_response):
    def _convert_time_in_millis_to_date_format(time_in_millis):
        time_in_date_format = None
        try:
            time_in_date_format = datetime.fromtimestamp(float(time_in_millis) / 1000).strftime(DATE_FORMAT)
            return time_in_date_format
        except TypeError:
            return time_in_date_format

    if isinstance(json_response, dict):
        end_time_date_format = _convert_time_in_millis_to_date_format(json_response.get("endTime"))
        start_time_date_format = _convert_time_in_millis_to_date_format(json_response.get("startTime"))
        parsed_json_response = {
            "ID": json_response.get("taskRunId"),
            "TemplateName": json_response.get("taskRunName"),
            "StartTime": start_time_date_format,
            "EndTime": end_time_date_format,
            "Status": json_response.get("status"),
        }
        return parsed_json_response
    return None


def pentera_authentication(client: Client):
    if not client.is_access_token_valid():
        try:
            client.authenticate()
        except Exception as auth_error:
            message = getattr(auth_error, "message", str(auth_error))
            raise DemistoException(f"Could not authenticate to Pentera, reason: {message}", exception=auth_error)


def increase_csv_field_size_limit():
    """
    This method will try to increase the csv field size limit as files might contain huge fields.
    :return: None
    """
    try:
        csv.field_size_limit(sys.maxsize)
    except OverflowError:
        pass


def main():
    params: dict = demisto.params()
    application_port = params["port"]
    base_url = params["url"].rstrip("/") + ":" + application_port
    client_id = params["clientId"]
    tgt = params["tgt"]
    verify_certificate = not params.get("insecure", False)
    proxy = params.get("proxy", False)
    client = Client(base_url=base_url, tgt=tgt, verify=verify_certificate, client_id=client_id, proxy=proxy, headers=HEADERS)
    command = demisto.command()
    demisto.debug(f"Got command: {command}")
    try:
        if demisto.command() == "test-module":
            demisto.results(pentera_test_module_command(client))
        else:
            pentera_authentication(client)
            if demisto.command() == "pentera-run-template-by-name":
                return_outputs(*pentera_run_template_command(client, demisto.args()))
            elif demisto.command() == "pentera-get-task-run-status":
                return_outputs(*pentera_get_task_run_status_command(client, demisto.args()))
            elif demisto.command() == "pentera-get-task-run-full-action-report":
                demisto.results(pentera_get_task_run_full_action_report_command(client, demisto.args()))
    except Exception as e:
        return_error(f'Failed to execute command: {command}, {getattr(e, "message", str(e))}', error=e)


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