Proofpoint Threat Protection

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

""" CONSTANTS """


CTX_PREFIX = "ProofpointThreatProtection"
OSL_HEADER = "Organizational Safe List"
OBL_HEADER = "Organizational Block List"
URL_SUFFIX_SAFELIST = "/emailProtection/modules/spam/orgSafeList"
URL_SUFFIX_BLOCKLIST = "/emailProtection/modules/spam/orgBlockList"
AUTH_HOST_BASE_URL = "https://auth.proofpoint.com/v1"
LIMIT_DEFAULT = 25
LIMIT_MAXIMUM = 100

""" INTEGRATION API CLIENT """


class Client(BaseClient):
    """Client class to interact with the service API"""

    def get_args(self):
        return demisto.args()

    def get_auth_host(self):
        return AUTH_HOST_BASE_URL

    def get_shared_integration_context(self):
        return get_integration_context()

    def set_shared_integration_context(self, context):
        return set_integration_context(context)

    def get_access_token(self, client_id, client_secret):
        """
        Get an access token that was previously created if it is still valid, else, generate a new access token from
        the API key and secret.
        """
        # Check if there is an existing valid access token
        integration_context = self.get_shared_integration_context()
        if integration_context.get("access_token") and integration_context.get("expiry_time") > date_to_timestamp(datetime.now()):
            self._headers = {"Authorization": f'Bearer {integration_context.get("access_token")}', "Accept": "application/json"}
            return integration_context.get("access_token")
        else:
            try:
                res = self._http_request(
                    method="POST",
                    full_url=f"{self.get_auth_host()}/token",
                    headers={"Content-Type": "application/x-www-form-urlencoded"},
                    data={"grant_type": "client_credentials", "client_id": client_id, "client_secret": client_secret},
                )
                access_token = res.get("access_token", None)
                if access_token is not None:
                    expiry_time = date_to_timestamp(datetime.now(), date_format="%Y-%m-%dT%H:%M:%S")
                    expiry_time += res.get("expires_in", 0) * 1000 - 10
                    context = {"access_token": access_token, "expiry_time": expiry_time}
                    self.set_shared_integration_context(context)
                    self._headers = {"Authorization": f"Bearer {access_token}", "Accept": "application/json"}
                    return res.get("access_token")

            except Exception as e:
                raise (
                    Exception(
                        f"Error occurred while creating an access token. "
                        f"Please check the instance configuration."
                        f"\n\n{e.args[0]}"
                    )
                )

    def get_safelist(self, cluster_id):
        return self._http_request("GET", url_suffix=URL_SUFFIX_SAFELIST, headers=self._headers, params={"clusterId": cluster_id})

    def safelist_add_delete(self, cluster_id, args, action):
        import json

        self._headers.update({"Content-Type": "application/json"})
        return self._http_request(
            "POST",
            resp_type="response",
            url_suffix=URL_SUFFIX_SAFELIST,
            headers=self._headers,
            params={"clusterId": cluster_id},
            data=json.dumps(
                {
                    "action": action,
                    "attribute": args.get("attribute"),
                    "operator": args.get("operator"),
                    "value": args.get("value"),
                    "comment": args.get("comment"),
                }
            ),
        )

    def get_blocklist(self, cluster_id):
        return self._http_request("GET", url_suffix=URL_SUFFIX_BLOCKLIST, headers=self._headers, params={"clusterId": cluster_id})

    def blocklist_add_delete(self, cluster_id, args, action):
        import json

        self._headers.update({"Content-Type": "application/json"})
        return self._http_request(
            "POST",
            resp_type="response",
            url_suffix=URL_SUFFIX_BLOCKLIST,
            headers=self._headers,
            params={"clusterId": cluster_id},
            data=json.dumps(
                {
                    "action": action,
                    "attribute": args.get("attribute"),
                    "operator": args.get("operator"),
                    "value": args.get("value"),
                    "comment": args.get("comment"),
                }
            ),
        )


""" HELPER FUNCTIONS """


