Merge pull request #849 from Ostorlab/feature/add-option-to-change-sc…

…an-timeout-in-local-scans

Add option to change scan timeout in local scans

.github/workflows/pytest-linux.yml

@@ -28,7 +28,6 @@ jobs:
     - uses: actions/[email protected]
       with:
         submodules: true
-        token: ${{ secrets.OXOTITAN_GITHUB_KEY }}
     - name: Set up Python ${{ matrix.python-version }}
       uses: actions/[email protected]
       with:

.github/workflows/pytest-windows.yml

@@ -28,7 +28,6 @@ jobs:
     - uses: actions/[email protected]
       with:
         submodules: true
-        token: ${{ secrets.OXOTITAN_GITHUB_KEY }}
     - name: Set up Python ${{ matrix.python-version }}
       uses: actions/[email protected]
       with:

.github/workflows/python-publish.yml

@@ -19,7 +19,6 @@ jobs:
     - uses: actions/[email protected]
       with:
         submodules: true
-        token: ${{ secrets.OXOTITAN_GITHUB_KEY }}
     - name: Setup Node.js
       uses: actions/setup-node@v2
       with:

.github/workflows/release_linux_mac_bin.yaml

@@ -16,7 +16,6 @@ jobs:
     - uses: actions/checkout@v4
       with:
         submodules: true
-        token: ${{ secrets.OXOTITAN_GITHUB_KEY }}
     - name: Setup Node.js
       uses: actions/setup-node@v2
       with:

.github/workflows/release_windows_bin.yaml

@@ -14,7 +14,6 @@ jobs:
     - uses: actions/checkout@v4
       with:
         submodules: true
-        token: ${{ secrets.OXOTITAN_GITHUB_KEY }}
     - name: Setup Node.js
       uses: actions/setup-node@v2
       with:

src/ostorlab/cli/scan/run/run.py

@@ -5,7 +5,7 @@
 
 import io
 import logging
-from typing import List
+from typing import List, Optional
 
 import click
 import httpx
@@ -89,6 +89,13 @@
     is_flag=True,
     required=False,
 )
