Add run-on to Boefje Setup page (#4061)

Co-authored-by: Donny Peeters <[email protected]> Co-authored-by: Donny Peeters <[email protected]> Co-authored-by: ammar92 <[email protected]> Co-authored-by: Jan Klopper <[email protected]>
minvws · Feb 7, 2025 · 3888974 · 3888974
1 parent 3248a1c
commit 3888974
Show file tree

Hide file tree

Showing 21 changed files with 335 additions and 49 deletions.
diff --git a/boefjes/boefjes/katalogus/plugins.py b/boefjes/boefjes/katalogus/plugins.py
@@ -15,6 +15,7 @@
     get_plugins_filter_parameters,
 )
 from boefjes.models import FilterParameters, PaginationParameters, PluginType
+from boefjes.sql.db_models import RunOn
 from boefjes.sql.plugin_storage import get_plugin_storage
 from boefjes.storage.interfaces import DuplicatePlugin, IntegrityError, NotAllowed, PluginStorage
 
@@ -130,6 +131,7 @@ class BoefjeIn(BaseModel):
     boefje_schema: dict | None = None
     cron: str | None = None
     interval: int | None = None
+    run_on: list[RunOn] | None = None
     oci_image: str | None = None
     oci_arguments: list[str] = Field(default_factory=list)
 

