AttackIQFireDrill

import demistomock as demisto
from CommonServerPython import *
from CommonServerUserPython import *

""" IMPORTS """
from requests import HTTPError
from typing import Any
from json.decoder import JSONDecodeError

import json
import traceback
import requests
import math
import urllib3

# Disable insecure warnings
urllib3.disable_warnings()

""" GLOBALS/PARAMS """

TOKEN = demisto.params().get("token_creds", {}).get("password") or demisto.params().get("token")
# Remove trailing slash to prevent wrong URL path to service
SERVER = (
    demisto.params().get("url")[:-1]
    if ("url" in demisto.params() and demisto.params()["url"].endswith("/"))
    else demisto.params().get("url", "")
)
# Should we use SSL
USE_SSL = not demisto.params().get("insecure", False)
# Headers to be sent in requests
HEADERS = {"Authorization": f"Token {TOKEN}", "Content-Type": "application/json", "Accept": "application/json"}

# Error messages
INVALID_ID_ERR_MSG = "Error in API call. This may be happen if you provided an invalid id."
API_ERR_MSG = "Error in API call to AttackIQ. "
DEFAULT_PAGE_SIZE = 10

# Transformation dicts
ASSESSMENTS_TRANS = {
    "id": "Id",
    "name": "Name",
    "user": "User",
    "users": "Users",
    "owner": "Owner",
    "groups": "Groups",
    "creator": "Creator",
    "created": "Created",
    "end_date": "EndDate",
    "modified": "Modified",
    "start_date": "StartDate",
    "description": "Description",
    "project_state": "AssessmentState",
    "master_job_count": "MasterJobCount",
    "default_schedule": "DefaultSchedule",
    "default_asset_count": "DefaultAssetCount",
    "project_template.id": "AssessmentTemplateId",
    "default_asset_group_count": "DefaultAssetGroupCount",
    "project_template.company": "AssessmentTemplateCompany",
    "project_template.created": "AssessmentTemplateCreated",
    "project_template.modified": "AssessmentTemplateModified",
    "project_template.template_name": "AssessmentTemplateName",
    "project_template.default_schedule": "AssessmentTemplateDefaultSchedule",
    "project_template.template_description": "AssessmentTemplateDescription",
}

TESTS_TRANS = {
    "id": "Id",
    "name": "Name",
    "description": "Description",
    "project": "Assessment",
    "total_asset_count": "TotalAssetCount",
    "cron_expression": "CronExpression",
    "runnable": "Runnable",
    "last_result": "LastResult",
    "user": "User",
    "created": "Created",
    "modified": "Modified",
    "using_default_schedule": "UsingDefaultSchedule",
    "using_default_assets": "UsingDefaultAssets",
    "latest_instance_id": "LatestInstanceId",
    "scenarios": {"name": "Name", "id": "Id"},
    "assets": {
        "id": "Id",
        "ipv4_address": "Ipv4Address",
        "hostname": "Hostname",
        "product_name": "ProductName",
        "modified": "Modified",
        "status": "Status",
    },
}

TEST_STATUS_TRANS = {
    "detected": "Detected",
    "failed": "Failed",
    "finished": "Finished",
    "passed": "Passed",
    "errored": "Errored",
    "total": "Total",
}

TEST_RESULT_TRANS = {
    "id": "Id",
    "modified": "Modified",
    "project_id": "Assessment.Id",
    "project_name": "Assessment.Name",
    "scenario.id": "Scenario.Id",
    "scenario.name": "Scenario.Name",
    "scenario.description": "Scenario.Description",
    "asset.id": "Asset.Id",
    "asset.ipv4_address": "Asset.Ipv4Address",
    "asset.hostname": "Asset.Hostname",
    "asset.product_name": "Asset.ProductName",
    "asset.modified": "Asset.Modified",
    "asset_group": "Asset.AssetGroup",
    "job_state_name": "JobState",
    "outcome_name": "Outcome",
}

""" HELPER FUNCTIONS """


