From 0cf9a5a55d99b12372a1aff5245f7c52d14f7422 Mon Sep 17 00:00:00 2001 From: originalsouth Date: Mon, 12 Aug 2024 12:03:46 +0200 Subject: [PATCH] Allow MuteFindings to expire by a user specified datetime (#3343) Co-authored-by: Jan Klopper Co-authored-by: Jeroen Dekkers Co-authored-by: ammar92 --- boefjes/boefjes/job_models.py | 3 ++- .../plugins/kat_manual/single_ooi/normalize.py | 3 ++- octopoes/octopoes/api/models.py | 2 ++ octopoes/octopoes/api/router.py | 2 +- octopoes/octopoes/core/service.py | 6 ++++-- .../octopoes/repositories/origin_repository.py | 4 ++-- octopoes/octopoes/xtdb/client.py | 2 +- .../tests/integration/test_api_connector.py | 2 +- octopoes/tests/test_octopoes_service.py | 4 ++-- rocky/rocky/locale/django.pot | 6 +++++- .../rocky/templates/findings/finding_list.html | 4 ++++ rocky/rocky/views/ooi_mute.py | 11 +++++++++-- rocky/rocky/views/ooi_view.py | 4 ++++ rocky/tests/objects/test_objects_add.py | 3 ++- rocky/tools/forms/ooi.py | 5 +++++ rocky/tools/forms/ooi_form.py | 12 ++++++++++++ rocky/tools/ooi_helpers.py | 18 ++++++++++++++---- 17 files changed, 72 insertions(+), 19 deletions(-) diff --git a/boefjes/boefjes/job_models.py b/boefjes/boefjes/job_models.py index 0d40cb497c3..8e419b79541 100644 --- 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 index 84ada07af3b..093ace2e3a1 100644 --- 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 index e2d2aabd962..7156231c6c9 100644 --- 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 index 4e6e0eae6a2..21a17cdeefa 100644 --- 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 index 52fef38b695..81f8ee9355f 100644 --- 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 index 252f166b1cd..0706e27ff4b 100644 --- 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 index 4e312e81f53..f78ef380b8b 100644 --- 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 index b57693e5fd5..34c28224f77 100644 --- 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 index 5eac373b197..dc7626c3b2b 100644 --- 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 index 573ef91b9de..1faac03bc93 100644 --- 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 \n" "Language-Team: LANGUAGE \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 index 6e34a5e7167..d9c5c47e65d 100644 --- a/rocky/rocky/templates/findings/finding_list.html +++ b/rocky/rocky/templates/findings/finding_list.html @@ -113,6 +113,10 @@

{% translate "Risk score" %}

+
+ + +
{% endif %} diff --git a/rocky/rocky/views/ooi_mute.py b/rocky/rocky/views/ooi_mute.py index a95759ffe71..89c8cef6106 100644 --- 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 index 983fdd709b6..5722858273a 100644 --- 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 index 8b69f6cdaf0..5c3a4accc81 100644 --- 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 index 849862f46af..6e910beda53 100644 --- 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 index ec7e3c19f0b..992edc08e96 100644 --- 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 index 7d2c6238254..ddf6e8d1b85 100644 --- 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)