Merge pull request #27 from Ostorlab/feature/adding_virustotal_key

Adding Virus_total_key argument

agent/subfinder_agent.py

@@ -1,7 +1,9 @@
 """Agent implementation for Subfinder : subdomain discovery tool that discovers valid subdomains for websites."""
 
 import logging
+import pathlib
 
+import ruamel.yaml
 from rich import logging as rich_logging
 import tld
 from ostorlab.agent import agent
@@ -22,6 +24,35 @@
 )
 logger = logging.getLogger(__name__)
 STORAGE_NAME = "agent_subfinder_storage"
+CONFIG_PATH = "/root/.config/subfinder/provider-config.yaml"
+
+
+def set_virustotal_api_key(
+    virustotal_key: str,
+    config_path: str = CONFIG_PATH,
+) -> None:
+    """Update the Subfinder provider configuration file with the VirusTotal API key."""
+    yaml = ruamel.yaml.YAML(typ="safe")
+    yaml.default_flow_style = False  # Ensure block-style lists
+    config_path = pathlib.Path(config_path)
+
+    if config_path.exists() is False:
+        logger.error("Configuration file not found at %s.", config_path)
+        return None
+
+    config = yaml.load(config_path.read_text()) or {}
+
+    if "virustotal" in config:
+        if virustotal_key not in config["virustotal"]:
+            config["virustotal"].append(virustotal_key)
+    else:
+        config["virustotal"] = [virustotal_key]
+
+    try:
+        with config_path.open("w") as file:
+            yaml.dump(config, file)
+    except (IOError, OSError) as write_error:
+        logger.error("Failed to write configuration file: %s", write_error)
 
 
 class SubfinderAgent(agent.Agent, agent_persist_mixin.AgentPersistMixin):
@@ -33,6 +64,12 @@ def __init__(
         agent_settings: runtime_definitions.AgentSettings,
     ) -> None:
         agent.Agent.__init__(self, agent_definition, agent_settings)
+
+        virustotal_key = self.args.get("virustotal_api_key")
+        if virustotal_key is not None:
+            logger.info("Updating configuration with VirusTotal API key.")
+            set_virustotal_api_key(virustotal_key)
+
         agent_persist_mixin.AgentPersistMixin.__init__(self, agent_settings)
 
     def process(self, message: m.Message) -> None:

ostorlab.yaml

@@ -60,3 +60,6 @@ args:
   - name: "max_subdomains"
     type: "number"
     description: "Maximum number of subdomains to return"
+  - name : "virustotal_api_key"
+    type: "string"
+    description: "Adding VirusTotal api key to get more data"

requirement.txt

@@ -1,3 +1,4 @@
 ostorlab[agent]
 rich
-tld
+tld
+ruamel.yaml

tests/conftest.py

@@ -49,3 +49,30 @@ def fixture_subfinder_agent_max_subdomains():
 
         agent = subfinder_agent.SubfinderAgent(definition, settings)
         return agent
+
+
+@pytest.fixture
+def subfinder_definition() -> agent_definitions.AgentDefinition:
+    with (pathlib.Path(__file__).parent.parent / "ostorlab.yaml").open() as yaml_o:
+        definition = agent_definitions.AgentDefinition.from_yaml(yaml_o)
+        definition.args = [
+            {
+                "name": "virustotal_api_key",
+                "value": "Justrandomvalue",
+                "type": "string",
+            },
+        ]
+        return definition
+
+
+@pytest.fixture
+def subfinder_settings() -> runtime_definitions.AgentSettings:
+    settings = runtime_definitions.AgentSettings(
+        key="agent/ostorlab/subfinder",
+        bus_url="NA",
+        bus_exchange_topic="NA",
+        args=[],
+        healthcheck_port=random.randint(5000, 6000),
+        redis_url="redis://guest:guest@localhost:6379",
+    )
+    return settings

tests/provider-config-no-virustotal.yaml

