Fix comments

agent/subfinder_agent.py

@@ -1,6 +1,7 @@
 """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
@@ -26,32 +27,33 @@
 CONFIG_PATH = "/root/.config/subfinder/provider-config.yaml"
 
 
-def update_provider_config(
+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
-    # Load existing configuration or initialize a new one
-    try:
-        with open(config_path, "r") as config_file:
-            config = yaml.load(config_file) or {}
-    except FileNotFoundError:
-        logger.error("Configuration file not found. Creating a new one.")
-        return
-    # Update the 'virustotal' section
+    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)  # Avoid duplicate keys
+            config["virustotal"].append(virustotal_key)
     else:
-        config["virustotal"] = [virustotal_key]  # Add a new entry
+        # add to sources
+        if "sources" in config and "virustotal" not in config["sources"]:
+            config["sources"].append("virustotal")
+        config["virustotal"] = [virustotal_key]
 
-    # Write back the updated configuration
     try:
-        with open(config_path, "w") as config_file:
-            yaml.dump(config, config_file)
-        logger.info("VirusTotal API key has been added to the configuration.")
+        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)
 
@@ -66,10 +68,10 @@ def __init__(
     ) -> None:
         agent.Agent.__init__(self, agent_definition, agent_settings)
 
-        virustotal_key = self.args.get("virustotal_key")
+        virustotal_key = self.args.get("virustotal_api_key")
         if virustotal_key is not None:
             logger.info("Updating configuration with VirusTotal API key.")
-            update_provider_config(virustotal_key)
+            set_virustotal_api_key(virustotal_key)
 
         agent_persist_mixin.AgentPersistMixin.__init__(self, agent_settings)
 

ostorlab.yaml

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

tests/conftest.py

@@ -57,7 +57,7 @@ def subfinder_definition() -> agent_definitions.AgentDefinition:
         definition = agent_definitions.AgentDefinition.from_yaml(yaml_o)
         definition.args = [
             {
-                "name": "virustotal_key",
+                "name": "virustotal_api_key",
                 "value": "Justrandomvalue",
                 "type": "string",
             },

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

@@ -4,6 +4,8 @@
 
 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
@@ -82,7 +84,7 @@ def testAgentSubfinder_whenMaxSubDomainsSet_emitsBackFindings(
     )
 
 
-def testAgentSubfinder_whenVirustotalKeyPassed_emitsBackFindings(
+def testAgentSubfinder_always_callsSetVirusTotalKeyInInit(
     subfinder_definition: agent_definitions.AgentDefinition,
     subfinder_settings: runtime_definitions.AgentSettings,
     mocker: plugin.MockerFixture,
@@ -91,35 +93,105 @@ def testAgentSubfinder_whenVirustotalKeyPassed_emitsBackFindings(
     Test that the Subfinder agent correctly updates the provider configuration
     with the VirusTotal key.
     """
-    mocker_update_provider_config = mocker.patch(
-        "agent.subfinder_agent.update_provider_config"
+    mocker_set_virustotal_api_key = mocker.patch(
+        "agent.subfinder_agent.set_virustotal_api_key"
     )
 
     sub_agent.SubfinderAgent(subfinder_definition, subfinder_settings)
 
-    assert mocker_update_provider_config.called is True
-    assert mocker_update_provider_config.call_args[0][0] == "Justrandomvalue"
+    assert mocker_set_virustotal_api_key.called is True
+    assert mocker_set_virustotal_api_key.call_args[0][0] == "Justrandomvalue"
 
 
-def testUpdateConfigurationFile_whenConfNotFound_handelFileNotFoundError(
+def testSetVirusTotalApiKey_whenConfFileNotFound_returnNoneAndLogError(
     caplog: pytest.LogCaptureFixture,
 ) -> None:
     """Test that the provider configuration correctly handles a missing configuration file."""
 
-    sub_agent.update_provider_config("existing_key", "test_not.yaml")
+    sub_agent.set_virustotal_api_key("existing_key", "test_not.yaml")
 
-    assert "Configuration file not found. Creating a new one." in caplog.text
+    assert "Configuration file not found at test_not.yaml." in caplog.text
 
 
-def testupdateconfigupdate_whenWriteConfigurationFail_handelWriteErro(
+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=FileNotFoundError)
 
-    sub_agent.update_provider_config(
+    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"]
+        assert "sources" in updated_config and "virustotal" in updated_config["sources"]
+
+
+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"]
+        assert "sources" in updated_config and "virustotal" in updated_config["sources"]
+
+
+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"]
+        assert "sources" in updated_config and "virustotal" in updated_config["sources"]

tests/test-requirement.txt

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