def http_request(method, url_suffix, params=None, data=None):
    url = urljoin(SERVER, url_suffix)
    LOG(
        f"AttackIQ is attempting {method} request sent to {url} with params:\n{json.dumps(params, indent=4)} \n "
        f'data:\n"{json.dumps(data)}'
    )
    try:
        res = requests.request(method, url, verify=USE_SSL, params=params, data=data, headers=HEADERS)
        if res.status_code == 204:
            return ""
        # Handle error responses gracefully
        if res.status_code not in {200, 201}:
            error_reason = get_http_error_reason(res)
            raise HTTPError(f"[{res.status_code}] - {error_reason}")
        try:
            return res.json()
        except JSONDecodeError:
            return_error(
                "Response contained no valid body. See logs for more information.",
                error=f"AttackIQ response body:\n{res.content!r}",
            )
    except requests.exceptions.ConnectionError as e:
        LOG(str(e))
        return_error("Encountered issue reaching the endpoint, please check that you entered the URL correctly.")


def get_http_error_reason(res):
    """
    Get error reason from an AttackIQ http error
    Args:
        res: AttackIQ response

    Returns: Reason for http error
    """
    err_reason = res.reason
    try:
        res_json = res.json()
        if "detail" in res_json:
            err_reason = f'{err_reason}. {res_json["detail"]}'
    except JSONDecodeError:
        pass
    return err_reason


def build_transformed_dict(src, trans_dict):
    """Builds a dictionary according to a conversion map

    Args:
        src (dict): original dictionary to build from
        trans_dict (dict): dict in the format { 'OldKey': 'NewKey', ...}

    Returns: src copy with changed keys
    """
    if isinstance(src, list):
        return [build_transformed_dict(x, trans_dict) for x in src]
    res: dict[str, Any] = {}
    for key, val in trans_dict.items():
        if isinstance(val, dict):
            # handle nested list
            sub_res = res
            item_val = [build_transformed_dict(item, val) for item in (demisto.get(src, key) or [])]
            key = underscoreToCamelCase(key)
            for sub_key in key.split(".")[:-1]:
                if sub_key not in sub_res:
                    sub_res[sub_key] = {}
                sub_res = sub_res[sub_key]
            sub_res[key.split(".")[-1]] = item_val
        elif "." in val:
            # handle nested vals
            update_nested_value(res, val, to_val=demisto.get(src, key))
        else:
            res[val] = demisto.get(src, key)
    return res


def create_invalid_id_err_msg(orig_err, error_codes):
    """
    Creates an 'invalid id' error message
    Args:
        orig_err (str): The original error message
        error_codes (list): List of error codes to look for

    Returns (str): Error message for invalid id
    """
    err_msg = API_ERR_MSG
    if any(err_code in orig_err for err_code in error_codes):
        err_msg += "This may be happen if you provided an invalid id.\n"
    err_msg += orig_err
    return err_msg


def update_nested_value(src_dict, to_key, to_val):
    """
    Updates nested value according to transformation dict structure where 'a.b' key will create {'a': {'b': val}}
    Args:
        src_dict (dict): The original dict
        to_key (str): Key to transform to (expected to contain '.' to mark nested)
        to_val: The value that'll be put under the nested key
    """
    sub_res = src_dict
    to_key_lst = to_key.split(".")
    for sub_to_key in to_key_lst[:-1]:
        if sub_to_key not in sub_res:
            sub_res[sub_to_key] = {}
        sub_res = sub_res[sub_to_key]
    sub_res[to_key_lst[-1]] = to_val


def get_page_number_and_page_size(args):
    """
    Get arguments page_number and page_size from args
    Args:
        args (dict): Argument dictionary, with possible page_number and page_size keys

    Returns (int, int): Return a tuple of (page_number, page_size)

    """
    page = args.get("page_number", 1)
    page_size = args.get("page_size", DEFAULT_PAGE_SIZE)
    err_msg_format = 'Error: Invalid {arg} value. "{val}" Is not a valid value. Please enter a positive integer.'
    try:
        page = int(page)
        if page <= 0:
            raise ValueError
    except (ValueError, TypeError):
        return_error(err_msg_format.format(arg="page_number", val=page))
    try:
        page_size = int(page_size)
        if page_size <= 0:
            raise ValueError
    except (ValueError, TypeError):
        return_error(err_msg_format.format(arg="page_size", val=page_size))
    return page, page_size