+@click.option(
+    "--timeout",
+    "-t",
+    type=int,
+    help="Timeout for the scan in seconds",
+    required=False,
+)
 @click.pass_context
 def run(
     ctx: click.core.Context,
@@ -101,6 +108,7 @@ def run(
     follow: List[str],
     no_follow: bool,
     no_asset: bool,
+    timeout: Optional[int] = None,
 ) -> None:
     """Start a new scan on your assets.\n
     Example:\n
@@ -146,6 +154,8 @@ def run(
             console.error(f"{e}")
             raise click.ClickException("Invalid asset Group Definition.") from e
     runtime_instance: runtime.Runtime = ctx.obj["runtime"]
+    if timeout is not None:
+        runtime_instance.timeout = timeout
 
     # Prepare and set the list of agents to follow.
     agent_keys = [agent.key for agent in agent_group.agents]

src/ostorlab/runtimes/local/runtime.py

@@ -22,6 +22,7 @@
 from rich import panel
 from sqlalchemy import case
 
+from ostorlab.utils import definitions as utils_definitions
 from ostorlab import exceptions
 from ostorlab.assets import asset as base_asset
 from ostorlab.cli import console as cli_console
@@ -199,6 +200,7 @@ def scan(
             title: Scan title
             agent_group_definition: Agent run definition defines the set of agents and how agents are configured.
             assets: the target asset to scan.
+            timeout: The timeout in seconds for the tracker agent to wait for all agents to finish.
 
         Returns:
             The scan object.
@@ -324,7 +326,7 @@ def stop(self, scan_id: str) -> None:
             config_labels = config.attrs["Spec"]["Labels"]
             if (
                 config_labels.get("ostorlab.universe") is not None
-                and int(config_labels.get("ostorlab.universe")) == scan_id
+                and config_labels.get("ostorlab.universe") == scan_id
             ):
                 logger.info("removing config %s", config_labels)
                 stopped_configs.append(config)
@@ -519,7 +521,18 @@ def _list_agent_services(self):
 
     def _start_tracker_agent(self):
         """Start the tracker agent to handle the scan lifecycle."""
-        tracker_agent_settings = definitions.AgentSettings(key=TRACKER_AGENT_DEFAULT)
+        tracker_agent_settings = definitions.AgentSettings(
+            key=TRACKER_AGENT_DEFAULT,
+        )
+
+        if self.timeout is not None:
+            tracker_agent_settings.args.extend(
+                [
+                    utils_definitions.Arg(
+                        name="scan_done_timeout_sec", type="number", value=self.timeout
+                    ),
+                ]
+            )
         self._start_agent(agent=tracker_agent_settings, extra_configs=[])
 
     def _start_persist_vulnz_agent(self):

src/ostorlab/runtimes/runtime.py

@@ -24,6 +24,7 @@ class Runtime(abc.ABC):
     """Runtime is in charge of preparing the environment to trigger a scan."""
 
     follow: list
+    timeout: Optional[int] = None
 
     @abc.abstractmethod
     def can_run(self, agent_group_definition: definitions.AgentGroupDefinition) -> bool:

tests/cli/scan/run/run_test.py

@@ -537,6 +537,82 @@ def testOstorlabScanRunCLI_whenFollow_shouldFollowSpecifiedAgents(
     assert "agent/ostorlab_inject_asset" not in spy_follow.spy_return
 
 
+def testScanRunCLI_whenTimeoutProvided_setsTrackerAgentTimeout(
+    mocker: plugin.MockerFixture,
+    run_scan_mock2: None,
+    db_engine_path: str,
+) -> None:
+    """Test that timeout parameter is passed to tracker agent."""
+    mocker.patch.object(models, "ENGINE_URL", db_engine_path)
+    mocker.patch(
+        "ostorlab.runtimes.local.runtime.LocalRuntime.can_run", return_value=True
+    )
+    mocker.patch("ostorlab.runtimes.local.runtime.LocalRuntime._create_network")
+    mocker.patch("ostorlab.runtimes.local.runtime.LocalRuntime._start_services")
+    mocker.patch("ostorlab.runtimes.local.runtime.LocalRuntime._start_pre_agents")
+    mock_start_agent = mocker.patch(
+        "ostorlab.runtimes.local.runtime.LocalRuntime._start_agent"
+    )
+    runner = CliRunner()
+
+    runner.invoke(
+        rootcli.rootcli,
+        [
+            "scan",
+            "--runtime=local",
+            "run",
+            "--install",
+            "--agent=agent/ostorlab/nmap",
+            "--timeout=3600",
+            "ip",
+            "8.8.8.8",
+        ],
+    )
+
+    assert mock_start_agent.call_count == 2
+    assert mock_start_agent.call_args[1].get("agent").key == "agent/ostorlab/tracker"
+    assert (
+        mock_start_agent.call_args[1].get("agent").args[0].name
+        == "scan_done_timeout_sec"
+    )
+    assert mock_start_agent.call_args[1].get("agent").args[0].value == 3600
+
+
+def testScanRunCLI_whenNoTimeoutProvided_usesDefaultTimeout(
+    mocker: plugin.MockerFixture,
+    run_scan_mock3: None,
+    db_engine_path: str,
+) -> None:
+    """Test that default timeout is used when no timeout is provided."""
+    mocker.patch.object(models, "ENGINE_URL", db_engine_path)
+    mocker.patch(
+        "ostorlab.runtimes.local.runtime.LocalRuntime.can_run", return_value=True
+    )
+    mocker.patch("ostorlab.runtimes.local.runtime.LocalRuntime._create_network")
+    mocker.patch("ostorlab.runtimes.local.runtime.LocalRuntime._start_services")
+    mocker.patch("ostorlab.runtimes.local.runtime.LocalRuntime._start_pre_agents")
+    mock_start_agent = mocker.patch(
+        "ostorlab.runtimes.local.runtime.LocalRuntime._start_agent"
+    )
+    runner = CliRunner()
+
+    runner.invoke(
+        rootcli.rootcli,
+        [
+            "scan",
+            "--runtime=local",
+            "run",
+            "--agent=agent/ostorlab/nmap",
+            "ip",
+            "8.8.8.8",
+        ],
+    )
+
+    assert mock_start_agent.call_count == 1
+    assert mock_start_agent.call_args[1].get("agent").key == "agent/ostorlab/tracker"
+    assert len(mock_start_agent.call_args[1].get("agent").args) == 0
+
+
 def testOstorlabScanRunCLI_whenTestflightAsset_shouldRunCOmmand(
     mocker: plugin.MockerFixture,
 ) -> None:

tests/conftest.py

@@ -1270,8 +1270,8 @@ def run_scan_mock(mocker: plugin.MockerFixture) -> None:
     mocker.patch("ostorlab.runtimes.local.runtime.LocalRuntime._create_network")
     mocker.patch("ostorlab.runtimes.local.runtime.LocalRuntime._start_services")
     mocker.patch("ostorlab.runtimes.local.runtime.LocalRuntime._start_pre_agents")
-    mocker.patch("ostorlab.runtimes.local.runtime.LocalRuntime._start_agents")
     mocker.patch("ostorlab.runtimes.local.runtime.LocalRuntime._start_post_agents")
+    mocker.patch("ostorlab.runtimes.local.runtime.LocalRuntime._start_agents")
     mocker.patch("ostorlab.runtimes.local.runtime.LocalRuntime._check_services_healthy")
     mocker.patch(
         "ostorlab.runtimes.local.runtime.LocalRuntime._check_agents_healthy",
@@ -1300,6 +1300,37 @@ def run_scan_mock2(mocker: plugin.MockerFixture) -> None:
     )
 
 
+@pytest.fixture
+def run_scan_mock3(mocker: plugin.MockerFixture) -> None:
+    """Mock functions required to run a scan."""
+    mocker.patch(
+        "ostorlab.cli.docker_requirements_checker.is_docker_installed",
+        return_value=True,
+    )
+    mocker.patch(
+        "ostorlab.cli.docker_requirements_checker.is_docker_working", return_value=True
+    )
+    mocker.patch(
+        "ostorlab.cli.docker_requirements_checker.is_swarm_initialized",
+        return_value=True,
+    )
+    mocker.patch("docker.from_env")
+
+    mocker.patch(
+        "ostorlab.runtimes.local.runtime.LocalRuntime.can_run", return_value=True
+    )
+    mocker.patch("ostorlab.runtimes.local.runtime.LocalRuntime._create_network")
+    mocker.patch("ostorlab.runtimes.local.runtime.LocalRuntime._start_services")
+    mocker.patch("ostorlab.runtimes.local.runtime.LocalRuntime._start_pre_agents")
+    mocker.patch("ostorlab.runtimes.local.runtime.LocalRuntime._start_agents")
+    mocker.patch("ostorlab.runtimes.local.runtime.LocalRuntime._check_services_healthy")
+    mocker.patch(
+        "ostorlab.runtimes.local.runtime.LocalRuntime._check_agents_healthy",
+        return_value=True,
+    )
+    mocker.patch("ostorlab.runtimes.local.runtime.LocalRuntime._inject_assets")
+
+
 @pytest.fixture
 def network_scan(
     clean_db: None, mocker: plugin.MockerFixture, db_engine_path: str