diff --git a/boefjes/boefjes/migrations/versions/fc0295b38184_add_run_on_field_to_boefje.py b/boefjes/boefjes/migrations/versions/fc0295b38184_add_run_on_field_to_boefje.py
@@ -0,0 +1,33 @@
+"""Add run on field to boefje
+
+Revision ID: fc0295b38184
+Revises: 9f48560b0000
+Create Date: 2025-02-04 16:43:59.171960
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "fc0295b38184"
+down_revision = "9f48560b0000"
+branch_labels = None
+depends_on = None
+
+
+run_on = sa.Enum("create", "update", "create_update", name="run_on")
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    run_on.create(op.get_bind())
+    op.add_column("boefje", sa.Column("run_on", run_on, nullable=True))
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column("boefje", "run_on")
+    run_on.drop(op.get_bind(), checkfirst=False)
+    # ### end Alembic commands ###
diff --git a/boefjes/boefjes/models.py b/boefjes/boefjes/models.py
@@ -1,5 +1,6 @@
 import datetime
 from enum import Enum
+from functools import total_ordering
 from typing import Literal
 
 from croniter import croniter
@@ -8,6 +9,18 @@
 from pydantic import BaseModel, Field, field_validator
 
 
+# This makes the RunOn sortable when in a list. This is convenient for e.g. the RunOnDB.from_run_ons method, that now
+# does not have to take the ordering of a boefje.run_on into account in its match statement. This is especially handy
+# once we introduce more RunOn values such as DELETE.
+@total_ordering
+class RunOn(Enum):
+    CREATE = "create"
+    UPDATE = "update"
+
+    def __lt__(self, other):
+        return self.value < other.value
+
+
 class Organisation(BaseModel):
     id: str
     name: str
@@ -34,6 +47,7 @@ class Boefje(Plugin):
     boefje_schema: dict | None = None
     cron: str | None = None
     interval: int | None = None
+    run_on: list[RunOn] | None = None
     runnable_hash: str | None = None
     oci_image: str | None = None
     oci_arguments: list[str] = Field(default_factory=list)

diff --git a/boefjes/boefjes/plugins/kat_export_http/boefje.json b/boefjes/boefjes/plugins/kat_export_http/boefje.json
@@ -5,8 +5,5 @@
   "consumes": [],
   "scan_level": 4,
   "oci_image": "ghcr.io/minvws/openkat/export-http:latest",
-  "run_on": [
-    "create",
-    "update"
-  ]
+  "run_on": ["create", "update"]
 }
diff --git a/boefjes/boefjes/sql/db_models.py b/boefjes/boefjes/sql/db_models.py
@@ -3,6 +3,7 @@
 from sqlalchemy import Boolean, Column, ForeignKey, Integer, String, UniqueConstraint, types
 from sqlalchemy.orm import relationship
 
+from boefjes.models import RunOn
 from boefjes.sql.db import SQL_BASE
 
 
@@ -14,6 +15,36 @@ class ScanLevel(Enum):
     L4 = 4
 
 
+class RunOnDB(Enum):
+    CREATE = "create"
+    UPDATE = "update"
+    CREATE_UPDATE = "create_update"
+
+    @classmethod
+    def from_run_ons(cls, run_ons: list[RunOn] | None):
+        if run_ons is None:
+            return None
+
+        match sorted(run_ons):
+            case [RunOn.CREATE]:
+                return cls.CREATE
+            case [RunOn.UPDATE]:
+                return cls.UPDATE
+            case [RunOn.CREATE, RunOn.UPDATE]:
+                return cls.CREATE_UPDATE
+            case _:
+                return None
+
+    def to_run_ons(self) -> list[RunOn]:
+        match self:
+            case RunOnDB.CREATE:
+                return [RunOn.CREATE]
+            case RunOnDB.UPDATE:
+                return [RunOn.UPDATE]
+            case RunOnDB.CREATE_UPDATE:
+                return [RunOn.CREATE, RunOn.UPDATE]
+
+
 class OrganisationInDB(SQL_BASE):
     __tablename__ = "organisation"
 
@@ -72,6 +103,7 @@ class BoefjeInDB(SQL_BASE):
     schema = Column(types.JSON(), nullable=True)
     cron = Column(types.String(length=128), nullable=True)
     interval = Column(types.Integer, nullable=True)
+    run_on = Column(types.Enum(*[x.value for x in RunOnDB], name="run_on"), nullable=True)
 
     # Image specifications
     oci_image = Column(types.String(length=256), nullable=True)

diff --git a/boefjes/boefjes/sql/plugin_storage.py b/boefjes/boefjes/sql/plugin_storage.py
@@ -6,7 +6,7 @@
 from boefjes.config import Settings, settings
 from boefjes.models import Boefje, Normalizer, PluginType
 from boefjes.sql.db import ObjectNotFoundException, session_managed_iterator
-from boefjes.sql.db_models import BoefjeInDB, NormalizerInDB
+from boefjes.sql.db_models import BoefjeInDB, NormalizerInDB, RunOnDB
 from boefjes.sql.session import SessionMixin
 from boefjes.storage.interfaces import NotAllowed, PluginNotFound, PluginStorage
 
@@ -98,6 +98,7 @@ def _db_normalizer_instance_by_id(self, normalizer_id: str) -> NormalizerInDB:
 
     @staticmethod
     def to_boefje_in_db(boefje: Boefje, pk: int | None = None) -> BoefjeInDB:
+        run_on_db = RunOnDB.from_run_ons(boefje.run_on)
         boefje = BoefjeInDB(
             plugin_id=boefje.id,
             created=boefje.created,
@@ -109,6 +110,7 @@ def to_boefje_in_db(boefje: Boefje, pk: int | None = None) -> BoefjeInDB:
             schema=boefje.boefje_schema,
             cron=boefje.cron,
             interval=boefje.interval,
+            run_on=run_on_db.value if run_on_db is not None else None,
             oci_image=boefje.oci_image,
             oci_arguments=boefje.oci_arguments,
             version=boefje.version,
@@ -152,6 +154,7 @@ def to_boefje(boefje_in_db: BoefjeInDB) -> Boefje:
             boefje_schema=boefje_in_db.schema,
             cron=boefje_in_db.cron,
             interval=boefje_in_db.interval,
+            run_on=RunOnDB(boefje_in_db.run_on).to_run_ons() if boefje_in_db.run_on else None,
             oci_image=boefje_in_db.oci_image,
             oci_arguments=boefje_in_db.oci_arguments,
             version=boefje_in_db.version,

diff --git a/boefjes/tests/integration/test_api.py b/boefjes/tests/integration/test_api.py
@@ -83,6 +83,22 @@ def test_enable_boefje(test_client, organisation, second_organisation):
     assert response.json()["enabled"] is False
 
 
+def test_run_on(test_client, organisation, second_organisation):
+    test_client.patch(f"/v1/organisations/{organisation.id}/plugins/export-to-http-api", json={"enabled": True})
+
+    response = test_client.get(f"/v1/organisations/{organisation.id}/plugins/export-to-http-api")
+    assert response.json()["enabled"] is True
+    assert response.json()["run_on"] == ["create", "update"]
+
+    boefje = Boefje(id="test_run_on", name="Run On", static=False, run_on=["create"])
+    response = test_client.post(f"/v1/organisations/{organisation.id}/plugins", content=boefje.model_dump_json())
+    assert response.status_code == 201
+
+    response = test_client.get(f"/v1/organisations/{organisation.id}/plugins/test_run_on")
+    assert response.json()["enabled"] is False
+    assert response.json()["run_on"] == [x.value for x in boefje.run_on]
+
+
 def test_cannot_add_static_plugin_with_duplicate_name(test_client, organisation):
     boefje = Boefje(id="test_plugin", name="DNS records", static=False)
     response = test_client.post(f"/v1/organisations/{organisation.id}/plugins", content=boefje.model_dump_json())

diff --git a/boefjes/tests/test_models.py b/boefjes/tests/test_models.py
@@ -0,0 +1,15 @@
+from boefjes.models import RunOn
+from boefjes.sql.db_models import RunOnDB
+
+
+def test_run_on():
+    assert RunOnDB.from_run_ons([RunOn.CREATE]) == RunOnDB.CREATE
+    assert RunOnDB.from_run_ons([RunOn.UPDATE]) == RunOnDB.UPDATE
+    assert RunOnDB.from_run_ons([RunOn.CREATE, RunOn.UPDATE]) == RunOnDB.CREATE_UPDATE
+    assert RunOnDB.from_run_ons([RunOn.UPDATE, RunOn.CREATE]) == RunOnDB.CREATE_UPDATE
+    assert RunOnDB.from_run_ons([1]) is None
+    assert RunOnDB.from_run_ons([]) is None
+
+    assert RunOnDB.CREATE.to_run_ons() == [RunOn.CREATE]
+    assert RunOnDB.UPDATE.to_run_ons() == [RunOn.UPDATE]
+    assert RunOnDB.CREATE_UPDATE.to_run_ons() == [RunOn.CREATE, RunOn.UPDATE]
diff --git a/mula/scheduler/models/ooi.py b/mula/scheduler/models/ooi.py
@@ -9,6 +9,11 @@ class MutationOperationType(Enum):
     DELETE = "delete"
 
 
+class RunOn(Enum):
+    CREATE = MutationOperationType.CREATE.value
+    UPDATE = MutationOperationType.UPDATE.value
+
+
 class ScanProfile(BaseModel):
     level: int
     reference: str

diff --git a/mula/scheduler/models/plugin.py b/mula/scheduler/models/plugin.py
@@ -2,6 +2,8 @@
 
 from pydantic import BaseModel
 
+from scheduler.models.ooi import RunOn
+
 
 class Plugin(BaseModel):
     id: str
@@ -19,4 +21,4 @@ class Plugin(BaseModel):
     produces: list[str]
     cron: str | None = None
     interval: int | None = None
-    run_on: list[str] | None = None
+    run_on: list[RunOn] | None = None
diff --git a/mula/scheduler/schedulers/schedulers/boefje.py b/mula/scheduler/schedulers/schedulers/boefje.py
@@ -21,6 +21,7 @@
     Task,
     TaskStatus,
 )
+from scheduler.models.ooi import RunOn
 from scheduler.schedulers import Scheduler
 from scheduler.schedulers.queue import PriorityQueue, QueueFullError
 from scheduler.schedulers.rankers import BoefjeRanker
@@ -214,9 +215,9 @@ def push_tasks_for_scan_profile_mutations(self, body: bytes) -> None:
                     create_schedule = False
                     run_task = False
                     if mutation.operation == MutationOperationType.CREATE:
-                        run_task = "create" in boefje.run_on
+                        run_task = RunOn.CREATE in boefje.run_on
                     elif mutation.operation == MutationOperationType.UPDATE:
-                        run_task = "update" in boefje.run_on
+                        run_task = RunOn.UPDATE in boefje.run_on
 
                 if not run_task:
                     self.logger.debug(

diff --git a/mula/tests/factories/plugin.py b/mula/tests/factories/plugin.py
@@ -1,5 +1,6 @@
 from factory import Factory, LazyFunction, Sequence, fuzzy
 from scheduler.models import Plugin
+from scheduler.models.ooi import RunOn
 
 
 class PluginFactory(Factory):
@@ -13,4 +14,4 @@ class Meta:
     enabled: bool = True
     cron: str | None = None
     interval: int | None = None
-    run_on: list[str] | None = None
+    run_on: RunOn | None = None
diff --git a/mula/tests/integration/test_boefje_scheduler.py b/mula/tests/integration/test_boefje_scheduler.py
@@ -4,6 +4,7 @@
 from unittest import mock
 
 from scheduler import clients, config, models, schedulers, storage
+from scheduler.models.ooi import RunOn
 from scheduler.storage import stores
 from structlog.testing import capture_logs
 
@@ -1283,7 +1284,7 @@ def test_push_tasks_for_scan_profile_mutations_op_create_run_on_create(self):
         # Arrange
         scan_profile = ScanProfileFactory(level=0)
         ooi = OOIFactory(scan_profile=scan_profile)
-        boefje = PluginFactory(scan_level=0, consumes=[ooi.object_type], run_on=["create"])
+        boefje = PluginFactory(scan_level=0, consumes=[ooi.object_type], run_on=[RunOn.CREATE])
         mutation = models.ScanProfileMutation(
             operation=models.MutationOperationType.CREATE, primary_key=ooi.primary_key, value=ooi
         ).model_dump_json()
@@ -1320,7 +1321,7 @@ def test_push_tasks_for_scan_profile_mutations_op_create_run_on_create_update(se
         # Arrange
         scan_profile = ScanProfileFactory(level=0)
         ooi = OOIFactory(scan_profile=scan_profile)
-        boefje = PluginFactory(scan_level=0, consumes=[ooi.object_type], run_on=["create", "update"])
+        boefje = PluginFactory(scan_level=0, consumes=[ooi.object_type], run_on=[RunOn.CREATE, RunOn.UPDATE])
         mutation = models.ScanProfileMutation(
             operation=models.MutationOperationType.CREATE, primary_key=ooi.primary_key, value=ooi
         ).model_dump_json()
@@ -1357,7 +1358,7 @@ def test_push_tasks_for_scan_profile_mutations_op_create_run_on_update(self):
         # Arrange
         scan_profile = ScanProfileFactory(level=0)
         ooi = OOIFactory(scan_profile=scan_profile)
-        boefje = PluginFactory(scan_level=0, consumes=[ooi.object_type], run_on=["update"])
+        boefje = PluginFactory(scan_level=0, consumes=[ooi.object_type], run_on=[RunOn.UPDATE])
         mutation = models.ScanProfileMutation(
             operation=models.MutationOperationType.CREATE, primary_key=ooi.primary_key, value=ooi
         ).model_dump_json()
@@ -1381,7 +1382,7 @@ def test_push_tasks_for_scan_profile_mutations_op_create_run_on_none(self):
         # Arrange
         scan_profile = ScanProfileFactory(level=0)
         ooi = OOIFactory(scan_profile=scan_profile)
-        boefje = PluginFactory(scan_level=0, consumes=[ooi.object_type], run_on=[])
+        boefje = PluginFactory(scan_level=0, consumes=[ooi.object_type], run_on=None)
         mutation = models.ScanProfileMutation(
             operation=models.MutationOperationType.CREATE, primary_key=ooi.primary_key, value=ooi
         ).model_dump_json()
@@ -1419,7 +1420,7 @@ def test_push_tasks_for_scan_profile_mutations_op_update_run_on_create(self):
         # Arrange
         scan_profile = ScanProfileFactory(level=0)
         ooi = OOIFactory(scan_profile=scan_profile)
-        boefje = PluginFactory(scan_level=0, consumes=[ooi.object_type], run_on=["create"])
+        boefje = PluginFactory(scan_level=0, consumes=[ooi.object_type], run_on=[RunOn.CREATE])
         mutation = models.ScanProfileMutation(
             operation=models.MutationOperationType.UPDATE, primary_key=ooi.primary_key, value=ooi
         ).model_dump_json()
@@ -1443,7 +1444,7 @@ def test_push_tasks_scan_profile_mutations_op_update_run_on_create_update(self):
         # Arrange
         scan_profile = ScanProfileFactory(level=0)
         ooi = OOIFactory(scan_profile=scan_profile)
-        boefje = PluginFactory(scan_level=0, consumes=[ooi.object_type], run_on=["create", "update"])
+        boefje = PluginFactory(scan_level=0, consumes=[ooi.object_type], run_on=[RunOn.CREATE, RunOn.UPDATE])
         mutation = models.ScanProfileMutation(
             operation=models.MutationOperationType.UPDATE, primary_key=ooi.primary_key, value=ooi
         ).model_dump_json()
@@ -1480,7 +1481,7 @@ def test_push_tasks_scan_profile_mutations_op_update_run_on_update(self):
         # Arrange
         scan_profile = ScanProfileFactory(level=0)
         ooi = OOIFactory(scan_profile=scan_profile)
-        boefje = PluginFactory(scan_level=0, consumes=[ooi.object_type], run_on=["update"])
+        boefje = PluginFactory(scan_level=0, consumes=[ooi.object_type], run_on=[RunOn.UPDATE])
         mutation = models.ScanProfileMutation(
             operation=models.MutationOperationType.UPDATE, primary_key=ooi.primary_key, value=ooi
         ).model_dump_json()
@@ -1517,7 +1518,7 @@ def test_push_tasks_scan_profile_mutations_op_update_run_on_none(self):
         # Arrange
         scan_profile = ScanProfileFactory(level=0)
         ooi = OOIFactory(scan_profile=scan_profile)
-        boefje = PluginFactory(scan_level=0, consumes=[ooi.object_type], run_on=[])
+        boefje = PluginFactory(scan_level=0, consumes=[ooi.object_type], run_on=None)
         mutation = models.ScanProfileMutation(
             operation=models.MutationOperationType.UPDATE, primary_key=ooi.primary_key, value=ooi
         ).model_dump_json()