Stamus

import json


def util_load_json(path):
    with open(path, encoding="utf-8") as f:
        return json.loads(f.read())


# TODO: ADD HERE unit tests for every command
def test_fetch_incidents(requests_mock):
    """
    Tests the fetch-incidents command function.

        Given:
            - requests_mock instance to generate the appropriate get_alert API response,
              loaded from a local JSON file.

        When:
            - Running the 'fetch_incidents' command.

        Then:
            - Checks the output of the command function with the expected output.
    """
    from Stamus import Client, fetch_incidents

    url = "https://stamus-test.com/rest/appliances/threat_history_incident/?timestamp=1689340848&page_size=200"
    mock_response = util_load_json("test_data/fetch_incidents.json")
    requests_mock.get(url, json=mock_response)

    client = Client(base_url="https://stamus-test.com", verify=False, headers={"Authentication": "Bearer some_api_key"})

    _, new_incidents = fetch_incidents(client=client, timestamp=1689340848)

    assert new_incidents == [
        {
            "name": "10.11.13.101_incident_0",
            "dbotMirrorId": "283115",
            "rawJSON": json.dumps(mock_response["results"][0]),
            "details": "description of my custom threat",
            "occurred": "2023-07-14T15:20:48.231617+02:00",
            "type": "ip",
        }
    ]


def test_fetch_by_ioc(requests_mock):
    """
    Given:
        - key/value to check.

    When:
        - Running the 'fetch_by_ioc'.

    Then:
        - Checks that results are the same than the ones built in tests.
    """
    from Stamus import Client, CommandResults, fetch_by_ioc, get_command_results, tableToMarkdown

    filter = ("src_ip", "10.7.5.101")
    url = f"https://stamus-test.com/rest/rules/es/events_tail/?qfilter={filter[0]}%3A{filter[1]}"
    mock_response = util_load_json("test_data/fetch_by_ioc.json")
    requests_mock.get(url, json=mock_response)

    client = Client(base_url="https://stamus-test.com", verify=False, headers={"Authentication": "Bearer some_api_key"})

    events = fetch_by_ioc(client, args={"indicator_key": filter[0], "indicator_value": filter[1]})
    results = mock_response.get("results", [])
    table = tableToMarkdown("IOC Matches", results, headers=["timestamp", "src_ip", "dest_ip", "event_type"])

    cmd_results = get_command_results(results, table, "IOC")

    for key in cmd_results.__dict__:
        assert cmd_results.__dict__.get(key) == events.__dict__.get(key)
    assert events.raw_response == mock_response["results"]
    assert isinstance(events, CommandResults)
    assert events.outputs_prefix == "StamusIntegration.IOC"


def test_fetch_events(requests_mock):
    """
    Given:
        - An id address to check.

    When:
        - Running the 'fetch_events'.

    Then:
        - Checks that results are the same than the ones built in tests.
    """
    from Stamus import Client, CommandResults, fetch_events, get_command_results, tableToMarkdown

    id_to_check = 3
    url = f"https://stamus-test.com/rest/appliances/threat_history_incident/{id_to_check}/get_events"
    mock_response = util_load_json("test_data/fetch_events.json")
    requests_mock.get(url, json=mock_response)

    client = Client(base_url="https://stamus-test.com", verify=False, headers={"Authentication": "Bearer some_api_key"})

    events = fetch_events(client, args={"id": id_to_check})
    results = mock_response.get("results", [])
    for result in results:
        result["method"] = result.get("alert", {}).get("signature", "algorithmic detection")
        result["info"] = ""
        if result.get("hostname_info"):
            result["info"] = f"Hostname: {result.get('hostname_info', {}).get('host', 'unknown')}"
        result["asset"] = result.get("stamus", {}).get("asset", "unknown")
        result["offender"] = result.get("stamus", {}).get("source", "unknown")
        result["killchain"] = result.get("stamus", {}).get("kill_chain", "unknown")
    headers = ["timestamp", "asset", "offender", "killchain", "method", "info", "src_ip", "dest_ip", "app_proto"]
    table = tableToMarkdown("Individual Events List", results, headers=headers)

    cmd_results = get_command_results(results, table, "RelatedEvents")

    for key in cmd_results.__dict__:
        assert cmd_results.__dict__.get(key) == events.__dict__.get(key)
    assert events.raw_response == mock_response["results"]
    assert isinstance(events, CommandResults)
    assert events.outputs_prefix == "StamusIntegration.RelatedEvents"


def test_fetch_host_id(requests_mock):
    """
    Given:
        - An ip address to check.

    When:
        - Running the 'fetch_host_id'.

    Then:
        - Checks that results are the same than the ones built in tests.
    """
    from Stamus import Client, CommandResults, fetch_host_id, get_command_results, linearize_host_id, tableToMarkdown

    ip = "217.116.0.227"
    url = f"https://stamus-test.com/rest/appliances/host_id/{ip}"
    mock_response = util_load_json("test_data/fetch_host_id.json")
    requests_mock.get(url, json=mock_response)

    client = Client(base_url="https://stamus-test.com", verify=False, headers={"Authentication": "Bearer some_api_key"})

    events = fetch_host_id(client, args={"ip": ip})
    results = mock_response
    headers = ["timestamp", "ip", "type", "value"]
    host_info = linearize_host_id(results)
    table = tableToMarkdown("Host Insight", host_info, headers=headers)

    cmd_results = get_command_results(results, table, "HostInsights")

    for key in cmd_results.__dict__:
        assert cmd_results.__dict__.get(key) == events.__dict__.get(key)
    assert cmd_results.readable_output == table
    assert cmd_results.outputs == results
    assert events.raw_response == results
    assert isinstance(events, CommandResults)
    assert events.outputs_prefix == "StamusIntegration.HostInsights"