Merge pull request #76 from Ostorlab/feature/add_support_for_proxy

Add support for proxy

agent/agent_nuclei.py

@@ -129,10 +129,11 @@ def __init__(
         agent.Agent.__init__(self, agent_definition, agent_settings)
         agent_persist_mixin.AgentPersistMixin.__init__(self, agent_settings)
         agent_report_vulnerability_mixin.AgentReportVulnMixin.__init__(self)
-        self._scope_urls_regex: Optional[str] = self.args.get("scope_urls_regex")
-        self._vpn_config: Optional[str] = self.args.get("vpn_config")
-        self._dns_config: Optional[str] = self.args.get("dns_config")
+        self._scope_urls_regex: str | None = self.args.get("scope_urls_regex")
+        self._vpn_config: str | None = self.args.get("vpn_config")
+        self._dns_config: str | None = self.args.get("dns_config")
         self._basic_credentials: list[BasicCredential] = []
+        self._proxy: str | None = self.args.get("proxy")
 
     def start(self) -> None:
         """Enable VPN configuration at the beginning if needed."""
@@ -454,6 +455,8 @@ def _run_command(
         ]
         for chunk in chunks:
             command = ["/nuclei/nuclei"]
+            if self._proxy is not None:
+                command.extend(["-proxy", self._proxy])
             for item in chunk:
                 command.extend(["-u", item])
             command.extend(["-j", "-irr", "-silent", "-o", OUTPUT_PATH])

ostorlab.yaml

@@ -107,4 +107,6 @@ args:
   - name: "basic_credentials"
     type: "array"
     description: "Credentials for basic authentication."
-
+  - name: "proxy"
+    type: "string"
+    description: "Proxy to use for the scan with nuclei."

tests/agent_test.py

@@ -724,3 +724,28 @@ def testPrepareTargets_whenIPv6AssetDoesNotReachCIDRLimit_doesNotRaiseValueError
     nuclei_agent: agent_nuclei.AgentNuclei,
 ) -> None:
     nuclei_agent.prepare_targets(scan_message_ipv6_with_mask112)
+
+
+@mock.patch("agent.agent_nuclei.OUTPUT_PATH", "./tests/result_nuclei.json")
+def testAgentNuclei_whenProxyIsProvided_shouldCallWithProxyArg(
+    scan_message: message.Message,
+    nuclei_agent_with_proxy: agent_nuclei.AgentNuclei,
+    mocker: plugin.MockerFixture,
+    requests_mock: rq_mock.mocker.Mocker,
+) -> None:
+    """Tests running the agent when proxy is provided."""
+    run_command_mock = mocker.patch("subprocess.run", return_value=None)
+    mocker.patch(
+        "agent.agent_nuclei.AgentNuclei.report_vulnerability", return_value=None
+    )
+    requests_mock.get("https://template1.yaml", json={})
+    requests_mock.get("https://template2.yaml", json={})
+
+    nuclei_agent_with_proxy.process(scan_message)
+
+    run_command_args = run_command_mock.call_args_list
+    command = " ".join(run_command_args[0].args[0])
+
+    assert "/nuclei/nuclei" in command
+    assert "-proxy" in command
+    assert "https://proxy.co" in command

tests/conftest.py

@@ -195,6 +195,34 @@ def nuclei_agent_with_basic_credentials(
         return agent_object
 
 
+@pytest.fixture
+def nuclei_agent_with_proxy(
+    agent_mock: list[message.Message],
+    agent_persist_mock: dict[str | bytes, str | bytes],
+) -> agent_nuclei.AgentNuclei:
+    del agent_mock
+    del agent_persist_mock
+    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/nuclei",
+            bus_url="NA",
+            bus_exchange_topic="NA",
+            args=[
+                utils_definitions.Arg(
+                    name="proxy",
+                    type="string",
+                    value=json.dumps("https://proxy.co").encode(),
+                )
+            ],
+            healthcheck_port=random.randint(5000, 6000),
+            redis_url="redis://guest:guest@localhost:6379",
+        )
+
+        agent_object = agent_nuclei.AgentNuclei(definition, settings)
+        return agent_object
+
+
 @pytest.fixture
 def nuclei_agent_with_custom_templates(
     agent_mock: list[message.Message],