Allow MuteFindings to expire by a user specified datetime (#3343)

Co-authored-by: Jan Klopper <[email protected]> Co-authored-by: Jeroen Dekkers <[email protected]> Co-authored-by: ammar92 <[email protected]>
minvws · Aug 12, 2024 · 0cf9a5a · 0cf9a5a
1 parent 71e74e2
commit 0cf9a5a
Show file tree

Hide file tree

Showing 17 changed files with 72 additions and 19 deletions.
diff --git a/boefjes/boefjes/job_models.py b/boefjes/boefjes/job_models.py
@@ -1,4 +1,4 @@
-from datetime import timedelta
+from datetime import datetime, timedelta
 from typing import Annotated, Literal, TypeAlias
 from uuid import UUID
 
@@ -85,6 +85,7 @@ class NormalizerObservation(BaseModel):
 class NormalizerDeclaration(BaseModel):
     type: Literal["declaration"] = "declaration"
     ooi: OOIType
+    end_valid_time: datetime | None = None
 
 
 class NormalizerAffirmation(BaseModel):

diff --git a/boefjes/boefjes/plugins/kat_manual/single_ooi/normalize.py b/boefjes/boefjes/plugins/kat_manual/single_ooi/normalize.py
@@ -6,4 +6,5 @@
 
 def run(input_ooi: dict, raw: bytes) -> Iterable[NormalizerOutput]:
     for declaration in json.loads(raw.decode()):
-        yield NormalizerDeclaration(ooi=declaration["ooi"])
+        end_valid_time = declaration.pop("end_valid_time", None)
+        yield NormalizerDeclaration(ooi=declaration["ooi"], end_valid_time=end_valid_time)
diff --git a/octopoes/octopoes/api/models.py b/octopoes/octopoes/api/models.py
@@ -41,6 +41,7 @@ class Declaration(BaseModel):
 
     ooi: OOIType
     valid_time: datetime
+    end_valid_time: datetime | None = None
     method: str | None = None
     source_method: str | None = None
     task_id: uuid.UUID | None = None
@@ -79,6 +80,7 @@ class ValidatedDeclaration(BaseModel):
 
     ooi: ValidatedOOIType
     valid_time: AwareDatetime
+    end_valid_time: AwareDatetime | None = None
     method: str | None = "manual"
     source_method: str | None = None
     task_id: uuid.UUID | None = Field(default_factory=uuid.uuid4)

diff --git a/octopoes/octopoes/api/router.py b/octopoes/octopoes/api/router.py
@@ -354,7 +354,7 @@ def save_declaration(
         result=[declaration.ooi.reference],
         task_id=declaration.task_id if declaration.task_id else uuid.uuid4(),
     )
-    octopoes.save_origin(origin, [declaration.ooi], declaration.valid_time)
+    octopoes.save_origin(origin, [declaration.ooi], declaration.valid_time, declaration.end_valid_time)
     octopoes.commit()
 
 

diff --git a/octopoes/octopoes/core/service.py b/octopoes/octopoes/core/service.py
@@ -152,7 +152,9 @@ def _delete_ooi(self, reference: Reference, valid_time: datetime) -> None:
         if not referencing_origins:
             self.ooi_repository.delete(reference, valid_time)
 
-    def save_origin(self, origin: Origin, oois: list[OOI], valid_time: datetime) -> None:
+    def save_origin(
+        self, origin: Origin, oois: list[OOI], valid_time: datetime, end_valid_time: datetime | None = None
+    ) -> None:
         origin.result = [ooi.reference for ooi in oois]
 
         # When an Origin is saved while the source OOI does not exist, reject saving the results
@@ -166,7 +168,7 @@ def save_origin(self, origin: Origin, oois: list[OOI], valid_time: datetime) ->
                 raise ValueError("Origin source of observation does not exist")
 
         for ooi in oois:
-            self.ooi_repository.save(ooi, valid_time=valid_time)
+            self.ooi_repository.save(ooi, valid_time=valid_time, end_valid_time=end_valid_time)
         self.origin_repository.save(origin, valid_time=valid_time)
 
     def _run_inference(self, origin: Origin, valid_time: datetime) -> None:

diff --git a/octopoes/octopoes/repositories/origin_repository.py b/octopoes/octopoes/repositories/origin_repository.py
@@ -60,14 +60,14 @@ def commit(self):
 
     @classmethod
     def serialize(cls, origin: Origin) -> dict[str, Any]:
-        data = origin.dict()
+        data = origin.model_dump()
         data[cls.pk_prefix] = origin.id
         data["type"] = origin.__class__.__name__
         return data
 
     @classmethod
     def deserialize(cls, data: dict[str, Any]) -> Origin:
-        return Origin.parse_obj(data)
+        return Origin.model_validate(data)
 
     def list_origins(
         self,

diff --git a/octopoes/octopoes/xtdb/client.py b/octopoes/octopoes/xtdb/client.py
@@ -136,7 +136,7 @@ def await_transaction(self, transaction_id: int) -> None:
     def submit_transaction(self, operations: list[Operation]) -> None:
         res = self._session.post(
             f"{self.client_url()}/submit-tx",
-            content=Transaction(operations=operations).json(by_alias=True),
+            content=Transaction(operations=operations).model_dump_json(by_alias=True),
             headers={"Content-Type": "application/json"},
         )
 

diff --git a/octopoes/tests/integration/test_api_connector.py b/octopoes/tests/integration/test_api_connector.py
@@ -53,7 +53,7 @@ def test_bulk_operations(octopoes_api_connector: OctopoesAPIConnector, valid_tim
     assert len(octopoes_api_connector.list_origins(task_id=uuid.uuid4(), valid_time=valid_time)) == 0
     origins = octopoes_api_connector.list_origins(task_id=task_id, valid_time=valid_time)
     assert len(origins) == 1
-    assert origins[0].dict() == {
+    assert origins[0].model_dump() == {
         "method": "normalizer_id",
         "origin_type": OriginType.OBSERVATION,
         "source": network.reference,

diff --git a/octopoes/tests/test_octopoes_service.py b/octopoes/tests/test_octopoes_service.py
@@ -133,5 +133,5 @@ def test_on_create_scan_profile(octopoes_service, new_data, old_data, bit_runner
     octopoes_service.process_event(event)
 
     assert octopoes_service.ooi_repository.save.call_count == 2
-    octopoes_service.ooi_repository.save.assert_any_call(mock_oois[0], valid_time=valid_time)
-    octopoes_service.ooi_repository.save.assert_any_call(mock_oois[1], valid_time=valid_time)
+    octopoes_service.ooi_repository.save.assert_any_call(mock_oois[0], valid_time=valid_time, end_valid_time=None)
+    octopoes_service.ooi_repository.save.assert_any_call(mock_oois[1], valid_time=valid_time, end_valid_time=None)
diff --git a/rocky/rocky/locale/django.pot b/rocky/rocky/locale/django.pot
@@ -8,7 +8,7 @@ msgid ""
 msgstr ""
 "Project-Id-Version: PACKAGE VERSION\n"
 "Report-Msgid-Bugs-To: \n"
-"POT-Creation-Date: 2024-08-05 13:42+0000\n"
+"POT-Creation-Date: 2024-08-07 13:24+0000\n"
 "PO-Revision-Date: YEAR-MO-DA HO:MI+ZONE\n"
 "Last-Translator: FULL NAME <EMAIL@ADDRESS>\n"
 "Language-Team: LANGUAGE <[email protected]>\n"
@@ -4873,6 +4873,10 @@ msgstr ""
 msgid "Reason:"
 msgstr ""
 
+#: rocky/templates/findings/finding_list.html
+msgid "Expires by (UTC):"
+msgstr ""
+
 #: rocky/templates/findings/finding_list.html
 msgid "Unmute Findings"
 msgstr ""

diff --git a/rocky/rocky/templates/findings/finding_list.html b/rocky/rocky/templates/findings/finding_list.html
@@ -113,6 +113,10 @@ <h2 class="heading-normal">{% translate "Risk score" %}</h2>
                                     <div>
                                         <textarea name="reason" cols="5" rows="3" id="id_reason"></textarea>
                                     </div>
+                                    <div>
+                                        <label for="id_end_valid_time">{% translate "Expires by (UTC):" %}</label>
+                                        <input type="datetime-local" name="end_valid_time" id="id_end_valid_time">
+                                    </div>
                                 </div>
                             </fieldset>
                         {% endif %}

diff --git a/rocky/rocky/views/ooi_mute.py b/rocky/rocky/views/ooi_mute.py
@@ -42,6 +42,11 @@ def post(self, request, *args, **kwargs):
         unmute = request.POST.get("unmute", None)
         selected_findings = request.POST.getlist("finding", None)
         reason = request.POST.get("reason", None)
+        end_valid_time = request.POST.get("end_valid_time", None)
+        if end_valid_time:
+            end_valid_time = datetime.strptime(end_valid_time, "%Y-%m-%dT%H:%M").replace(tzinfo=timezone.utc)
+        else:
+            end_valid_time = None
 
         if not selected_findings:
             messages.add_message(self.request, messages.WARNING, _("Please select at least one finding."))
@@ -54,8 +59,10 @@ def post(self, request, *args, **kwargs):
             return redirect(reverse("finding_list", kwargs={"organization_code": self.organization.code}))
         else:
             for finding in selected_findings:
-                ooi = self.ooi_class.parse_obj({"finding": finding, "reason": reason})
-                create_ooi(self.octopoes_api_connector, self.bytes_client, ooi, datetime.now(timezone.utc))
+                ooi = self.ooi_class.model_validate({"finding": finding, "reason": reason})
+                create_ooi(
+                    self.octopoes_api_connector, self.bytes_client, ooi, datetime.now(timezone.utc), end_valid_time
+                )
 
             messages.add_message(self.request, messages.SUCCESS, _("Finding(s) successfully muted."))
             return redirect(reverse("finding_list", kwargs={"organization_code": self.organization.code}))
diff --git a/rocky/rocky/views/ooi_view.py b/rocky/rocky/views/ooi_view.py
@@ -181,12 +181,16 @@ def get_form_kwargs(self):
     def form_valid(self, form):
         # Transform into OOI
         try:
+            end_valid_time = form.cleaned_data.pop("end_valid_time", None)
+            if end_valid_time is not None:
+                end_valid_time = end_valid_time.replace(tzinfo=timezone.utc)
             new_ooi = self.ooi_class.model_validate(form.cleaned_data)
             create_ooi(
                 self.octopoes_api_connector,
                 self.bytes_client,
                 new_ooi,
                 datetime.now(timezone.utc),
+                end_valid_time,
             )
             sleep(1)
             return redirect(self.get_ooi_success_url(new_ooi))

diff --git a/rocky/tests/objects/test_objects_add.py b/rocky/tests/objects/test_objects_add.py
@@ -31,7 +31,8 @@ def test_add_ooi(rf, client_member, mock_organization_view_octopoes, mock_bytes_
 
 def test_add_bad_schema(rf, client_member):
     request = setup_request(
-        rf.post("ooi_add", {"ooi_type": "Network", "testnamewrong": "testnetwork"}), client_member.user
+        rf.post("ooi_add", {"ooi_type": "Network", "testnamewrong": "testnetwork"}),
+        client_member.user,
     )
 
     response = OOIAddView.as_view()(request, organization_code=client_member.organization.code, ooi_type="Network")

diff --git a/rocky/tools/forms/ooi.py b/rocky/tools/forms/ooi.py
@@ -117,3 +117,8 @@ class MuteFindingForm(forms.Form):
     finding = forms.CharField(widget=forms.HiddenInput(), required=False)
     ooi_type = forms.CharField(widget=forms.HiddenInput(), required=False)
     reason = forms.CharField(widget=forms.Textarea(attrs={"name": "reason", "rows": "3", "cols": "5"}), required=False)
+    end_valid_time = forms.DateTimeField(
+        label="Expires by (UTC)",
+        widget=forms.DateTimeInput(attrs={"name": "end_valid_time", "type": "datetime-local"}),
+        required=False,
+    )
diff --git a/rocky/tools/forms/ooi_form.py b/rocky/tools/forms/ooi_form.py
@@ -100,6 +100,18 @@ def generate_form_fields(
             else:
                 fields[name] = forms.CharField(max_length=256, **default_attrs)
 
+        # ruff: noqa: ERA001
+        # Currently we are not ready to use the following line as the
+        # event manager is not aware of the deletion of a generic OOI
+        # it does work for 'end-point'-OOIs like MutedFinding and the
+        # field is hidden for now
+        # fields["end_valid_time"] = forms.DateTimeField(
+        #     label="Expires by",
+        #     widget=forms.DateTimeInput(attrs={"type": "datetime-local"}),
+        #     required=False,
+        # )
+        fields["end_valid_time"] = forms.DateTimeField(widget=forms.HiddenInput(), required=False)
+
         return fields
 
 

diff --git a/rocky/tools/ooi_helpers.py b/rocky/tools/ooi_helpers.py
@@ -237,18 +237,28 @@ def get_finding_type_from_finding(finding: Finding) -> FindingType:
 
 
 def get_or_create_ooi(
-    api_connector: OctopoesAPIConnector, bytes_client: BytesClient, ooi: OOI, observed_at: datetime
+    api_connector: OctopoesAPIConnector,
+    bytes_client: BytesClient,
+    ooi: OOI,
+    observed_at: datetime,
+    end_valid_time: datetime | None = None,
 ) -> tuple[OOI, bool]:
     try:
         return api_connector.get(ooi.reference, observed_at), False
     except ObjectNotFoundException:
-        create_ooi(api_connector, bytes_client, ooi, observed_at)
+        create_ooi(api_connector, bytes_client, ooi, observed_at, end_valid_time)
         return ooi, True
 
 
-def create_ooi(api_connector: OctopoesAPIConnector, bytes_client: BytesClient, ooi: OOI, observed_at: datetime) -> None:
+def create_ooi(
+    api_connector: OctopoesAPIConnector,
+    bytes_client: BytesClient,
+    ooi: OOI,
+    observed_at: datetime,
+    end_valid_time: datetime | None = None,
+) -> None:
     task_id = uuid4()
-    declaration = Declaration(ooi=ooi, valid_time=observed_at, task_id=str(task_id))
+    declaration = Declaration(ooi=ooi, valid_time=observed_at, task_id=task_id, end_valid_time=end_valid_time)
     bytes_client.add_manual_proof(task_id, BytesClient.raw_from_declarations([declaration]))
 
     api_connector.save_declaration(declaration)