@@ -0,0 +1,20 @@
+resolvers:
+    - 1.1.1.1
+    - 1.0.0.1
+    - 8.8.8.8
+    - 8.8.4.4
+    - 9.9.9.9
+
+sources:
+    - github
+    - zoomeyeapi
+    - quake
+github:
+  - ghp_lkyJGU3jv1xmwk4SDXavrLDJ4dl2pSJMzj4X
+  - ghp_gkUuhkIYdQPj13ifH4KA3cXRn8JD2lqir2d4
+zoomeyeapi:
+  - 4f73021d-ff95-4f53-937f-83d6db719eec
+quake:
+  - 0cb9030c-0a40-48a3-b8c4-fca28e466ba3
+
+subfinder-version: 2.4.5

tests/provider-config-virustotal.yaml

@@ -0,0 +1,24 @@
+resolvers:
+    - 1.1.1.1
+    - 1.0.0.1
+    - 8.8.8.8
+    - 8.8.4.4
+    - 9.9.9.9
+
+sources:
+    - github
+    - zoomeyeapi
+    - quake
+    - virustotal
+
+github:
+  - ghp_lkyJGU3jv1xmwk4SDXavrLDJ4dl2pSJMzj4X
+  - ghp_gkUuhkIYdQPj13ifH4KA3cXRn8JD2lqir2d4
+zoomeyeapi:
+  - 4f73021d-ff95-4f53-937f-83d6db719eec
+quake:
+  - 0cb9030c-0a40-48a3-b8c4-fca28e466ba3
+virustotal:
+  - example-api-key
+
+subfinder-version: 2.4.5

tests/subfinder_agent_test.py

@@ -1,6 +1,16 @@
 """Unittests for the Subfinder Agent."""
 
+import pathlib
+
+import pytest
+from pytest_mock import plugin
+import ruamel.yaml
+from pyfakefs import fake_filesystem_unittest
+
 from ostorlab.agent.message import message
