Merge pull request #2 from Ostorlab/feature/persist-json

Persist to json file

Dockerfile

@@ -12,4 +12,4 @@ ENV PYTHONPATH=/app
 COPY agent /app/agent
 COPY ostorlab.yaml /app/agent/ostorlab.yaml
 WORKDIR /app
-CMD ["python3", "/app/agent/template_agent.py"]
+CMD ["python3", "/app/agent/nebula_agent.py"]

README.md

@@ -71,4 +71,7 @@ git clone https://github.com/Ostorlab/agent_nebula.git && cd agent_nebula
 	 * If you specified an organization when building the image:
     ```shell
     ostorlab scan run --agent agent/[ORGANIZATION]/nebula link --url www.yourdomain.com --method GET
-    ```
+    ```
+
+### License
+    [Apache-2.0](./LICENSE)

agent/nebula_agent.py

@@ -1,6 +1,11 @@
 """Nebula agent: responsible for persisting all types of messages."""
 
+import base64
+import json
 import logging
+import os
+from datetime import datetime
+from typing import Any
 
 from ostorlab.agent import agent, definitions as agent_definitions
 from ostorlab.agent.message import message as m
@@ -16,6 +21,15 @@
 )
 logger = logging.getLogger(__name__)
 
+SUPPORTED_FILE_TYPES = ["json"]
+
+
+class CustomEncoder(json.JSONEncoder):
+    def default(self, obj: Any) -> Any:
+        if isinstance(obj, bytes) is True:
+            return base64.b64encode(obj).decode("utf-8")
+        return json.JSONEncoder.default(self, obj)
+
 
 class NebulaAgent(agent.Agent):
     """Agent responsible for persisting all types of messages to file type specified in the agent definition."""
@@ -26,22 +40,40 @@ def __init__(
         agent_settings: runtime_definitions.AgentSettings,
     ) -> None:
         super().__init__(agent_definition, agent_settings)
-        self._file_type = self.args.get("file_type")
-        self._file_path = self.args.get("file_path")
+        self._file_type = self.args.get("file_type", "json")
+        if self._file_type.lower() not in SUPPORTED_FILE_TYPES:
+            raise ValueError(
+                f"File type {self._file_type} is not supported. Supported file types are {SUPPORTED_FILE_TYPES}"
+            )
+
+        self._output_folder = (
+            f"/output/messages_{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}"
+        )
+        os.makedirs(self._output_folder, exist_ok=True)
 
     def process(self, message: m.Message) -> None:
         """Process the message and persist it to the file type and location specified in the agent definition.
 
         Args:
             message: The message to process.
         """
-        logger.info("processing message of selector : %s", message.selector)
-        # TODO (elyousfi5): add the logic to persist the message to the file type and location specified in the agent
-        logger.info(
-            "message persisted to file type: %s at location: %s",
-            self._file_type,
-            self._file_path,
-        )
+        logger.info("Processing message of selector : %s", message.selector)
+
+        if self._file_type == "json":
+            self._persist_to_json(message)
+
+    def _persist_to_json(self, message_to_persist: m.Message) -> None:
+        """Persist message to JSON file.
+
+        Args:
+            message_to_persist: The message to persist.
+        """
+        data = message_to_persist.data
+        selector = message_to_persist.selector
+        file_name = f"{self._output_folder}/{selector}_messages.json"
+
+        with open(file_name, "a") as file:
+            file.write(json.dumps(data, cls=CustomEncoder) + "\n")
 
 
 if __name__ == "__main__":

ostorlab.yaml

@@ -1,6 +1,7 @@
 kind: Agent
 name: nebula
 version: 0.0.1
+image: images/logo.png
 description: |
   _The Nebula Agent is responsible for persisting all types of messages locally._
 
@@ -61,18 +62,21 @@ description: |
       ```shell
       ostorlab scan run --agent agent/[ORGANIZATION]/nebula link --url www.yourdomain.com --method GET
       ```
+  ### License
+    [Apache-2.0](./LICENSE)
+
+license: Apache-2.0
+source: https://github.com/Ostorlab/agent_nebula
 in_selectors:
-  - v3.asset.ip.v4
-  - v3.asset.ip.v6
-  - v3.asset.domain_name
-  - v3.asset.link
+  - v3
 out_selectors: []
 docker_file_path : Dockerfile
 docker_build_root : .
+mounts:
+  - '$CONFIG_HOME/:/output/'
 args:
   - name: "file_type"
     type: "string"
     description: "The type of the file where the message will be persisted."
-  - name: "file_path"
-    type: "string"
-    description: "The path of the file where the message will be persisted."
+    value: "json"
+

tests/conftest.py

@@ -1 +1,63 @@
 """Pytest fixtures for Agent Nebula."""
+
+import pathlib
+import random
+
+import pytest
+from ostorlab.agent import definitions as agent_definitions
+from ostorlab.agent.message import message as msg
+from ostorlab.runtimes import definitions as runtime_definitions
+
+
+@pytest.fixture(scope="function", name="agent_definition")
+def agent_definition() -> agent_definitions.AgentDefinition:
+    """NebulaAgent definition fixture for testing purposes."""
+    with (pathlib.Path(__file__).parent.parent / "ostorlab.yaml").open() as yaml_o:
+        return agent_definitions.AgentDefinition.from_yaml(yaml_o)
+
+
+@pytest.fixture(scope="function", name="agent_settings")
+def agent_settings() -> runtime_definitions.AgentSettings:
+    """NebulaAgent settings fixture for testing purposes."""
+    return runtime_definitions.AgentSettings(
+        key="agent/ostorlab/nebula",
+        bus_url="NA",
+        bus_exchange_topic="NA",
+        healthcheck_port=random.randint(5000, 6000),
+        redis_url="redis://guest:guest@localhost:6379",
+    )
+
+
+@pytest.fixture
+def link_message() -> msg.Message:
+    """Creates a dummy message of type v3.asset.link to be used by the agent for testing purposes."""
+    selector = "v3.asset.link"
+    msg_data = {"url": "https://ostorlab.co", "method": b"GET"}
+    return msg.Message.from_data(selector, data=msg_data)
+
+
+@pytest.fixture
+def multiple_link_messages() -> list[msg.Message]:
+    """Creates dummy messages of type v3.asset.link to be used by the agent for testing purposes."""
+    selector = "v3.asset.link"
+    return [
+        msg.Message.from_data(
+            selector, data={"url": f"https://www.domain{i}.com", "method": b"GET"}
+        )
+        for i in range(0, 5)
+    ]
+
+
+@pytest.fixture
+def multiple_messages() -> list[msg.Message]:
+    """Creates dummy messages of type v3.asset.link, v3.asset.domain, v3.asset.ip to be used by the agent for testing
+    purposes."""
+    return [
+        msg.Message.from_data(
+            "v3.asset.link", data={"url": "https://www.domain.com", "method": b"GET"}
+        ),
+        msg.Message.from_data("v3.asset.domain_name", data={"name": "www.domain.com"}),
+        msg.Message.from_data(
+            "v3.asset.ip", data={"host": "192.168.1.1", "mask": "24"}
+        ),
+    ]

tests/nebula_agent_test.py

@@ -1,9 +1,137 @@
 """Unit tests for Nebula agent."""
 
+import json
+import os
+import pathlib
 
-# TODO (elyousfi5): add tests for the Nebula agent
+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
+from ostorlab.utils import defintions as utils_definitions
+from pyfakefs import fake_filesystem_unittest
 
+from agent import nebula_agent
 
-def testNebulaAgent_always_persistMessages() -> None:
-    """Test Nebula agent."""
-    pass
+
+def testAgentNebula_whenUnsupportedFileType_raisesValueError() -> None:
+    """Test that NebulaAgent raises ValueError when file type is not supported."""
+    with pytest.raises(ValueError):
+        with (pathlib.Path(__file__).parent.parent / "ostorlab.yaml").open() as yaml_o:
+            definition = agent_definitions.AgentDefinition.from_yaml(yaml_o)
+            settings = runtime_definitions.AgentSettings(
+                key="agent/ostorlab/nebula",
+                bus_url="NA",
+                bus_exchange_topic="NA",
+                args=[
+                    utils_definitions.Arg(
+                        name="file_type",
+                        type="string",
+                        value=json.dumps("txt").encode(),
+                    ),
+                ],
+                healthcheck_port=5301,
+                redis_url="redis://guest:guest@localhost:6379",
+            )
+            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."""
+    with fake_filesystem_unittest.Patcher():
+        expected_output = json.dumps(
+            {"url": "https://ostorlab.co", "method": b"GET"},
+            cls=nebula_agent.CustomEncoder,
+        )
+        nebula_test_agent = nebula_agent.NebulaAgent(agent_definition, agent_settings)
+
+        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
+        with open(
+            "/output/messages_2024-03-05_12-00-00/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."""
+    with fake_filesystem_unittest.Patcher():
+        expected_output = [
+            json.dumps(
+                {"url": f"https://www.domain{i}.com", "method": b"GET"},
+                cls=nebula_agent.CustomEncoder,
+            )
+            for i in range(0, 5)
+        ]
+        nebula_test_agent = nebula_agent.NebulaAgent(agent_definition, agent_settings)
+
+        for message in multiple_link_messages:
+            nebula_test_agent.process(message)
+
+        file_path = "/output/messages_2023-03-05_12-00-00"
+        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:
+            lines = file.readlines()
+        assert len(lines) == len(expected_output)
+        for line, expected_line in zip(lines, expected_output):
+            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."""
+    with fake_filesystem_unittest.Patcher():
+        expected_output = [
+            json.dumps(
+                {"url": "https://www.domain.com", "method": b"GET"},
+                cls=nebula_agent.CustomEncoder,
+            ),
+            json.dumps({"name": "www.domain.com"}, cls=nebula_agent.CustomEncoder),
+            json.dumps(
+                {"host": "192.168.1.1", "mask": "24"},
+                cls=nebula_agent.CustomEncoder,
+            ),
+        ]
+        nebula_test_agent = nebula_agent.NebulaAgent(agent_definition, agent_settings)
+
+        for message in multiple_messages:
+            nebula_test_agent.process(message)
+
+        file_path = "/output/messages_2023-03-05_12-00-00"
+        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
+        with open(f"{file_path}/v3.asset.link_messages.json", "r") as file:
+            lines = file.readlines()
+        assert len(lines) == 1
+        assert lines[0].strip() == expected_output[0].strip()
+        assert os.path.exists(f"{file_path}/v3.asset.domain_name_messages.json") is True
+        with open(f"{file_path}/v3.asset.domain_name_messages.json", "r") as file:
+            lines = file.readlines()
+        assert len(lines) == 1
+        assert lines[0].strip() == expected_output[1].strip()
+        assert os.path.exists(f"{file_path}/v3.asset.ip_messages.json") is True
+        with open(f"{file_path}/v3.asset.ip_messages.json", "r") as file:
+            lines = file.readlines()
+        assert len(lines) == 1
+        assert lines[0].strip() == expected_output[2].strip()

tests/test-requirement.txt

@@ -2,3 +2,6 @@ pytest
 ruff
 mypy
 typing-extensions
+pytest-mock
+freezegun
+pyfakefs