def make_return_command_results(header_str: str, listname: str, op: str | dict) -> CommandResults:
    readable_op = tableToMarkdown(header_str, op, headers={} if not isinstance(op, str) else ["Outcome"], removeNull=True)
    return_cr = CommandResults(readable_output=readable_op, outputs_prefix=CTX_PREFIX, outputs={listname: op})

    return return_cr


def get_limit_args(client: Client):
    args = client.get_args()
    limit = args.get("limit", str(LIMIT_DEFAULT))
    limit = int(limit) if limit.isnumeric() else LIMIT_DEFAULT
    limit = min(LIMIT_MAXIMUM, limit)
    all_results = args.get("all_results", "false")
    all_results = all_results.lower() == "true"

    return limit, all_results


""" COMMAND FUNCTIONS """


def module_test_command(client: Client, cluster_id: str) -> str:
    client.get_blocklist(cluster_id)
    return "ok"


def safelist_list_command(client: Client, cluster_id: str) -> CommandResults:
    limit, all_results = get_limit_args(client)
    res = client.get_safelist(cluster_id)
    res_rc = make_return_command_results(
        OSL_HEADER, "Safelist", res.get("entries") if all_results else res.get("entries")[-limit:]
    )
    return res_rc


def safelist_add_command(client: Client, cluster_id: str) -> CommandResults:
    client.safelist_add_delete(cluster_id, client.get_args(), "add")
    res_rc = make_return_command_results(OSL_HEADER, "Safelist Entry Added", "Success")
    return res_rc


def safelist_delete_command(client: Client, cluster_id: str) -> CommandResults:
    client.safelist_add_delete(cluster_id, client.get_args(), "delete")
    res_rc = make_return_command_results(OSL_HEADER, "Safelist Entry Deleted", "Success")
    return res_rc


def blocklist_list_command(client: Client, cluster_id: str) -> CommandResults:
    limit, all_results = get_limit_args(client)
    res = client.get_blocklist(cluster_id)
    res_rc = make_return_command_results(
        OBL_HEADER, "Blocklist", res.get("entries") if all_results else res.get("entries")[-limit:]
    )
    return res_rc


def blocklist_add_command(client: Client, cluster_id: str) -> CommandResults:
    client.blocklist_add_delete(cluster_id, client.get_args(), "add")
    res_rc = make_return_command_results(OBL_HEADER, "Blocklist Entry Added", "Success")
    return res_rc


def blocklist_delete_command(client: Client, cluster_id: str) -> CommandResults:
    client.blocklist_add_delete(cluster_id, client.get_args(), "delete")
    res_rc = make_return_command_results(OBL_HEADER, "Blocklist Entry Deleted", "Success")
    return res_rc


"""         UPDATE COMMAND MAPPINGS
    (After command function declarations) """

COMMANDS = {
    "test-module": module_test_command,
    "proofpoint-tp-safelist-list": safelist_list_command,
    "proofpoint-tp-safelist-add-entry": safelist_add_command,
    "proofpoint-tp-safelist-delete-entry": safelist_delete_command,
    "proofpoint-tp-blocklist-list": blocklist_list_command,
    "proofpoint-tp-blocklist-add-entry": blocklist_add_command,
    "proofpoint-tp-blocklist-delete-entry": blocklist_delete_command,
}


""" MAIN FUNCTION """


def parse_params(params: dict):
    client_id = params.get("credentials", {}).get("username", params.get("credentials", {}).get("identifier"))
    client_secret = params.get("credentials", {}).get("password")
    base_url = params.get("url")
    cluster_id = params.get("cluster_id")
    verify_certificate = not params.get("insecure", False)
    proxy = params.get("proxy", False)

    return client_id, client_secret, base_url, cluster_id, verify_certificate, proxy


def main() -> None:
    client_id, client_secret, base_url, cluster_id, verify_certificate, proxy = parse_params(demisto.params())

    command = demisto.command()

    try:
        if command in COMMANDS:
            client = Client(base_url=base_url, verify=verify_certificate, proxy=proxy)

            client.get_access_token(client_id, client_secret)
            return_results(COMMANDS[command](client, cluster_id))
        else:
            raise NotImplementedError(f"Command {command} is not implemented")

    except Exception as e:
        return_error(f"Failed to execute {command} command.\nError:\n{e!s}")


""" ENTRY POINT """


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