Refactor KATalogus client in Rocky

octopoes/octopoes/connector/katalogus.py

@@ -1,7 +1,7 @@
 import httpx
 
 
-class KATalogusClientV1:
+class KATalogusClient:
     def __init__(self, base_uri: str):
         self.base_uri = f"{base_uri.rstrip('/')}/v1"
 

octopoes/octopoes/tasks/tasks.py

@@ -12,7 +12,7 @@
 from pydantic import TypeAdapter
 
 from octopoes.config.settings import QUEUE_NAME_OCTOPOES, Settings
-from octopoes.connector.katalogus import KATalogusClientV1
+from octopoes.connector.katalogus import KATalogusClient
 from octopoes.core.app import bootstrap_octopoes, close_rabbit_channel, get_xtdb_client
 from octopoes.events.events import DBEvent, DBEventType
 from octopoes.events.manager import get_rabbit_channel
@@ -80,7 +80,7 @@ def handle_event(event: dict):
 @app.task(queue=QUEUE_NAME_OCTOPOES)
 def schedule_scan_profile_recalculations():
     try:
-        orgs = KATalogusClientV1(str(settings.katalogus_api)).get_organisations()
+        orgs = KATalogusClient(str(settings.katalogus_api)).get_organisations()
     except HTTPError:
         logger.exception("Failed getting organizations")
         raise

rocky/katalogus/health.py

@@ -1,15 +1,15 @@
 import structlog
 from httpx import HTTPError
 
-from katalogus.client import get_katalogus
+from katalogus.client import get_katalogus_client
 from rocky.health import ServiceHealth
 
 logger = structlog.get_logger(__name__)
 
 
 def get_katalogus_health() -> ServiceHealth:
     try:
-        katalogus_client = get_katalogus("")  # For the health endpoint the organization has no effect
+        katalogus_client = get_katalogus_client()  # For the health endpoint the organization has no effect
         katalogus_health = katalogus_client.health()
     except HTTPError:
         logger.exception("Error while retrieving KATalogus health state")

rocky/katalogus/utils.py

@@ -5,4 +5,6 @@
 
 
 def get_enabled_boefjes_for_ooi_class(ooi_class: type[OOI], organization: Organization) -> list[Boefje]:
-    return [boefje for boefje in get_katalogus(organization.code).get_enabled_boefjes() if ooi_class in boefje.consumes]
+    katalogus = get_katalogus(organization.code)
+
+    return [boefje for boefje in katalogus.get_enabled_boefjes() if ooi_class in boefje.consumes]

rocky/katalogus/views/mixins.py

@@ -8,13 +8,13 @@
 from httpx import HTTPError, HTTPStatusError
 from rest_framework.status import HTTP_404_NOT_FOUND
 
-from katalogus.client import KATalogusClientV1, Plugin, get_katalogus
+from katalogus.client import Boefje, KATalogus, Plugin, get_katalogus
 
 logger = structlog.get_logger(__name__)
 
 
 class SinglePluginView(OrganizationView):
-    katalogus_client: KATalogusClientV1
+    katalogus_client: KATalogus
     plugin: Plugin
 
     def setup(self, request, *args, plugin_id: str, **kwargs):
@@ -26,7 +26,7 @@ def setup(self, request, *args, plugin_id: str, **kwargs):
 
         try:
             self.plugin = self.katalogus_client.get_plugin(plugin_id)
-            self.plugin_schema = self.katalogus_client.get_plugin_schema(plugin_id)
+            self.plugin_schema = self.plugin.boefje_schema if isinstance(self.plugin, Boefje) else None
         except HTTPError as exc:
             if isinstance(exc, HTTPStatusError) and exc.response.status_code == HTTP_404_NOT_FOUND:
                 raise Http404(f"Plugin {plugin_id} not found.")

rocky/onboarding/views.py