""" COMMANDS + REQUESTS FUNCTIONS """


def test_module():
    """
    Performs basic get request to get item samples
    """
    http_request("GET", "/v1/assessments")
    demisto.results("ok")


""" COMMANDS MANAGER / SWITCH PANEL """


def activate_assessment_command():
    """Implements attackiq-activate-assessment command"""
    ass_id = demisto.getArg("assessment_id")
    try:
        raw_res = http_request("POST", f"/v1/assessments/{ass_id}/activate")
        hr = raw_res.get("message", f"Assessment {ass_id} activation was sent successfully.")
        demisto.results(hr)
    except HTTPError as e:
        return_error(create_invalid_id_err_msg(str(e), ["403"]))


def get_assessment_execution_status_command():
    """Implements attackiq-get-assessment-execution-status command"""
    ass_id = demisto.getArg("assessment_id")
    try:
        raw_res = http_request("GET", f"/v1/assessments/{ass_id}/is_on_demand_running")
        ex_status = raw_res.get("message")
        hr = f'Assessment {ass_id} execution is {"" if ex_status else "not "}running.'
        ec = {"AttackIQ.Assessment(val.Id === obj.Id)": {"Running": ex_status, "Id": ass_id}}
        return_outputs(hr, ec, raw_res)
    except HTTPError as e:
        return_error(create_invalid_id_err_msg(str(e), ["403"]))


def get_test_execution_status_command():
    """Implements attackiq-get-test-execution-status command"""
    test_id = demisto.getArg("test_id")
    try:
        raw_test_status = http_request("GET", f"/v1/tests/{test_id}/get_status")
        test_status = build_transformed_dict(raw_test_status, TEST_STATUS_TRANS)
        test_status["Id"] = test_id
        hr = tableToMarkdown(f"Test {test_id} status", test_status)
        return_outputs(hr, {"AttackIQTest(val.Id === obj.Id)": test_status}, raw_test_status)
    except HTTPError as e:
        return_error(create_invalid_id_err_msg(str(e), ["500"]))


def build_test_results_hr(test_results, test_id, page, tot_pages):
    """
    Creates test results human readable
    Args:
        page (int): Current page
        tot_pages (int): Total pages
        test_results (list): Results of the test (after being transformed)
        test_id (str): ID of the test

    Returns: Human readable of test results
    """
    keys = ["Assessment Name", "Scenario Name", "Hostname", "Asset IP", "Job State", "Modified", "Outcome"]
    test_results_mod = []
    for t_res in test_results:
        assessment = t_res.get("Assessment")
        asset = t_res.get("Asset")
        scenario = t_res.get("Scenario")
        hr_items = {
            keys[0]: assessment.get("Name"),
            keys[1]: scenario.get("Name"),
            keys[2]: asset.get("Hostname"),
            keys[3]: asset.get("Ipv4Address"),
            keys[4]: demisto.get(t_res, "JobState"),
            keys[5]: t_res.get("Modified"),
            keys[6]: demisto.get(t_res, "Outcome.Name"),
        }
        test_results_mod.append(hr_items)
    return tableToMarkdown(f"Test Results for {test_id}\n ### Page {page}/{tot_pages}", test_results_mod, keys)


def get_test_results(page, page_size, test_id, show_last_res):
    """
    Get test results response
    Args:
        page (int): Page number
        page_size (int): Page size
        test_id (int): ID of test
        show_last_res (bool): Flag for showing only last result

    Returns: Test results
    """
    params = {"page": page, "page_size": page_size, "test_id": test_id, "show_last_result": show_last_res}
    return http_request("GET", "/v1/results", params=params)


