From 35e923ed4938fde246921616614cbc2b728bf6fb Mon Sep 17 00:00:00 2001 From: ostorlab Date: Tue, 17 Dec 2024 11:30:54 +0100 Subject: [PATCH 1/4] Stop scan crashing when casting none to int --- src/ostorlab/runtimes/local/runtime.py | 8 +++- tests/runtimes/local/runtime_test.py | 66 ++++++++++++++++++++++++++ 2 files changed, 72 insertions(+), 2 deletions(-) diff --git a/src/ostorlab/runtimes/local/runtime.py b/src/ostorlab/runtimes/local/runtime.py index dece11a37..756672b89 100644 --- a/src/ostorlab/runtimes/local/runtime.py +++ b/src/ostorlab/runtimes/local/runtime.py @@ -309,9 +309,13 @@ def stop(self, scan_id: str) -> None: networks = self._docker_client.networks.list() for network in networks: network_labels = network.attrs["Labels"] + if network_labels is None: + continue + universe = network_labels.get("ostorlab.universe") if ( - network_labels is not None - and int(network_labels.get("ostorlab.universe")) == scan_id + universe is not None + and network_labels is not None + and int(universe) == scan_id ): logger.info("removing network %s", network_labels) stopped_network.append(network) diff --git a/tests/runtimes/local/runtime_test.py b/tests/runtimes/local/runtime_test.py index e3e2f1519..436ecedf8 100644 --- a/tests/runtimes/local/runtime_test.py +++ b/tests/runtimes/local/runtime_test.py @@ -6,6 +6,7 @@ import docker import pytest from docker.models import services as services_model +from docker.models import networks as networks_model from pytest_mock import plugin import ostorlab @@ -364,3 +365,68 @@ def testCheckServicesMethod_whenServicesAreStopped_shouldExit( assert exit_mock.call_count == 1 exit_with = exit_mock.call_args_list[0][0][0] assert exit_with == 0 + + +@pytest.mark.docker +def testRuntimeScanStop_whenUnrelatedNetworks_RemovesScanServiceWithoutCrash( + mocker, db_engine_path +): + """Unittest for the scan stop method when there are networks not related to the scan, the process shouldn't crash""" + mocker.patch.object(models, "ENGINE_URL", db_engine_path) + create_scan_db = models.Scan.create("test") + + def docker_services(): + """Method for mocking the services list response.""" + with models.Database() as session: + scan = session.query(models.Scan).first() + services = [ + { + "ID": "0099i5n1y3gycuekvksyqyxav", + "CreatedAt": "2021-12-27T13:37:02.795789947Z", + "Spec": {"Labels": {"ostorlab.universe": scan.id}}, + }, + { + "ID": "0099i5n1y3gycuekvksyqyxav", + "CreatedAt": "2021-12-27T13:37:02.795789947Z", + "Spec": {"Labels": {"ostorlab.universe": 9999}}, + }, + ] + + return [services_model.Service(attrs=service) for service in services] + + def docker_networks(): + """Method for mocking the services list response.""" + with models.Database() as session: + scan = session.query(models.Scan).first() + networks = [ + { + "ID": "0099i5n1y3gycuekvksyqyxav", + "CreatedAt": "2021-12-27T13:37:02.795789947Z", + "Labels": {"ostorlab.universe": scan.id}, + }, + { + "ID": "0099i5n1y3gycuekvksyqyxav", + "CreatedAt": "2021-12-27T13:37:02.795789947Z", + "Labels": {}, + }, + ] + + return [networks_model.Network(attrs=network) for network in networks] + + mocker.patch( + "docker.DockerClient.services", return_value=services_model.ServiceCollection() + ) + mocker.patch("docker.DockerClient.services.list", side_effect=docker_services) + mocker.patch( + "docker.models.networks.NetworkCollection.list", return_value=docker_networks() + ) + mocker.patch("docker.models.configs.ConfigCollection.list", return_value=[]) + docker_network_remove = mocker.patch("docker.models.networks.Network.remove") + docker_service_remove = mocker.patch( + "docker.models.services.Service.remove", return_value=None + ) + + local_runtime.LocalRuntime().stop(scan_id=create_scan_db.id) + + docker_service_remove.assert_called_once() + docker_network_remove.assert_called_once() From 9ef1082e2bf7f9197e234bb09f7e5990b814dcc0 Mon Sep 17 00:00:00 2001 From: Falcon <110392140+oussamaessaji@users.noreply.github.com> Date: Tue, 17 Dec 2024 11:54:20 +0100 Subject: [PATCH 2/4] Update tests/runtimes/local/runtime_test.py Co-authored-by: Hamza AIT BEN YISSA <144700714+benyissa@users.noreply.github.com> --- tests/runtimes/local/runtime_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/runtimes/local/runtime_test.py b/tests/runtimes/local/runtime_test.py index 436ecedf8..9c96bbf03 100644 --- a/tests/runtimes/local/runtime_test.py +++ b/tests/runtimes/local/runtime_test.py @@ -368,7 +368,7 @@ def testCheckServicesMethod_whenServicesAreStopped_shouldExit( @pytest.mark.docker -def testRuntimeScanStop_whenUnrelatedNetworks_RemovesScanServiceWithoutCrash( +def testRuntimeScanStop_whenUnrelatedNetworks_removesScanServiceWithoutCrash( mocker, db_engine_path ): """Unittest for the scan stop method when there are networks not related to the scan, the process shouldn't crash""" From c5c5385b7c028263e1a61166a19baceeea7bcbee Mon Sep 17 00:00:00 2001 From: ostorlab Date: Tue, 17 Dec 2024 14:10:39 +0100 Subject: [PATCH 3/4] Address comments --- src/ostorlab/runtimes/local/runtime.py | 6 +----- tests/runtimes/local/runtime_test.py | 8 ++++---- 2 files changed, 5 insertions(+), 9 deletions(-) diff --git a/src/ostorlab/runtimes/local/runtime.py b/src/ostorlab/runtimes/local/runtime.py index 756672b89..e8e90d92a 100644 --- a/src/ostorlab/runtimes/local/runtime.py +++ b/src/ostorlab/runtimes/local/runtime.py @@ -312,11 +312,7 @@ def stop(self, scan_id: str) -> None: if network_labels is None: continue universe = network_labels.get("ostorlab.universe") - if ( - universe is not None - and network_labels is not None - and int(universe) == scan_id - ): + if universe is not None and int(universe) == scan_id: logger.info("removing network %s", network_labels) stopped_network.append(network) network.remove() diff --git a/tests/runtimes/local/runtime_test.py b/tests/runtimes/local/runtime_test.py index 436ecedf8..a440c4837 100644 --- a/tests/runtimes/local/runtime_test.py +++ b/tests/runtimes/local/runtime_test.py @@ -375,7 +375,7 @@ def testRuntimeScanStop_whenUnrelatedNetworks_RemovesScanServiceWithoutCrash( mocker.patch.object(models, "ENGINE_URL", db_engine_path) create_scan_db = models.Scan.create("test") - def docker_services(): + def docker_services() -> list[services_model.Service]: """Method for mocking the services list response.""" with models.Database() as session: scan = session.query(models.Scan).first() @@ -394,7 +394,7 @@ def docker_services(): return [services_model.Service(attrs=service) for service in services] - def docker_networks(): + def docker_networks() -> list[networks_model.Network]: """Method for mocking the services list response.""" with models.Database() as session: scan = session.query(models.Scan).first() @@ -402,12 +402,12 @@ def docker_networks(): { "ID": "0099i5n1y3gycuekvksyqyxav", "CreatedAt": "2021-12-27T13:37:02.795789947Z", - "Labels": {"ostorlab.universe": scan.id}, + "Labels": {}, }, { "ID": "0099i5n1y3gycuekvksyqyxav", "CreatedAt": "2021-12-27T13:37:02.795789947Z", - "Labels": {}, + "Labels": {"ostorlab.universe": scan.id}, }, ] From 1e542dbc14f5a7153db9a45d3e515b7f94df4641 Mon Sep 17 00:00:00 2001 From: ostorlab Date: Tue, 17 Dec 2024 18:51:50 +0100 Subject: [PATCH 4/4] Address comments --- src/ostorlab/runtimes/local/runtime.py | 12 +++++++----- tests/runtimes/local/runtime_test.py | 2 +- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/src/ostorlab/runtimes/local/runtime.py b/src/ostorlab/runtimes/local/runtime.py index e8e90d92a..b0c8dc809 100644 --- a/src/ostorlab/runtimes/local/runtime.py +++ b/src/ostorlab/runtimes/local/runtime.py @@ -310,12 +310,14 @@ def stop(self, scan_id: str) -> None: for network in networks: network_labels = network.attrs["Labels"] if network_labels is None: + logger.debug("Skipping network with no labels") continue - universe = network_labels.get("ostorlab.universe") - if universe is not None and int(universe) == scan_id: - logger.info("removing network %s", network_labels) - stopped_network.append(network) - network.remove() + if isinstance(network_labels, dict): + universe = network_labels.get("ostorlab.universe") + if universe is not None and int(universe) == scan_id: + logger.info("removing network %s", network_labels) + stopped_network.append(network) + network.remove() configs = self._docker_client.configs.list() for config in configs: diff --git a/tests/runtimes/local/runtime_test.py b/tests/runtimes/local/runtime_test.py index fe2fe807e..c29cb77f8 100644 --- a/tests/runtimes/local/runtime_test.py +++ b/tests/runtimes/local/runtime_test.py @@ -369,7 +369,7 @@ def testCheckServicesMethod_whenServicesAreStopped_shouldExit( @pytest.mark.docker def testRuntimeScanStop_whenUnrelatedNetworks_removesScanServiceWithoutCrash( - mocker, db_engine_path + mocker: plugin.MockerFixture, db_engine_path: str ): """Unittest for the scan stop method when there are networks not related to the scan, the process shouldn't crash""" mocker.patch.object(models, "ENGINE_URL", db_engine_path)