"""Armis Integration for Cortex XSOAR - Unit Tests file This file contains the Pytest Tests for the Armis Integration """ import json import pytest import time import CommonServerPython import demistomock as demisto def test_untag_device_success(requests_mock): from Armis import Client, untag_device_command mock_token = {"data": {"access_token": "example", "expiration_utc": time.ctime(time.time() + 10000)}} requests_mock.post("https://test.com/api/v1/access_token/", json=mock_token) requests_mock.delete("https://test.com/api/v1/devices/1/tags/", json={}) client = Client("secret-example", "https://test.com/api/v1", verify=False, proxy=False) result = untag_device_command(client, {"device_id": "1", "tags": "test-tag"}) assert result == "Successfully Untagged device: 1 with tags: ['test-tag']" def test_untag_device_failure(requests_mock): from Armis import Client, untag_device_command mock_token = {"data": {"access_token": "example", "expiration_utc": time.ctime(time.time() + 10000)}} requests_mock.post("https://test.com/api/v1/access_token/", json=mock_token) requests_mock.delete("https://test.com/api/v1/devices/1/tags/", json={}, status_code=400) client = Client("secret-example", "https://test.com/api/v1", verify=False, proxy=False) with pytest.raises(CommonServerPython.DemistoException): untag_device_command(client, {"device_id": "1", "tags": "test-tag"}) def test_tag_device(requests_mock): from Armis import Client, tag_device_command mock_token = {"data": {"access_token": "example", "expiration_utc": time.ctime(time.time() + 10000)}} requests_mock.post("https://test.com/api/v1/access_token/", json=mock_token) requests_mock.post("https://test.com/api/v1/devices/1/tags/", json={}) client = Client("secret-example", "https://test.com/api/v1", verify=False, proxy=False) result = tag_device_command(client, {"device_id": "1", "tags": "test-tag"}) assert result == "Successfully Tagged device: 1 with tags: ['test-tag']" def test_update_alert_status(requests_mock): from Armis import Client, update_alert_status_command mock_token = {"data": {"access_token": "example", "expiration_utc": time.ctime(time.time() + 10000)}} requests_mock.post("https://test.com/api/v1/access_token/", json=mock_token) requests_mock.patch("https://test.com/api/v1/alerts/1/", json={}) client = Client("secret-example", "https://test.com/api/v1", verify=False, proxy=False) args = {"alert_id": "1", "status": "UNHANDLED"} assert update_alert_status_command(client, args) == "Successfully Updated Alert: 1 to status: UNHANDLED" def test_search_alerts(requests_mock): from Armis import Client, search_alerts_command mock_token = {"data": {"access_token": "example", "expiration_utc": time.ctime(time.time() + 10000)}} requests_mock.post("https://test.com/api/v1/access_token/", json=mock_token) url = "https://test.com/api/v1/search/?aql=" url += "+".join( [ "in%3Aalerts", "timeFrame%3A%223+days%22", "riskLevel%3AHigh%2CMedium", "status%3AUNHANDLED%2CRESOLVED", "type%3A%22Policy+Violation%22", "alertId%3A%281%29", ] ) mock_results = {"data": {"results": []}} requests_mock.get(url, json=mock_results) client = Client("secret-example", "https://test.com/api/v1", verify=False, proxy=False) args = { "severity": "High,Medium", "status": "UNHANDLED,RESOLVED", "alert_type": "Policy Violation", "alert_id": "1", "max_results": "20", "time_frame": "3 days", } response = search_alerts_command(client, args) assert response == "No results found" example_alerts = [ { "activityIds": [19625045, 19625223, 19625984, 19626169, 19626680, 19626818, 19628162, 19628359], "activityUUIDs": [ "1-uS23YBAAAC-vCTQOhA", "7eut23YBAAAC-vCTkOhB", "Oes13HYBAAAC-vCTcel0", "T-tU3HYBAAAC-vCTyunu", "mevb3HYBAAAC-vCT9-nn", "uev33HYBAAAC-vCTa-mg", "P-u33XYBAAAC-vCTlOpq", "SevT3XYBAAAC-vCTA-o_", ], "alertId": 1, "connectionIds": [845993, 846061, 846157, 846308], "description": "Smart TV started connection to Corporate Network", "deviceIds": [165722, 532], "severity": "Medium", "status": "Unhandled", "time": "2021-01-07T06:39:13.320893+00:00", "title": "Smart TV connected to Corporate network", "type": "System Policy Violation", } ] mock_results["data"]["results"] = example_alerts requests_mock.get(url, json=mock_results) response = search_alerts_command(client, args) assert response.outputs == example_alerts def test_search_alerts_by_aql(requests_mock): from Armis import Client, search_alerts_by_aql_command mock_token = {"data": {"access_token": "example", "expiration_utc": time.ctime(time.time() + 10000)}} requests_mock.post("https://test.com/api/v1/access_token/", json=mock_token) url = "https://test.com/api/v1/search/?aql=" url += "+".join( [ "in%3Aalerts", "timeFrame%3A%223+days%22", "riskLevel%3AHigh%2CMedium", "status%3AUNHANDLED%2CRESOLVED", "type%3A%22Policy+Violation%22", ] ) mock_results = {"data": {"results": []}} requests_mock.get(url, json=mock_results) client = Client("secret-example", "https://test.com/api/v1", verify=False, proxy=False) args = {"aql_string": 'timeFrame:"3 days" riskLevel:High,Medium status:UNHANDLED,RESOLVED type:"Policy Violation"'} response = search_alerts_by_aql_command(client, args) assert response == "No alerts found" example_alerts = [ { "activityIds": [19625045, 19625223, 19625984, 19626169, 19626680, 19626818, 19628162, 19628359], "activityUUIDs": [ "1-uS23YBAAAC-vCTQOhA", "7eut23YBAAAC-vCTkOhB", "Oes13HYBAAAC-vCTcel0", "T-tU3HYBAAAC-vCTyunu", "mevb3HYBAAAC-vCT9-nn", "uev33HYBAAAC-vCTa-mg", "P-u33XYBAAAC-vCTlOpq", "SevT3XYBAAAC-vCTA-o_", ], "alertId": 1, "connectionIds": [845993, 846061, 846157, 846308], "description": "Smart TV started connection to Corporate Network", "deviceIds": [165722, 532], "severity": "Medium", "status": "Unhandled", "time": "2021-01-07T06:39:13.320893+00:00", "title": "Smart TV connected to Corporate network", "type": "System Policy Violation", } ] mock_results["data"]["results"] = example_alerts requests_mock.get(url, json=mock_results) response = search_alerts_by_aql_command(client, args) assert response.outputs == example_alerts def test_search_devices(requests_mock): from Armis import Client, search_devices_command mock_token = {"data": {"access_token": "example", "expiration_utc": time.ctime(time.time() + 10000)}} requests_mock.post("https://test.com/api/v1/access_token/", json=mock_token) url = "https://test.com/api/v1/search/?aql=in%3Adevices+timeFrame%3A%223+days%22+deviceId%3A%281%29" mock_results = {"data": {"results": []}} requests_mock.get(url, json=mock_results) client = Client("secret-example", "https://test.com/api/v1", verify=False, proxy=False) args = {"device_id": "1", "time_frame": "3 days"} response = search_devices_command(client, args) assert response == "No devices found" example_alerts = [ { "accessSwitch": None, "category": "Network Equipment", "dataSources": [ { "firstSeen": "2021-01-15T03:26:56+00:00", "lastSeen": "2021-01-16T18:16:32+00:00", "name": "Meraki", "types": ["WLC"], } ], "firstSeen": "2021-01-15T03:26:56+00:00", "id": 1, "ipAddress": None, "ipv6": None, "lastSeen": "2021-01-16T18:16:32+00:00", "macAddress": "f8:ca:59:53:91:ce", "manufacturer": "NetComm Wireless", "model": "NetComm device", "name": "Aussie Broadband 0079", "operatingSystem": None, "operatingSystemVersion": None, "riskLevel": 5, "sensor": {"name": "win-wap-tom-Upstairs", "type": "Access Point"}, "site": {"location": "51 Longview Court, Thomastown Vic 3074", "name": "Winslow Workshop - Thomastown"}, "tags": ["Access Point", "Off Network", "SSID=Aussie Broadband 0079"], "type": "Access Point Interface", "user": "", "visibility": "Full", } ] mock_results["data"]["results"] = example_alerts requests_mock.get(url, json=mock_results) response = search_devices_command(client, args) assert response.outputs == example_alerts def test_search_devices_by_aql(requests_mock): from Armis import Client, search_devices_by_aql_command mock_token = {"data": {"access_token": "example", "expiration_utc": time.ctime(time.time() + 10000)}} requests_mock.post("https://test.com/api/v1/access_token/", json=mock_token) url = "https://test.com/api/v1/search/?aql=in%3Adevices+timeFrame%3A%223+days%22+deviceId%3A%281%29" mock_results = {"data": {"results": []}} requests_mock.get(url, json=mock_results) client = Client("secret-example", "https://test.com/api/v1", verify=False, proxy=False) args = {"aql_string": 'timeFrame:"3 days" deviceId:(1)'} response = search_devices_by_aql_command(client, args) assert response == "No devices found" example_alerts = [ { "accessSwitch": None, "category": "Network Equipment", "dataSources": [ { "firstSeen": "2021-01-15T03:26:56+00:00", "lastSeen": "2021-01-16T18:16:32+00:00", "name": "Meraki", "types": ["WLC"], } ], "firstSeen": "2021-01-15T03:26:56+00:00", "id": 1, "ipAddress": None, "ipv6": None, "lastSeen": "2021-01-16T18:16:32+00:00", "macAddress": "f8:ca:59:53:91:ce", "manufacturer": "NetComm Wireless", "model": "NetComm device", "name": "Aussie Broadband 0079", "operatingSystem": None, "operatingSystemVersion": None, "riskLevel": 5, "sensor": {"name": "win-wap-tom-Upstairs", "type": "Access Point"}, "site": {"location": "51 Longview Court, Thomastown Vic 3074", "name": "Winslow Workshop - Thomastown"}, "tags": ["Access Point", "Off Network", "SSID=Aussie Broadband 0079"], "type": "Access Point Interface", "user": "", "visibility": "Full", } ] mock_results["data"]["results"] = example_alerts requests_mock.get(url, json=mock_results) response = search_devices_by_aql_command(client, args) assert response.outputs == example_alerts def test_fetch_incidents_no_duplicates(mocker): """ Given: - 'client': Armis client. - 'last_run': Last run parameters. When: - Performing two consecutive calls to fetch incidents Then: - Ensure incident that was already fetched is not fetched again. """ from Armis import Client, fetch_incidents client = Client("secret-example", "https://test.com/api/v1", verify=False, proxy=False) last_fetch = "2021-03-09T01:00:00.000001+00:00" armis_incident = {"time": "2021-03-09T01:00:00.000001+00:00", "type": "System Policy Violation"} response = {"results": [armis_incident], "next": "more data"} mocker.patch.object(client, "search_alerts", return_value=response) next_run, incidents = fetch_incidents(client, {"last_fetch": last_fetch}, None, "Low", [], [], "", 1) assert next_run["last_fetch"] == last_fetch assert incidents[0]["rawJSON"] == json.dumps(armis_incident) _, incidents = fetch_incidents(client, next_run, None, "Low", [], [], "", 1) assert not incidents class MockClient: def __init__(self, secret: str, base_url: str, verify: bool, proxy): pass def test_url_parameter(mocker): """ Given: - Instance parameters with a base URL without `api/v1` prefix. When: - Running the main function and configured the client class. Then: - Ensure that hte base URL in the client class is with teh `api/v1` prefix. """ from Armis import main mocker.patch.object(demisto, "params", return_value={"url": "test.com"}) mock_client = mocker.patch("Armis.Client", side_effect=MockClient) main() assert mock_client.call_args.kwargs["base_url"] == "test.com/api/v1/" def test_get_api_token_when_found_in_integration_context(mocker): """Test cases for scenario when there is api_token and expiration_time in integration context.""" from Armis import Client test_integration_context = {"token": "1234567890", "token_expiration": time.ctime(time.time() + 10000)} mocker.patch.object(demisto, "getIntegrationContext", return_value=test_integration_context) client = Client("secret-example", "https://test.com/api/v1", verify=False, proxy=False) api_token = client._get_token() assert api_token == test_integration_context["token"] def test_get_api_token_when_expired_token_found_in_integration_context(mocker, requests_mock): """Test cases for scenario when there is an expired api_token in integration context.""" from Armis import Client mock_token = {"data": {"access_token": "example", "expiration_utc": time.ctime(time.time() + 10000)}} requests_mock.post("https://test.com/api/v1/access_token/", json=mock_token) client = Client("secret-example", "https://test.com/api/v1", verify=False, proxy=False) api_token = client._get_token() assert api_token == mock_token["data"]["access_token"] def test_retry_for_401_error(mocker, requests_mock): from Armis import Client, search_alerts_by_aql_command test_integration_context = {"token": "invalid_token", "token_expiration": time.ctime(time.time() - 10000)} mocker.patch.object(demisto, "getIntegrationContext", return_value=test_integration_context) url = "https://test.com/api/v1/search/?aql=" url += "+".join( [ "in%3Aalerts", "timeFrame%3A%223+days%22", "riskLevel%3AHigh%2CMedium", "status%3AUNHANDLED%2CRESOLVED", "type%3A%22Policy+Violation%22", ] ) mock_results = {"message": "Invalid access token.", "success": False} mock_token = {"data": {"access_token": "example", "expiration_utc": time.ctime(time.time() + 10000)}} requests_mock.post("https://test.com/api/v1/access_token/", json=mock_token) example_alerts = [ { "accessSwitch": None, "category": "Dummy Category", "dataSources": [ { "firstSeen": "2025-01-01T00:00:00+00:00", "lastSeen": "2025-01-02T00:00:00+00:00", "name": "Dummy Source", "types": ["Dummy Type"], } ], "firstSeen": "2025-01-01T00:00:00+00:00", "id": 100, "ipAddress": "0.0.0.1", "ipv6": "0000:0000:0000:0000:0000:0000:0000:0001", "lastSeen": "2025-01-02T00:00:00+00:00", "macAddress": "00:00:00:00:00:01", "manufacturer": "Dummy Manufacturer", "model": "Dummy Model", "name": "Dummy Device", "operatingSystem": "Dummy OS", "operatingSystemVersion": "1.0", "riskLevel": 1, "sensor": {"name": "Dummy Sensor", "type": "Dummy Sensor Type"}, "site": {"location": "Dummy Location", "name": "Dummy Site"}, "tags": ["Dummy Tag 1", "Dummy Tag 2", "Dummy Tag 3"], "type": "Dummy Type", "user": "Dummy User", "visibility": "Dummy Visibility", } ] new_mock_results = {"data": {"results": example_alerts}} requests_mock.register_uri( "GET", url, [{"status_code": 401, "json": mock_results}, {"status_code": 200, "json": new_mock_results}] ) client = Client("secret-example", "https://test.com/api/v1", verify=False, proxy=False) args = {"aql_string": 'timeFrame:"3 days" riskLevel:High,Medium status:UNHANDLED,RESOLVED type:"Policy Violation"'} response = search_alerts_by_aql_command(client, args) assert response.outputs == example_alerts def test_test_module_when_is_fetch_is_true(mocker): """ Given: - 'client': Armis client. - 'params': A dictionary containing the parameters provided by the user. When: - Performing calls to test_module Then: - Ensure test_module returns 'ok' """ from Armis import Client, test_module as armis_test_module params = { "isFetch": True, "min_severity": "Low", "alert_type": [], "alert_status": [], "free_fetch_string": "", "first_fetch": "3 days", "max_fetch": 10, } mocker.patch.object(demisto, "params", return_value=params) client = Client("secret-example", "https://test.com/api/v1", verify=False, proxy=False) armis_incident = {"time": "2025-03-09T01:00:00.000001+00:00", "type": "test_type"} response = {"results": [armis_incident], "next": "more data"} mocker.patch.object(client, "search_alerts", return_value=response) assert armis_test_module(client, params) == "ok"