+from agent import subfinder_agent as sub_agent
+from ostorlab.agent import definitions as agent_definitions
+from ostorlab.runtimes import definitions as runtime_definitions
 
 
 def testAgentSubfinder_whenFindsSubDomains_emitsBackFindings(
@@ -72,3 +82,133 @@ def testAgentSubfinder_whenMaxSubDomainsSet_emitsBackFindings(
     assert agent_mock[0].selector == "v3.asset.domain_name", (
         agent_mock[0].data["name"] == "subdomain1"
     )
+
+
+def testAgentSubfinder_always_callsSetVirusTotalKeyInInit(
+    subfinder_definition: agent_definitions.AgentDefinition,
+    subfinder_settings: runtime_definitions.AgentSettings,
+    mocker: plugin.MockerFixture,
+) -> None:
+    """
+    Test that the Subfinder agent correctly updates the provider configuration
+    with the VirusTotal key.
+    """
+    mocker_set_virustotal_api_key = mocker.patch(
+        "agent.subfinder_agent.set_virustotal_api_key"
+    )
+
+    sub_agent.SubfinderAgent(subfinder_definition, subfinder_settings)
+
+    assert mocker_set_virustotal_api_key.called is True
+    assert mocker_set_virustotal_api_key.call_args[0][0] == "Justrandomvalue"
+
+
+def testSetVirusTotalApiKey_whenConfFileNotFound_returnNoneAndLogError(
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test that the provider configuration correctly handles a missing configuration file."""
+
+    sub_agent.set_virustotal_api_key("existing_key", "test_not.yaml")
+
+    assert "Configuration file not found at test_not.yaml." in caplog.text
+
+
+def testSetVirusTotalApiKey_whenWriteConfigurationFail_handleWriteError(
+    mocker: plugin.MockerFixture,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test that the provider configuration handles a write error correctly and logs the failure."""
+    mocker.patch("ruamel.yaml.main.YAML.dump", side_effect=IOError)
+
+    sub_agent.set_virustotal_api_key(
+        "existing_key",
+        str(pathlib.Path(__file__).parent / "provider-config.yaml"),
+    )
+
+    assert "Failed to write configuration file" in caplog.text
+
+
+def testSetVirusTotalApiKey_createsSectionAndAddsKeyWhenNoSectionExists() -> None:
+    """
+    Test that the function creates a `virustotal` section and adds the key
+    when it does not exist in the configuration.
+    """
+    real_file_path = (
+        pathlib.Path(__file__).parent / "provider-config-no-virustotal.yaml"
+    )
+    fake_file_path = "/fake/path/provider-config-no-virustotal.yaml"
+    file_contents = real_file_path.read_text()
+
+    with fake_filesystem_unittest.Patcher() as patcher:
+        patcher.fs.create_file(fake_file_path, contents=file_contents)
+
+        sub_agent.set_virustotal_api_key("new_key", fake_file_path)
+
+        yaml = ruamel.yaml.YAML(typ="safe")
+        fake_file = pathlib.Path(fake_file_path)
+        updated_config = yaml.load(fake_file.read_text()) or {}
+        assert "virustotal" in updated_config
+        assert updated_config["virustotal"] == ["new_key"]
+
+
+def testSetVirusTotalApiKey_whenVirusTotalSectionExists_addsKeyToExistingSection() -> (
+    None
+):
+    """
+    Test that the function adds the key to the existing `virustotal` section
+    when it already exists in the configuration.
+    """
+    real_file_path = pathlib.Path(__file__).parent / "provider-config-virustotal.yaml"
+    fake_file_path = "/fake/path/provider-config-virustotal.yaml"
+    file_contents = real_file_path.read_text()
+
+    with fake_filesystem_unittest.Patcher() as patcher:
+        patcher.fs.create_file(fake_file_path, contents=file_contents)
+
+        sub_agent.set_virustotal_api_key("new_key", fake_file_path)
+
+        yaml = ruamel.yaml.YAML(typ="safe")
+        fake_file = pathlib.Path(fake_file_path)
+        updated_config = yaml.load(fake_file.read_text()) or {}
+        assert "virustotal" in updated_config
+        assert updated_config["virustotal"] == ["example-api-key", "new_key"]
+
+
+def testSetVirusTotalApiKey_whenKeyAlreadyExists_doesNotAddKeyAgain() -> None:
+    """
+    Test that the function does not add the key to the `virustotal` section
+    when it already exists in the configuration.
+    """
+    real_file_path = pathlib.Path(__file__).parent / "provider-config-virustotal.yaml"
+    fake_file_path = "/fake/path/provider-config-virustotal.yaml"
+    file_contents = real_file_path.read_text()
+
+    with fake_filesystem_unittest.Patcher() as patcher:
+        patcher.fs.create_file(fake_file_path, contents=file_contents)
+
+        sub_agent.set_virustotal_api_key("example-api-key", fake_file_path)
+
+        yaml = ruamel.yaml.YAML(typ="safe")
+        fake_file = pathlib.Path(fake_file_path)
+        updated_config = yaml.load(fake_file.read_text()) or {}
+        assert "virustotal" in updated_config
+        assert updated_config["virustotal"] == ["example-api-key"]
+
+
+def testSetVirusTotalApiKey_whenFileEmpty_addVirusTotalKey() -> None:
+    """
+    Test that the function adds the `virustotal` key
+    when the configuration file is empty.
+    """
+    fake_file_path = "/fake/path/provider-config-empty.yaml"
+
+    with fake_filesystem_unittest.Patcher() as patcher:
+        patcher.fs.create_file(fake_file_path, contents="")
+
+        sub_agent.set_virustotal_api_key("new_key", fake_file_path)
+
+        yaml = ruamel.yaml.YAML(typ="safe")
+        fake_file = pathlib.Path(fake_file_path)
+        updated_config = yaml.load(fake_file.read_text()) or {}
+        assert "virustotal" in updated_config
+        assert updated_config["virustotal"] == ["new_key"]

tests/test-requirement.txt

@@ -2,4 +2,5 @@ ruff
 pytest
 pytest-cov
 pytest-mock
-tld
+tld
+pyfakefs