def get_test_results_command(args=demisto.args()):
    """Implements attackiq-get-test-results command"""
    test_id = args.get("test_id")
    outcome_filter = args.get("outcome_filter")
    page, page_size = get_page_number_and_page_size(demisto.args())
    try:
        raw_test_res = get_test_results(page, page_size, test_id, args.get("show_last_result") == "True")
        test_cnt = raw_test_res.get("count")
        if test_cnt == 0:
            return_outputs("No results were found", {})
        else:
            total_pages = math.ceil(test_cnt / page_size)
            remaining_pages = total_pages - page
            if remaining_pages < 0:
                remaining_pages = 0
            test_res = build_transformed_dict(raw_test_res["results"], TEST_RESULT_TRANS)
            if outcome_filter:
                test_res = list(filter(lambda x: x.get("Outcome") == outcome_filter, test_res))
            context = {
                "AttackIQTestResult(val.Id === obj.Id)": test_res,
                "AttackIQTestResult(val.Count).Count": test_cnt,
                "AttackIQTestResult(val.RemainingPages).RemainingPages": remaining_pages,
            }
            hr = build_test_results_hr(test_res, test_id, page, total_pages)
            return_outputs(hr, context, raw_test_res)
    except HTTPError as e:
        return_error(create_invalid_id_err_msg(str(e), ["500"]))


def get_assessments(page="1", assessment_id=None, page_size=DEFAULT_PAGE_SIZE):
    """
    Fetches assessments from attackIQ
    Args:
        page (str or int): Page number to fetch
        assessment_id (str): (Optional) If provided will fetch only the assessment with matching ID

    Returns: Assessments from attackIQ
    """
    params = {"page_size": page_size, "page": page}
    if assessment_id:
        return http_request("GET", f"/v1/assessments/{assessment_id}")
    return http_request("GET", "/v1/assessments", params=params)


def list_assessments_command():
    """Implements attackiq-list-assessments command"""
    page, page_size = get_page_number_and_page_size(demisto.args())
    raw_assessments = get_assessments(page=page, page_size=page_size)
    assessments_res = build_transformed_dict(raw_assessments.get("results"), ASSESSMENTS_TRANS)
    ass_cnt = raw_assessments.get("count")
    total_pages = math.ceil(ass_cnt / page_size)
    remaining_pages = total_pages - page
    if remaining_pages < 0:
        remaining_pages = 0
    context = {
        "AttackIQ.Assessment(val.Id === obj.Id)": assessments_res,
        "AttackIQ.Assessment(val.Count).Count": ass_cnt,
        "AttackIQ.Assessment(val.RemainingPages).RemainingPages": remaining_pages,
    }
    hr = tableToMarkdown(
        f"AttackIQ Assessments Page {page}/{total_pages}",
        assessments_res,
        headers=["Id", "Name", "Description", "User", "Created", "Modified"],
    )
    return_outputs(hr, context, raw_assessments)


def get_assessment_by_id_command():
    """Implements attackiq-get-assessment-by-id command"""
    assessment_id = demisto.getArg("assessment_id")
    try:
        raw_assessments = get_assessments(assessment_id=assessment_id)
        assessments_res = build_transformed_dict(raw_assessments, ASSESSMENTS_TRANS)
        hr = tableToMarkdown(
            f"AttackIQ Assessment {assessment_id}",
            assessments_res,
            headers=["Id", "Name", "Description", "User", "Created", "Modified"],
        )
        return_outputs(hr, {"AttackIQ.Assessment(val.Id === obj.Id)": assessments_res}, raw_assessments)
    except HTTPError as e:
        return_error(create_invalid_id_err_msg(str(e), ["403"]))


def build_tests_hr(tests_res, ass_id, page_num, tot_pages):
    """
    Creates tests human readable
    Args:
        tot_pages (int): Total pages
        page_num (int): Current page
        ass_id (str): Assignment ID
        tests_res (list): Transformed result of test

    Returns: Human readable string (md format) of tests
    """
    hr = f"# Assessment {ass_id} tests\n## Page {page_num} / {tot_pages}\n"
    for test in tests_res:
        test = dict(test)
        assets = test.pop("Assets", {})
        scenarios = test.pop("Scenarios", {})
        test_name = test.get("Name")
        hr += tableToMarkdown(
            f"Test - {test_name}",
            test,
            headers=["Id", "Name", "Created", "Modified", "Runnable", "LastResult"],
            headerTransform=pascalToSpace,
        )
        hr += tableToMarkdown(f"Assets ({test_name})", assets)
        hr += tableToMarkdown(f"Scenarios ({test_name})", scenarios)
    if not hr:
        hr = "Found no tests"
    return hr


