Proofpoint Threat Protection

from datetime import datetime, timedelta
import pytest
import urllib3

from CommonServerPython import date_to_timestamp
from ProofpointThreatProtection import *

urllib3.disable_warnings()

""" CONSTANTS """

TEST_SERVER_IP_BINDING = "127.0.0.1"
TEST_SERVER_TCP_PORT = 8000
TEST_SERVER_BASE_URL = f"http://{TEST_SERVER_IP_BINDING}:{TEST_SERVER_TCP_PORT}/api/v1"
TEST_AUTH_HOST = f"http://{TEST_SERVER_IP_BINDING}:{TEST_SERVER_TCP_PORT}/v1"

ACCESS_TOKEN_VALUE = "TOKEN"

GOOD_ACCESS_TOKEN = {"access_token": ACCESS_TOKEN_VALUE, "expiry_time": date_to_timestamp(datetime.now() + timedelta(hours=1))}

EXPIRED_ACCESS_TOKEN = {
    "access_token": ACCESS_TOKEN_VALUE,
    "expiry_time": date_to_timestamp(datetime.now() + timedelta(hours=-1)),
}

GOOD_ACCESS_TOKEN_FROM_API = {"access_token": ACCESS_TOKEN_VALUE, "expires_in": 3600}

MOCK_SAFEBLOCK_ADD_ENTRY = {
    "action": "add",
    "attribute": "$hfrom",
    "operator": "equal",
    "value": "test@mydomain.com",
    "comment": "comment 1",
}

MOCK_SAFEBLOCK_LIST_API_RETURN = [
    {"attribute": "$hfrom", "operator": "equal", "value": "test@mydomain.com", "comment": "comment 1"},
    {"attribute": "$from", "operator": "equal", "value": "sample@example.com", "comment": "comment B"},
]

""" HELPER FUNCTIONS """


def mock_main(mocker, command, args={}):
    rc = Client(base_url=TEST_SERVER_BASE_URL, verify=False, proxy=False)

    # Prevent any accidental real HTTP calls
    mocker.patch.object(rc, "_http_request", return_value={})
    mocker.patch.object(rc, "get_access_token", return_value=ACCESS_TOKEN_VALUE)
    mocker.patch.object(rc, "get_args", return_value=args)

    mocker.patch.object(rc, "get_safelist", return_value={"entries": MOCK_SAFEBLOCK_LIST_API_RETURN})
    mocker.patch.object(rc, "get_blocklist", return_value={"entries": MOCK_SAFEBLOCK_LIST_API_RETURN})
    mocker.patch.object(rc, "safelist_add_delete", return_value={"status": "success"})
    mocker.patch.object(rc, "blocklist_add_delete", return_value={"status": "success"})

    return COMMANDS[command](rc, "CLUSTERID1")


""" TEST FUNCTIONS """


def test_non_expired_access_token_present(mocker):
    c = Client(base_url=TEST_SERVER_BASE_URL, verify=False)
    mocker.patch.object(c, "get_shared_integration_context", return_value=GOOD_ACCESS_TOKEN)
    obtained_token = c.get_access_token("CLID1", "CLSECRET1")
    assert obtained_token == ACCESS_TOKEN_VALUE


def test_expired_access_token(mocker):
    c = Client(base_url=TEST_SERVER_BASE_URL, verify=False)

    fixed_now_dt = datetime(2025, 6, 8, 10, 0, 0)
    expired_time = date_to_timestamp(fixed_now_dt - timedelta(hours=1), date_format="%Y-%m-%dT%H:%M:%S")
    expired_token_context = {"access_token": ACCESS_TOKEN_VALUE, "expiry_time": expired_time}

    mocker.patch.object(c, "get_shared_integration_context", return_value=expired_token_context)
    mocker.patch.object(c, "_http_request", return_value=GOOD_ACCESS_TOKEN_FROM_API)

    mocker.patch("ProofpointThreatProtection.datetime", wraps=datetime)
    mocker.patch("ProofpointThreatProtection.datetime.now", return_value=fixed_now_dt)

    expected_expiry_time = date_to_timestamp(fixed_now_dt, date_format="%Y-%m-%dT%H:%M:%S")
    expected_expiry_time += GOOD_ACCESS_TOKEN_FROM_API["expires_in"] * 1000 - 10
    expected_context_to_set = {"access_token": ACCESS_TOKEN_VALUE, "expiry_time": expected_expiry_time}

    mock_set_context = mocker.patch("ProofpointThreatProtection.set_integration_context")
    mocker.patch("ProofpointThreatProtection.get_integration_context", return_value=expected_context_to_set)

    obtained_token = c.get_access_token("CLID1", "CLSECRET1")

    mock_set_context.assert_called_once_with(expected_context_to_set)
    assert obtained_token == ACCESS_TOKEN_VALUE


