From 91fee60530ac38688330b5d0bc465901fe6e7b18 Mon Sep 17 00:00:00 2001 From: mohsinenar Date: Thu, 28 Mar 2024 14:10:33 +0000 Subject: [PATCH 1/5] Name scan folder by scan id. --- agent/nebula_agent.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/agent/nebula_agent.py b/agent/nebula_agent.py index a85690f..88b7afe 100644 --- a/agent/nebula_agent.py +++ b/agent/nebula_agent.py @@ -47,7 +47,7 @@ def __init__( ) self._output_folder = ( - f"/output/messages_{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}" + f"/output/scan_{self.universe}_messages" ) os.makedirs(self._output_folder, exist_ok=True) From 200f664bdbbbcac37b47f02d09008ca4d9d14956 Mon Sep 17 00:00:00 2001 From: mohsinenar Date: Thu, 28 Mar 2024 14:29:30 +0000 Subject: [PATCH 2/5] Fix unitest. --- tests/nebula_agent_test.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/tests/nebula_agent_test.py b/tests/nebula_agent_test.py index 6074a5f..61f8eb4 100644 --- a/tests/nebula_agent_test.py +++ b/tests/nebula_agent_test.py @@ -13,7 +13,7 @@ from pyfakefs import fake_filesystem_unittest from agent import nebula_agent - +import os def testAgentNebula_whenUnsupportedFileType_raisesValueError() -> None: """Test that NebulaAgent raises ValueError when file type is not supported.""" @@ -37,13 +37,13 @@ def testAgentNebula_whenUnsupportedFileType_raisesValueError() -> None: nebula_agent.NebulaAgent(definition, settings) -@freeze_time("2024-03-05 12:00:00") def testAgentNebula_whenFileTypeIsJson_persistMessage( agent_definition: agent_definitions.AgentDefinition, agent_settings: runtime_definitions.AgentSettings, link_message: msg.Message, ) -> None: """Test that NebulaAgent persists message to json file.""" + os.environ["UNIVERSE"] = "43" with fake_filesystem_unittest.Patcher(): expected_output = json.dumps( {"url": "https://ostorlab.co", "method": b"GET"}, @@ -53,23 +53,23 @@ def testAgentNebula_whenFileTypeIsJson_persistMessage( nebula_test_agent.process(link_message) - assert os.path.exists("/output/messages_2024-03-05_12-00-00") - assert len(os.listdir("/output/messages_2024-03-05_12-00-00")) == 1 + assert os.path.exists("/output/scan_43_messages") + assert len(os.listdir("/output/scan_43_messages")) == 1 with open( - "/output/messages_2024-03-05_12-00-00/v3.asset.link_messages.json" + "/output/scan_43_messages/v3.asset.link_messages.json" ) as file: assert sorted(json.load(file).items()) == sorted( json.loads(expected_output).items() ) -@freeze_time("2023-03-05 12:00:00") def testAgentNebula_whenFileTypeIsJson_persistMultipleLinkMessages( agent_definition: agent_definitions.AgentDefinition, agent_settings: runtime_definitions.AgentSettings, multiple_link_messages: list[msg.Message], ) -> None: """Test that NebulaAgent persists multiple link messages to json file.""" + os.environ["UNIVERSE"] = "43" with fake_filesystem_unittest.Patcher(): expected_output = [ json.dumps( @@ -83,7 +83,7 @@ def testAgentNebula_whenFileTypeIsJson_persistMultipleLinkMessages( for message in multiple_link_messages: nebula_test_agent.process(message) - file_path = "/output/messages_2023-03-05_12-00-00" + file_path = "/output/scan_43_messages" assert os.path.exists(file_path) assert len(os.listdir(file_path)) == 1 with open(f"{file_path}/v3.asset.link_messages.json", "r") as file: @@ -93,13 +93,13 @@ def testAgentNebula_whenFileTypeIsJson_persistMultipleLinkMessages( assert line.strip() == expected_line.strip() -@freeze_time("2023-03-05 12:00:00") def testAgentNebula_whenFileTypeIsJson_persistMultipleMessages( agent_definition: agent_definitions.AgentDefinition, agent_settings: runtime_definitions.AgentSettings, multiple_messages: list[msg.Message], ) -> None: """Test that NebulaAgent persists multiple messages of different types to json files.""" + os.environ["UNIVERSE"] = "43" with fake_filesystem_unittest.Patcher(): expected_output = [ json.dumps( @@ -117,7 +117,7 @@ def testAgentNebula_whenFileTypeIsJson_persistMultipleMessages( for message in multiple_messages: nebula_test_agent.process(message) - file_path = "/output/messages_2023-03-05_12-00-00" + file_path = "/output/scan_43_messages" assert os.path.exists(file_path) assert len(os.listdir(file_path)) == 3 assert os.path.exists(f"{file_path}/v3.asset.link_messages.json") is True From 9de85a7c19de989a167900b6668d13e304aee631 Mon Sep 17 00:00:00 2001 From: mohsinenar Date: Thu, 28 Mar 2024 14:29:42 +0000 Subject: [PATCH 3/5] Apply ruff. --- agent/nebula_agent.py | 4 +--- tests/nebula_agent_test.py | 5 ++--- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/agent/nebula_agent.py b/agent/nebula_agent.py index 88b7afe..4f74cfe 100644 --- a/agent/nebula_agent.py +++ b/agent/nebula_agent.py @@ -46,9 +46,7 @@ def __init__( f"File type {self._file_type} is not supported. Supported file types are {SUPPORTED_FILE_TYPES}" ) - self._output_folder = ( - f"/output/scan_{self.universe}_messages" - ) + self._output_folder = f"/output/scan_{self.universe}_messages" os.makedirs(self._output_folder, exist_ok=True) def process(self, message: m.Message) -> None: diff --git a/tests/nebula_agent_test.py b/tests/nebula_agent_test.py index 61f8eb4..b47e085 100644 --- a/tests/nebula_agent_test.py +++ b/tests/nebula_agent_test.py @@ -15,6 +15,7 @@ from agent import nebula_agent import os + def testAgentNebula_whenUnsupportedFileType_raisesValueError() -> None: """Test that NebulaAgent raises ValueError when file type is not supported.""" with pytest.raises(ValueError): @@ -55,9 +56,7 @@ def testAgentNebula_whenFileTypeIsJson_persistMessage( assert os.path.exists("/output/scan_43_messages") assert len(os.listdir("/output/scan_43_messages")) == 1 - with open( - "/output/scan_43_messages/v3.asset.link_messages.json" - ) as file: + with open("/output/scan_43_messages/v3.asset.link_messages.json") as file: assert sorted(json.load(file).items()) == sorted( json.loads(expected_output).items() ) From a28d6a2da5debe73e645b6c57ca3868e9c4b53ce Mon Sep 17 00:00:00 2001 From: mohsinenar Date: Thu, 28 Mar 2024 14:32:55 +0000 Subject: [PATCH 4/5] Apply ruff. --- agent/nebula_agent.py | 7 +++---- tests/nebula_agent_test.py | 20 +++++++++----------- 2 files changed, 12 insertions(+), 15 deletions(-) diff --git a/agent/nebula_agent.py b/agent/nebula_agent.py index 4f74cfe..648d856 100644 --- a/agent/nebula_agent.py +++ b/agent/nebula_agent.py @@ -4,7 +4,6 @@ import json import logging import os -from datetime import datetime from typing import Any from ostorlab.agent import agent, definitions as agent_definitions @@ -35,9 +34,9 @@ class NebulaAgent(agent.Agent): """Agent responsible for persisting all types of messages to file type specified in the agent definition.""" def __init__( - self, - agent_definition: agent_definitions.AgentDefinition, - agent_settings: runtime_definitions.AgentSettings, + self, + agent_definition: agent_definitions.AgentDefinition, + agent_settings: runtime_definitions.AgentSettings, ) -> None: super().__init__(agent_definition, agent_settings) self._file_type = self.args.get("file_type", "json") diff --git a/tests/nebula_agent_test.py b/tests/nebula_agent_test.py index b47e085..6f1c5de 100644 --- a/tests/nebula_agent_test.py +++ b/tests/nebula_agent_test.py @@ -5,7 +5,6 @@ import pathlib import pytest -from freezegun import freeze_time from ostorlab.agent import definitions as agent_definitions from ostorlab.agent.message import message as msg from ostorlab.runtimes import definitions as runtime_definitions @@ -13,7 +12,6 @@ from pyfakefs import fake_filesystem_unittest from agent import nebula_agent -import os def testAgentNebula_whenUnsupportedFileType_raisesValueError() -> None: @@ -39,9 +37,9 @@ def testAgentNebula_whenUnsupportedFileType_raisesValueError() -> None: def testAgentNebula_whenFileTypeIsJson_persistMessage( - agent_definition: agent_definitions.AgentDefinition, - agent_settings: runtime_definitions.AgentSettings, - link_message: msg.Message, + agent_definition: agent_definitions.AgentDefinition, + agent_settings: runtime_definitions.AgentSettings, + link_message: msg.Message, ) -> None: """Test that NebulaAgent persists message to json file.""" os.environ["UNIVERSE"] = "43" @@ -63,9 +61,9 @@ def testAgentNebula_whenFileTypeIsJson_persistMessage( def testAgentNebula_whenFileTypeIsJson_persistMultipleLinkMessages( - agent_definition: agent_definitions.AgentDefinition, - agent_settings: runtime_definitions.AgentSettings, - multiple_link_messages: list[msg.Message], + agent_definition: agent_definitions.AgentDefinition, + agent_settings: runtime_definitions.AgentSettings, + multiple_link_messages: list[msg.Message], ) -> None: """Test that NebulaAgent persists multiple link messages to json file.""" os.environ["UNIVERSE"] = "43" @@ -93,9 +91,9 @@ def testAgentNebula_whenFileTypeIsJson_persistMultipleLinkMessages( def testAgentNebula_whenFileTypeIsJson_persistMultipleMessages( - agent_definition: agent_definitions.AgentDefinition, - agent_settings: runtime_definitions.AgentSettings, - multiple_messages: list[msg.Message], + agent_definition: agent_definitions.AgentDefinition, + agent_settings: runtime_definitions.AgentSettings, + multiple_messages: list[msg.Message], ) -> None: """Test that NebulaAgent persists multiple messages of different types to json files.""" os.environ["UNIVERSE"] = "43" From f6a8182c63198c68e33f1c2608768378abe217b5 Mon Sep 17 00:00:00 2001 From: mohsinenar Date: Thu, 28 Mar 2024 14:33:00 +0000 Subject: [PATCH 5/5] Apply ruff. --- agent/nebula_agent.py | 6 +++--- tests/nebula_agent_test.py | 18 +++++++++--------- 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/agent/nebula_agent.py b/agent/nebula_agent.py index 648d856..15f772e 100644 --- a/agent/nebula_agent.py +++ b/agent/nebula_agent.py @@ -34,9 +34,9 @@ class NebulaAgent(agent.Agent): """Agent responsible for persisting all types of messages to file type specified in the agent definition.""" def __init__( - self, - agent_definition: agent_definitions.AgentDefinition, - agent_settings: runtime_definitions.AgentSettings, + self, + agent_definition: agent_definitions.AgentDefinition, + agent_settings: runtime_definitions.AgentSettings, ) -> None: super().__init__(agent_definition, agent_settings) self._file_type = self.args.get("file_type", "json") diff --git a/tests/nebula_agent_test.py b/tests/nebula_agent_test.py index 6f1c5de..9f49748 100644 --- a/tests/nebula_agent_test.py +++ b/tests/nebula_agent_test.py @@ -37,9 +37,9 @@ def testAgentNebula_whenUnsupportedFileType_raisesValueError() -> None: def testAgentNebula_whenFileTypeIsJson_persistMessage( - agent_definition: agent_definitions.AgentDefinition, - agent_settings: runtime_definitions.AgentSettings, - link_message: msg.Message, + agent_definition: agent_definitions.AgentDefinition, + agent_settings: runtime_definitions.AgentSettings, + link_message: msg.Message, ) -> None: """Test that NebulaAgent persists message to json file.""" os.environ["UNIVERSE"] = "43" @@ -61,9 +61,9 @@ def testAgentNebula_whenFileTypeIsJson_persistMessage( def testAgentNebula_whenFileTypeIsJson_persistMultipleLinkMessages( - agent_definition: agent_definitions.AgentDefinition, - agent_settings: runtime_definitions.AgentSettings, - multiple_link_messages: list[msg.Message], + agent_definition: agent_definitions.AgentDefinition, + agent_settings: runtime_definitions.AgentSettings, + multiple_link_messages: list[msg.Message], ) -> None: """Test that NebulaAgent persists multiple link messages to json file.""" os.environ["UNIVERSE"] = "43" @@ -91,9 +91,9 @@ def testAgentNebula_whenFileTypeIsJson_persistMultipleLinkMessages( def testAgentNebula_whenFileTypeIsJson_persistMultipleMessages( - agent_definition: agent_definitions.AgentDefinition, - agent_settings: runtime_definitions.AgentSettings, - multiple_messages: list[msg.Message], + agent_definition: agent_definitions.AgentDefinition, + agent_settings: runtime_definitions.AgentSettings, + multiple_messages: list[msg.Message], ) -> None: """Test that NebulaAgent persists multiple messages of different types to json files.""" os.environ["UNIVERSE"] = "43"