def list_tests_by_assessment(params):
    return http_request("GET", "/v1/tests", params=params)


def list_tests_by_assessment_command():
    """Implements attackiq-list-tests-by-assessment command"""
    page, page_size = get_page_number_and_page_size(demisto.args())
    ass_id = demisto.getArg("assessment_id")
    params = {"project": ass_id, "page_size": page_size, "page": page}
    raw_res = list_tests_by_assessment(params)
    test_cnt = raw_res.get("count")
    if test_cnt == 0:
        return_outputs("No results were found", {})
    else:
        tests_res = build_transformed_dict(raw_res.get("results"), TESTS_TRANS)
        total_pages = math.ceil(test_cnt / page_size)
        remaining_pages = total_pages - page
        if remaining_pages < 0:
            remaining_pages = 0
        context = {
            "AttackIQTest(val.Id === obj.Id)": tests_res,
            "AttackIQTest(val.Count).Count": test_cnt,
            "AttackIQTest(val.RemainingPages).RemainingPages": remaining_pages,
        }
        hr = build_tests_hr(tests_res, ass_id, page, total_pages)
        return_outputs(hr, context, raw_res)


def run_all_tests_in_assessment_command():
    """Implements attackiq-run-all-tests-in-assessment command"""
    args = demisto.args()
    ass_id = args.get("assessment_id")
    on_demand_only = args.get("on_demand_only")
    try:
        params = {"on_demand_only": on_demand_only == "True"}
        raw_res = http_request("POST", f"/v1/assessments/{ass_id}/run_all_tests", params=params)
        hr = raw_res.get("message", f"Request to run all tests for assessment {ass_id} was sent successfully.")
        demisto.results(hr)
    except HTTPError as e:
        return_error(create_invalid_id_err_msg(str(e), ["403"]))


@logger
def list_templates_command():
    """
    Returns:
        A list of all assesment templates.
    """

    res = http_request("GET", "/v1/project_template_types")
    templates = []
    for template_group in res.get("results", []):
        for template in template_group.get("project_templates", []):
            template_dict = {
                "ID": template.get("id"),
                "Name": template.get("template_name"),
                "Description": template.get("template_description"),
                "ProjectName": template.get("project_name"),
                "ProjectDescription": template.get("project_description"),
                "Hidden": template.get("hidden"),
            }
            templates.append(template_dict)
    ec = {"AttackIQ.Template(val.ID && val.ID === obj.ID)": templates}
    hr = tableToMarkdown("Templates:", templates, ["ID", "Name", "Description", "ProjectName", "ProjectDescription"])
    return_outputs(hr, ec, res)


@logger
def list_assets_command():
    """

    Returns:
        A list of all configured assets.
    """
    res = http_request("GET", "/v1/assets")
    assets = []
    for asset in res.get("results", []):
        asset_dict = {
            "ID": asset.get("id", ""),
            "Description": asset.get("description", ""),
            "IPv4": asset.get("ipv4_address", ""),
            "IPv6": asset.get("ipv6_address", ""),
            "MacAddress": asset.get("mac_address", ""),
            "ProcessorArch": asset.get("processor_arch", ""),
            "ProductName": asset.get("product_name", ""),
            "Hostname": asset.get("hostname", ""),
            "Domain": asset.get("domain_name", ""),
            "User": asset.get("user", ""),
            "Status": asset.get("status", ""),
        }
        groups = []
        for group in asset.get("asset_groups", []):
            temp_group = {"ID": group.get("id"), "Name": group.get("name")}
            groups.append(temp_group)
        asset_dict["Groups"] = groups
        assets.append(asset_dict)
    ec = {"AttackIQ.Asset(val.ID && val.ID === obj.ID)": assets}
    hr = tableToMarkdown("Assets:", assets, ["ID", "Hostname", "IPv4", "MacAddress", "Domain", "Description", "User", "Status"])

    return_outputs(hr, ec, res)


