Armorblox

import collections
import json
from typing import Any

import dateparser
import demistomock as demisto  # noqa: F401
import urllib3
from armorblox.client import Client as AbxBaseClient  # pylint: disable=E0401,E0611
from CommonServerPython import *  # noqa: F401

# disable insecure warnings
urllib3.disable_warnings()

DATE_FORMAT = "%Y-%m-%dT%H:%M:%SZ"
MAX_INCIDENTS_TO_FETCH = demisto.params().get("max_fetch")
FIRST_FETCH = demisto.params().get("first_fetch")
TENANT_NAME = demisto.params().get("tenantName")
INSECURE = demisto.params().get("insecure")
PROXY = demisto.params().get("proxy")
API_KEY = demisto.params().get("apikey")
verify_certificate = not demisto.params().get("insecure", False)
proxy = demisto.params().get("proxy", False)


class Client(AbxBaseClient):
    """Client class to interact with the service API
    This Client implements API calls, and does not contain any Demisto logic.
    Should only do requests and return data.
    """

    def get_incidents(self, orderBy="ASC", pageSize=None, pageToken=None, first_fetch=None):
        request_params = {"orderBy": orderBy}

        if pageToken == -1 and first_fetch:
            request_params["timeFilter"] = first_fetch
        elif pageToken and first_fetch:
            request_params["timeFilter"] = first_fetch
            request_params["pageToken"] = pageToken

        if pageSize:
            request_params["pageSize"] = pageSize

        response_json, next_page_token, total_count = self.incidents.list(params=request_params)
        return response_json, next_page_token

    def get_incident_details(self, incident_id):
        return self.incidents.get(incident_id)


def makehash():
    return collections.defaultdict(makehash)


def test_module(client: Client) -> str:  # pragma: no coverage
    """Tests API connectivity and authentication'
    Returning 'ok' indicates that the integration works like it is supposed to.
    Connection to the service is successful.
    Raises exceptions if something goes wrong.
    :type client: ``Client``
    :param Client: Armorblox client to use
    :type name: ``str``
    :return: 'ok' if test passed, anything else will fail the test.
    :rtype: ``str``
    """

    try:
        client.get_incidents(pageSize=1)

    except DemistoException as e:
        if "Forbidden" in str(e):
            return "Authorization Error: make sure API Key is correctly set"
        else:
            raise e
    return "ok"


def get_incidents_list(client, pageToken, first_fetch):
    """
    Hits the Armorblox API and returns the list of fetched incidents.
    """
    results, next_page_token = client.get_incidents(pageSize=MAX_INCIDENTS_TO_FETCH, pageToken=pageToken, first_fetch=first_fetch)
    # For each incident, get the details and extract the message_id
    for result in results:
        result["message_ids"] = get_incident_message_ids(client, result["id"])
    return results, next_page_token


def get_incident_message_ids(client, incident_id):
    """
    Returns the message ids for all the events for the input incident.
    """

    detail_response = client.get_incident_details(incident_id)
    message_ids = []
    # loop through all the events of this incident and collect the message ids
    if "events" in detail_response:
        for event in detail_response["events"]:
            message_ids.append(event["message_id"])

    if "abuse_events" in detail_response:
        for event in detail_response["abuse_events"]:
            message_ids.append(event["message_id"])
    return message_ids


def get_remediation_action(client, incident_id):
    """
    Returns the remediation action(s) for the input incident.
    """

    detail_response = client.get_incident_details(incident_id)
    remediation_actions = None
    if "remediation_actions" in detail_response:
        remediation_actions = detail_response["remediation_actions"][0]
    else:
        remediation_actions = None
    contxt = makehash()
    human_readable = makehash()
    human_readable["incident_id"] = incident_id
    human_readable["remediation_actions"] = remediation_actions
    contxt["incident_id"] = incident_id
    contxt["remediation_actions"] = remediation_actions
    return CommandResults(outputs_prefix="Armorblox.Threat", outputs=contxt)


def fetch_incidents_command(client):
    last_run = demisto.getLastRun()
    start_time: Any
    # pageToken fetched from demisto lastRun
    pageToken = 0
    incidents = []
    if "start_time" not in last_run:
        pageToken = -1
        response, next_page_token = client.get_incidents(pageSize=1, pageToken=pageToken, first_fetch=FIRST_FETCH)
        if response:
            response = response[0]
            start_time = response.get("date")
            start_time = dateparser.parse(start_time)
            message_ids = get_incident_message_ids(client, response.get("id"))
            response["message_ids"] = message_ids
            curr_incident = {"rawJSON": json.dumps(response), "details": json.dumps(response)}
            incidents.append(curr_incident)

    if last_run and "pageToken" in last_run:
        pageToken = last_run.get("pageToken")

    if last_run and "start_time" in last_run:
        start_time = dateparser.parse(last_run.get("start_time"))

    start_time = start_time.timestamp()
    incidents_data, pageToken = get_incidents_list(client, pageToken=pageToken, first_fetch=FIRST_FETCH)
    last_time = start_time

    for incident in incidents_data:
        dt = incident.get("date")
        parsed_date = dateparser.parse(dt)
        assert parsed_date is not None, f"failed parsing {dt}"
        dt = int(parsed_date.timestamp())
        # Update last run and add incident if the incident is newer than last fetch
        if dt > int(start_time):
            curr_incident = {"rawJSON": json.dumps(incident), "details": json.dumps(incident)}
            last_time = dt
            incidents.append(curr_incident)
    # Save the next_run as a dict with the start_time key to be stored
    demisto.setLastRun({"start_time": str(last_time), "pageToken": pageToken})
    return incidents


def main():  # pragma: no coverage
    """EXECUTION"""
    demisto.info(f"Command being called is {demisto.command()}")
    try:
        client = Client(api_key=API_KEY, instance_name=TENANT_NAME)
        if demisto.command() == "fetch-incidents":
            incident_results = fetch_incidents_command(client)
            demisto.incidents(incident_results)
            return_results("Incidents fetched successfully!!")
            # return_results(fetch_incidents_command(client))
        if demisto.command() == "armorblox-check-remediation-action":
            incident_id = demisto.args().get("incident_id")
            return_results(get_remediation_action(client, incident_id))

        elif demisto.command() == "test-module":
            result = test_module(client)
            return_results(result)
    except Exception as e:
        return_error(str(e))


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