@@ -268,8 +268,9 @@ class OnboardingSetupScanSelectPluginsView(
 
     def get_plugins(self) -> dict[str, list[Plugin]]:
         all_plugins = {}
+        katalogus = get_katalogus(self.organization.code)
         for required_optional, plugin_ids in self.plugins.items():
-            plugins = [get_katalogus(self.organization.code).get_plugin(plugin_id) for plugin_id in plugin_ids]  # type: ignore
+            plugins = katalogus.get_plugins(ids=plugin_ids)  # type: ignore
             all_plugins[required_optional] = plugins
 
         return all_plugins

rocky/reports/views/base.py

@@ -17,7 +17,7 @@
 from django.utils.translation import gettext_lazy as _
 from django.views.generic import TemplateView
 from django_weasyprint import WeasyTemplateResponseMixin
-from katalogus.client import Boefje, KATalogusClientV1, KATalogusError, Plugin, get_katalogus
+from katalogus.client import Boefje, KATalogus, KATalogusError, Plugin, get_katalogus
 from pydantic import RootModel, TypeAdapter
 from tools.ooi_helpers import create_ooi
 from tools.view_helpers import BreadcrumbsMixin, PostRedirect, url_with_querystring
@@ -120,7 +120,7 @@ def get(self, request: HttpRequest, *args: Any, **kwargs: Any) -> HttpResponse:
         return redirect(reverse("report_history", kwargs=self.get_kwargs()))
 
 
-def hydrate_plugins(report_types: list[type["BaseReport"]], katalogus: KATalogusClientV1) -> dict[str, list[Plugin]]:
+def hydrate_plugins(report_types: list[type["BaseReport"]], katalogus: KATalogus) -> dict[str, list[Plugin]]:
     plugins: dict[str, list[Plugin]] = {"required": [], "optional": []}
     merged_plugins = report_plugins_union(report_types)
 

rocky/tests/conftest.py

@@ -90,7 +90,7 @@ def create_user(django_user_model, email, password, name, device_name, superuser
 
 
 def create_organization(name, organization_code):
-    katalogus_client = "katalogus.client.KATalogusClientV1"
+    katalogus_client = "katalogus.client.KATalogusClient"
     octopoes_node = "tools.models.OctopoesAPIConnector"
     with patch(katalogus_client), patch(octopoes_node):
         return Organization.objects.create(name=name, code=organization_code)
@@ -308,7 +308,7 @@ def blocked_member(django_user_model, organization):
 
 @pytest.fixture
 def mock_models_katalogus(mocker):
-    return mocker.patch("tools.models.get_katalogus")
+    return mocker.patch("tools.models.get_katalogus_client")
 
 
 @pytest.fixture
@@ -1015,7 +1015,7 @@ def tree_data_tls_findings_and_suites():
 
 
 @pytest.fixture
-def plugin_details():
+def plugin_details(plugin_schema):
     return parse_plugin(
         {
             "id": "test-boefje",
@@ -1027,15 +1027,15 @@ def plugin_details():
             "consumes": ["Network"],
             "produces": ["Network"],
             "enabled": True,
-            "boefje_schema": {},
+            "boefje_schema": plugin_schema,
             "oci_image": None,
             "oci_arguments": ["-test", "-arg"],
         }
     )
 
 
 @pytest.fixture
-def plugin_details_with_container():
+def plugin_details_with_container(plugin_schema):
     return parse_plugin(
         {
             "id": "test-boefje",
@@ -1047,7 +1047,7 @@ def plugin_details_with_container():
             "consumes": ["Network"],
             "produces": ["Network"],
             "enabled": True,
-            "boefje_schema": {},
+            "boefje_schema": plugin_schema,
             "oci_image": "ghcr.io/test/image:123",
             "oci_arguments": ["-test", "-arg"],
         }
@@ -1353,7 +1353,7 @@ def mock_mixins_katalogus(mocker):
 
 @pytest.fixture()
 def mock_katalogus_client(mocker):
-    return mocker.patch("katalogus.client.KATalogusClientV1")
+    return mocker.patch("katalogus.client.KATalogusClient")
 
 
 @pytest.fixture

rocky/tests/integration/conftest.py

@@ -27,7 +27,7 @@ def valid_time():
 
 @pytest.fixture
 def katalogus_mock(mocker):
-    katalogus = mocker.patch("katalogus.client.KATalogusClientV1")
+    katalogus = mocker.patch("katalogus.client.KATalogusClient")
     katalogus().health.return_value = ServiceHealth(service="katalogus", healthy=True)
 
     return katalogus

rocky/tests/katalogus/test_katalogus.py

@@ -1,7 +1,7 @@
 import pytest
 from django.core.exceptions import PermissionDenied
 from django.urls import resolve
-from katalogus.client import KATalogusClientV1, parse_plugin, valid_organization_code, valid_plugin_id
+from katalogus.client import KATalogus, KATalogusClient, parse_plugin, valid_organization_code, valid_plugin_id
 from katalogus.views.katalogus import AboutPluginsView, BoefjeListView, KATalogusView, NormalizerListView
 from katalogus.views.katalogus_settings import ConfirmCloneSettingsView, KATalogusSettingsView
 from katalogus.views.plugin_enable_disable import PluginEnableDisableView
@@ -148,7 +148,7 @@ def test_katalogus_plugin_listing_no_enable_disable_perm(rf, client_member, mock
 
 def test_katalogus_settings_one_organization(redteam_member, rf, mocker):
     # Mock katalogus calls: return right boefjes and settings
-    mock_katalogus = mocker.patch("katalogus.client.KATalogusClientV1")
+    mock_katalogus = mocker.patch("katalogus.client.KATalogusClient")
     boefjes_data = get_boefjes_data()
     mock_katalogus().get_boefjes.return_value = [parse_plugin(b) for b in boefjes_data if b["type"] == "boefje"]
     mock_katalogus().get_plugin_settings.return_value = {"BINARYEDGE_API": "test", "Second": "value"}
@@ -169,7 +169,7 @@ def test_katalogus_settings_one_organization(redteam_member, rf, mocker):
 
 def test_katalogus_settings_list_multiple_organization(redteam_member, organization_b, rf, mocker):
     # Mock katalogus calls: return right boefjes and settings
-    mock_katalogus = mocker.patch("katalogus.client.KATalogusClientV1")
+    mock_katalogus = mocker.patch("katalogus.client.KATalogusClient")
     boefjes_data = get_boefjes_data()
     mock_katalogus().get_boefjes.return_value = [parse_plugin(b) for b in boefjes_data if b["type"] == "boefje"]
     mock_katalogus().get_plugin_settings.return_value = {"BINARYEDGE_API": "test"}
@@ -193,7 +193,7 @@ def test_katalogus_settings_list_multiple_organization(redteam_member, organizat
 
 
 def test_katalogus_confirm_clone_settings(redteam_member, organization_b, rf, mock_models_octopoes, mocker):
-    mocker.patch("katalogus.client.KATalogusClientV1")
+    mocker.patch("katalogus.client.KATalogusClient")
 
     create_member(redteam_member.user, organization_b)
 
@@ -214,7 +214,7 @@ def test_katalogus_confirm_clone_settings(redteam_member, organization_b, rf, mo
 
 def test_katalogus_clone_settings(redteam_member, organization_b, rf, mocker, mock_models_octopoes):
     # Mock katalogus calls: return right boefjes and settings
-    mock_katalogus = mocker.patch("katalogus.client.KATalogusClientV1")
+    mock_katalogus = mocker.patch("katalogus.client.KATalogusClient")
 
     create_member(redteam_member.user, organization_b)
 
@@ -224,14 +224,16 @@ def test_katalogus_clone_settings(redteam_member, organization_b, rf, mocker, mo
     )
     assert response.status_code == 302
 
-    mock_katalogus().clone_all_configuration_to_organization.assert_called_once_with(organization_b.code)
+    mock_katalogus().clone_all_configuration_to_organization.assert_called_once_with(
+        redteam_member.organization.code, organization_b.code
+    )
 
 
 def test_katalogus_clone_settings_not_accessible_without_perms(
     client_member, organization_b, rf, mocker, mock_models_octopoes
 ):
     # Mock katalogus calls: return right boefjes and settings
-    mocker.patch("katalogus.client.KATalogusClientV1")
+    mocker.patch("katalogus.client.KATalogusClient")
 
     create_member(client_member.user, organization_b)
 
@@ -242,33 +244,30 @@ def test_katalogus_clone_settings_not_accessible_without_perms(
         )
 
 
-def test_katalogus_client_organization_not_exists(mocker):
-    mock_requests = mocker.patch("katalogus.client.httpx")
-    mock_requests.Client().get().status_code = 404
-
-    client = KATalogusClientV1("test", "test")
+def test_katalogus_client_organization_not_exists(httpx_mock):
+    httpx_mock.add_response(status_code=404)
 
-    assert client.organization_exists() is False
+    assert KATalogusClient("http://test").organization_exists("test") is False
 
 
 def test_katalogus_client_organization_exists(mocker):
     mock_requests = mocker.patch("katalogus.client.httpx")
     mock_requests.Client().get().status_code = 200
 
-    client = KATalogusClientV1("test", "test")
+    client = KATalogusClient("test")
 
-    assert client.organization_exists() is True
+    assert client.organization_exists("test") is True
 
 
 def test_katalogus_client_invalid_organization():
     with pytest.raises(ValueError):
-        KATalogusClientV1("test", "test$$$1123")
+        KATalogus(KATalogusClient("test"), "test$$$1123")
 
 
 def test_katalogus_client(httpx_mock):
     httpx_mock.add_response(json={"service": "test", "healthy": False, "version": None, "additional": 2, "results": []})
 
-    client = KATalogusClientV1("http://test", "test")
+    client = KATalogusClient("http://test")
 
     assert isinstance(client.health(), ServiceHealth)
     assert client.health().service == "test"

rocky/tests/katalogus/test_katalogus_boefje_setup.py

@@ -22,7 +22,7 @@ def test_boefje_setup(rf, superuser_member):
 
 
 def test_boefje_variant_setup(rf, superuser_member, boefje_dns_records, mocker):
-    katalogus_mocker = mocker.patch("katalogus.client.KATalogusClientV1")()
+    katalogus_mocker = mocker.patch("katalogus.client.KATalogusClient")()
     katalogus_mocker.get_plugin.return_value = boefje_dns_records
     request = setup_request(rf.get("boefje_variant_setup"), superuser_member.user)
     response = AddBoefjeVariantView.as_view()(
@@ -43,7 +43,7 @@ def test_boefje_variant_setup(rf, superuser_member, boefje_dns_records, mocker):
 
 
 def test_edit_boefje_view(rf, superuser_member, boefje_dns_records, mocker):
-    katalogus_mocker = mocker.patch("katalogus.client.KATalogusClientV1")()
+    katalogus_mocker = mocker.patch("katalogus.client.KATalogusClient")()
     katalogus_mocker.get_plugin.return_value = boefje_dns_records
     request = setup_request(rf.get("edit_boefje"), superuser_member.user)
     response = EditBoefjeView.as_view()(

rocky/tests/katalogus/test_katalogus_plugin_add.py

@@ -154,8 +154,8 @@ def test_plugin_settings_add_error_message_about_integer_too_big(
 
 def test_plugin_single_settings_add_view_no_schema(rf, superuser_member, mock_mixins_katalogus, plugin_details):
     mock_katalogus = mock_mixins_katalogus()
+    plugin_details.boefje_schema = None
     mock_katalogus.get_plugin.return_value = plugin_details
-    mock_katalogus.get_plugin_schema.return_value = None
     mock_katalogus.get_plugin_settings.return_value = None
 
     request = setup_request(rf.post("plugin_settings_add", data={"boefje_id": 123}), superuser_member.user)

rocky/tests/katalogus/test_katalogus_plugin_detail.py

@@ -16,7 +16,7 @@ def test_plugin_detail_view(
     mocker,
 ):
     mock_mixins_katalogus().get_plugin.return_value = plugin_details
-    katalogus_mocker = mocker.patch("katalogus.client.KATalogusClientV1")()
+    katalogus_mocker = mocker.patch("katalogus.client.KATalogusClient")()
     katalogus_mocker.get_plugins.return_value = [boefje_dns_records, boefje_nmap_tcp]
 
     request = setup_request(rf.get("boefje_detail"), superuser_member.user)
@@ -46,7 +46,7 @@ def test_plugin_detail_view_with_container_image(
     mock_scheduler_client_task_list,
 ):
     mock_mixins_katalogus().get_plugin.return_value = plugin_details_with_container
-    katalogus_mocker = mocker.patch("katalogus.client.KATalogusClientV1")()
+    katalogus_mocker = mocker.patch("katalogus.client.KATalogusClient")()
     katalogus_mocker.get_plugins.return_value = [boefje_dns_records, boefje_nmap_tcp]
 
     request = setup_request(rf.get("boefje_detail"), superuser_member.user)
@@ -84,7 +84,7 @@ def test_plugin_detail_view_no_consumes(
 ):
     plugin_details.consumes = []
     mock_mixins_katalogus().get_plugin.return_value = plugin_details
-    katalogus_mocker = mocker.patch("katalogus.client.KATalogusClientV1")()
+    katalogus_mocker = mocker.patch("katalogus.client.KATalogusClient")()
     katalogus_mocker.get_plugins.return_value = [boefje_dns_records, boefje_nmap_tcp]
 
     request = setup_request(rf.get("boefje_detail"), superuser_member.user)