@logger
def create_assessment_command():
    """
    name - The name of the assesment to create.

    Returns:

    """
    body = {"project_name": demisto.args().get("name"), "template": demisto.args().get("template_id")}

    try:
        res = http_request("POST", "/v1/assessments/project_from_template", data=json.dumps(body))
    except Exception as e:
        raise ValueError(f"Could not create an assessment. Specifically: {str(e)}")

    assessment_id = res.get("project_id")
    raw_assessments = get_assessments(assessment_id=assessment_id)
    assessments_res = build_transformed_dict(raw_assessments, ASSESSMENTS_TRANS)
    hr = tableToMarkdown(
        f"Created Assessment: {assessment_id} successfully.",
        assessments_res,
        headers=["Id", "Name", "Description", "User", "Created", "Modified"],
    )
    return_outputs(hr, {"AttackIQ.Assessment(val.Id === obj.Id)": assessments_res}, raw_assessments)


@logger
def add_assets_to_assessment():
    assessment_id = demisto.args().get("assessment_id")
    assets = demisto.args().get("assets")
    asset_groups = demisto.args().get("asset_groups")

    data = {}
    if assets:
        data["assets"] = assets
    if asset_groups:
        data["asset_groups"] = asset_groups

    if data == {}:
        raise ValueError("No asset or asset groups were specified.")
    try:
        res = http_request("POST", f"/v1/assessments/{assessment_id}/update_defaults", data=json.dumps(data))
        demisto.results(res.get("message", ""))
    except Exception as e:
        if "403" in str(e):
            raise ValueError("Could not find either the assessment or one of the assets/asset groups.")
        else:
            raise


@logger
def delete_assessment_command():
    assessment_id = demisto.args().get("assessment_id")
    try:
        http_request("DELETE", f"/v1/assessments/{assessment_id}")
        demisto.results(f"Deleted assessment {assessment_id} successfully.")
    except Exception as e:
        if "403" in str(e):
            raise ValueError(f"Could not find the assessment {assessment_id}")
        else:
            raise


def main():
    handle_proxy()
    command = demisto.command()
    LOG(f"Command being called is {command}")
    try:
        if command == "test-module":
            test_module()
        elif command == "attackiq-activate-assessment":
            activate_assessment_command()
        elif command == "attackiq-get-assessment-execution-status":
            get_assessment_execution_status_command()
        elif command == "attackiq-get-test-execution-status":
            get_test_execution_status_command()
        elif command == "attackiq-get-test-results":
            get_test_results_command()
        elif command == "attackiq-list-assessments":
            list_assessments_command()
        elif command == "attackiq-get-assessment-by-id":
            get_assessment_by_id_command()
        elif command == "attackiq-list-tests-by-assessment":
            list_tests_by_assessment_command()
        elif command == "attackiq-run-all-tests-in-assessment":
            run_all_tests_in_assessment_command()
        elif command == "attackiq-list-assessment-templates":
            list_templates_command()
        elif command == "attackiq-list-assets":
            list_assets_command()
        elif command == "attackiq-create-assessment":
            create_assessment_command()
        elif command == "attackiq-add-assets-to-assessment":
            add_assets_to_assessment()
        elif command == "attackiq-delete-assessment":
            delete_assessment_command()
        else:
            return_error(f"Command {command} is not supported.")
    except HTTPError as e:
        # e is expected to contain parsed error message
        err = f"{API_ERR_MSG}{str(e)}"
        return_error(err)
    except Exception as e:
        message = f"Unexpected error: {str(e)}, traceback: {traceback.format_exc()}"
        return_error(message)


# python2 uses __builtin__ python3 uses builtins
if __name__ in ("__builtin__", "builtins", "__main__"):
    main()