def test_non_existent_access_token(mocker):
    c = Client(base_url=TEST_SERVER_BASE_URL, verify=False)
    mocker.patch.object(c, "get_shared_integration_context", return_value={})
    mocker.patch.object(c, "_http_request", return_value=GOOD_ACCESS_TOKEN_FROM_API)

    # Fixed datetime object
    fixed_now_dt = datetime(2025, 6, 8, 10, 0, 0)

    # Mock datetime.datetime.now() correctly
    mocked_datetime = mocker.patch("ProofpointThreatProtection.datetime")
    mocked_datetime.now.return_value = fixed_now_dt
    mocked_datetime.side_effect = lambda *args, **kw: datetime(*args, **kw)  # still allow datetime(…) calls

    expected_expiry_time = date_to_timestamp(fixed_now_dt, date_format="%Y-%m-%dT%H:%M:%S")
    expected_expiry_time += GOOD_ACCESS_TOKEN_FROM_API.get("expires_in", 0) * 1000 - 10
    expected_context_to_set = {"access_token": ACCESS_TOKEN_VALUE, "expiry_time": expected_expiry_time}

    mock_set_context = mocker.patch.object(c, "set_shared_integration_context")
    mocker.patch("CommonServerPython.get_integration_context", return_value=expected_context_to_set)

    obtained_token = c.get_access_token("CLID1", "CLSECRET1")

    mock_set_context.assert_called_once_with(expected_context_to_set)
    assert obtained_token == ACCESS_TOKEN_VALUE


def test_bad_get_access_token_request(mocker):
    c = Client(base_url=TEST_SERVER_BASE_URL, verify=False)
    mocker.patch.object(c, "get_shared_integration_context", return_value={})
    mocker.patch.object(c, "get_auth_host", return_value=TEST_AUTH_HOST)
    mocker.patch.object(
        c,
        "_http_request",
        side_effect=Exception("Error occurred while creating an access token. Please check the instance configuration."),
    )
    mock_set_context = mocker.patch.object(c, "set_shared_integration_context")

    with pytest.raises(Exception) as error_info:
        c.get_access_token("CLID1", "CLSECRET1")

    assert str(error_info.value).startswith("Error occurred while creating an access token")
    mock_set_context.assert_not_called()


def test_list_safelist(mocker):
    return_obj = mock_main(mocker, "proofpoint-tp-safelist-list").outputs["Safelist"]
    assert return_obj == MOCK_SAFEBLOCK_LIST_API_RETURN


def test_list_blocklist(mocker):
    return_obj = mock_main(mocker, "proofpoint-tp-blocklist-list").outputs["Blocklist"]
    assert return_obj == MOCK_SAFEBLOCK_LIST_API_RETURN


def test_add_to_safelist(mocker):
    return_obj = mock_main(mocker, "proofpoint-tp-safelist-add-entry", args=MOCK_SAFEBLOCK_ADD_ENTRY)
    assert return_obj.outputs["Safelist Entry Added"] == "Success"


def test_add_to_blocklist(mocker):
    return_obj = mock_main(mocker, "proofpoint-tp-blocklist-add-entry", args=MOCK_SAFEBLOCK_ADD_ENTRY)
    assert return_obj.outputs["Blocklist Entry Added"] == "Success"


def test_delete_from_safelist(mocker):
    return_obj = mock_main(mocker, "proofpoint-tp-safelist-delete-entry", args=MOCK_SAFEBLOCK_ADD_ENTRY)
    assert return_obj.outputs["Safelist Entry Deleted"] == "Success"


def test_delete_from_blocklist(mocker):
    return_obj = mock_main(mocker, "proofpoint-tp-blocklist-delete-entry", args=MOCK_SAFEBLOCK_ADD_ENTRY)
    assert return_obj.outputs["Blocklist Entry Deleted"] == "Success"


def test_parse_params(mocker):
    client_id, client_secret, base_url, cluster_id, verify_certificate, proxy = parse_params(
        {
            "credentials": {"username": "client_id", "password": "client_secret"},
            "url": "base_url",
            "cluster_id": "cluster_id",
            "verify_certificate": True,
            "proxy": True,
        }
    )

    assert client_id == "client_id"
    assert client_secret == "client_secret"
    assert base_url == "base_url"
    assert cluster_id == "cluster_id"
    assert verify_certificate is True
    assert proxy is True