Endpoint rerun bits (#899)

Co-authored-by: Jesse Lisser <[email protected]>
Co-authored-by: ammar92 <[email protected]>
Co-authored-by: Donny Peeters <[email protected]>

octopoes/octopoes/api/router.py

@@ -305,3 +305,12 @@ def delete_node(xtdb_session_: XTDBSession = Depends(xtdb_session)) -> None:
         )
     except XTDBException as e:
         raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Deleting node failed") from e
+
+
+@router.post("/bits/recalculate")
+def recalculate_bits(
+    xtdb_session_: XTDBSession = Depends(xtdb_session), octopoes: OctopoesService = Depends(octopoes_service)
+) -> int:
+    inference_count = octopoes.recalculate_bits()
+    xtdb_session_.commit()
+    return inference_count

octopoes/octopoes/connector/octopoes.py

@@ -167,3 +167,6 @@ def get_finding_type_count(self, valid_time: Optional[datetime] = None) -> Dict[
         params = {"valid_time": valid_time}
         res = self.session.get(f"/{self.client}/finding_types/count", params=params)
         return res.json()
+
+    def recalculate_bits(self) -> int:
+        return self.session.post(f"/{self.client}/bits/recalculate").json()

octopoes/octopoes/core/service.py

@@ -1,5 +1,5 @@
 import json
-from datetime import datetime
+from datetime import datetime, timezone
 from logging import getLogger
 from typing import Callable, Dict, List, Optional, Set, Type
 
@@ -509,3 +509,10 @@ def get_scan_profile_inheritance(
                 return expl
 
         return inheritance_chain
+
+    def recalculate_bits(self) -> int:
+        valid_time = datetime.now(timezone.utc)
+        origins = self.origin_repository.list(origin_type=OriginType.INFERENCE, valid_time=valid_time)
+        for origin in origins:
+            self._run_inference(origin, valid_time)
+        return len(origins)

octopoes/octopoes/repositories/origin_repository.py

@@ -10,7 +10,7 @@
 from octopoes.events.manager import EventManager
 from octopoes.models import Reference
 from octopoes.models.exception import ObjectNotFoundException
-from octopoes.models.origin import Origin
+from octopoes.models.origin import Origin, OriginType
 from octopoes.xtdb import FieldSet
 from octopoes.xtdb.client import OperationType as XTDBOperationType
 from octopoes.xtdb.client import XTDBSession
@@ -38,6 +38,9 @@ def list_by_source(self, reference: Reference, valid_time: datetime) -> List[Ori
     def delete(self, origin: Origin, valid_time: datetime) -> None:
         raise NotImplementedError
 
+    def list(self, origin_type: OriginType, valid_time: datetime) -> List[Origin]:
+        raise NotImplementedError
+
 
 class XTDBOriginRepository(OriginRepository):
     xtdb_type: XTDBType = XTDBType.CRUX
@@ -123,3 +126,14 @@ def delete(self, origin: Origin, valid_time: datetime) -> None:
             old_data=origin,
         )
         self.session.listen_post_commit(lambda: self.event_manager.publish(event))
+
+    def list(self, origin_type: OriginType, valid_time: datetime) -> List[Origin]:
+        query = generate_pull_query(
+            self.xtdb_type,
+            FieldSet.ALL_FIELDS,
+            {
+                "origin_type": origin_type.value,
+            },
+        )
+        results = self.session.client.query(query, valid_time=valid_time)
+        return [self.deserialize(r[0]) for r in results]

octopoes/tests/robot/07_rerun_bits.robot

@@ -0,0 +1,24 @@
+*** Settings ***
+Resource            robot.resource
+
+Test Setup          Setup Test
+Test Teardown       Teardown Test
+
+
+*** Test Cases ***
+Rerun bits
+    Insert Observation    tests/fixtures/normalizer_output.json
+    Await Sync
+    ${response}    Post    ${OCTOPOES_URI}/bits/recalculate
+    Should Be Equal As Integers    ${response.status_code}    200
+    Should Be Equal As Integers    ${response.content}    19
+
+
+*** Keywords ***
+Setup Test
+    Start Monitoring    ${QUEUE_URI}
+
+Teardown Test
+    Cleanup
+    Await Sync
+    Stop Monitoring

rocky/rocky/templates/organizations/organization_settings.html

@@ -4,73 +4,84 @@
 {% load static %}
 
 {% block content %}
-    {% include "header.html" %}
+  {% include "header.html" %}
 
-    <main id="main-content">
-        <section>
-            <div>
-                <h1>{% translate "Organization" %}: {{ organization.name }}</h1>
-                <p>
-                    {% blocktranslate with organization_name=organization.name %}
-                    An overview of "{{ organization_name }}". This shows general information and its settings.
-                    {% endblocktranslate %}
-                </p>
-                {% if not indemnification_present %}
-                    <p class="warning"
-                       role="group"
-                       aria-label="{% translate "indemnification warning" %}">
-                        {% url "organization_settings" organization.code as organization_settings %}
-                        {% blocktranslate %}
-                            <strong>Warning:</strong>
-                            Indemnification is not set for this organization.
-                        {% endblocktranslate %}
-                    </p>
-                {% endif %}
-                <h2>{% translate "Organization details" %}</h2>
-                <div class="horizontal-scroll">
-                    <table>
-                        <thead>
-                            <tr>
-                                <th>{% translate "Name" %}</th>
-                                <th>{% translate "Code" %}</th>
-                                <th>{% translate "Edit" %}</th>
-                            </tr>
-                        </thead>
-                        <tbody>
-                            <tr>
-                                <td>{{ organization.name }}</td>
-                                <td>{{ organization.code }}</td>
-                                <td>
-                                    {% if perms.tools.can_change_organization %}
-                                        {% spaceless %}
-                                            <a href="{% url "organization_edit" organization.code %}">
-                                                <button class="icon ti-edit action-button">{% translate "Edit" %}</button>
-                                            </a>
-                                        {% endspaceless %}
-                                    {% endif %}
-                                </td>
-                            </tr>
-                        </tbody>
-                    </table>
-                </div>
-                {% if perms.tools.add_indemnification %}
-                    <a href="{% url "indemnification_add" organization.code %}"
-                       class="button ghost">{% translate "Add indemnification" %}</a>
+  <main id="main-content">
+    <section>
+      <div>
+        <h1>{% translate "Organization" %}: {{ organization.name }}</h1>
+        <p>
+          {% blocktranslate with organization_name=organization.name %}
+            An overview of "{{ organization_name }}". This shows general information and its settings.
+          {% endblocktranslate %}
+        </p>
+        {% if not indemnification_present %}
+          <p class="warning"
+             role="group"
+             aria-label="{% translate "indemnification warning" %}">
+            {% url "organization_settings" organization.code as organization_settings %}
+            {% blocktranslate %}
+              <strong>Warning:</strong>
+              Indemnification is not set for this organization.
+            {% endblocktranslate %}
+          </p>
+        {% endif %}
+        <h2>{% translate "Organization details" %}</h2>
+        <div class="horizontal-scroll">
+          <table>
+            <thead>
+            <tr>
+              <th>{% translate "Name" %}</th>
+              <th>{% translate "Code" %}</th>
+              <th>{% translate "Edit" %}</th>
+            </tr>
+            </thead>
+            <tbody>
+            <tr>
+              <td>{{ organization.name }}</td>
+              <td>{{ organization.code }}</td>
+              <td>
+                {% if perms.tools.can_change_organization %}
+                  {% spaceless %}
+                    <a href="{% url "organization_edit" organization.code %}">
+                      <button class="icon ti-edit action-button">{% translate "Edit" %}</button>
+                    </a>
+                  {% endspaceless %}
                 {% endif %}
-            </div>
-        </section>
-        {% if organization.tags.all %}
-            <section>
-                <div>
-                    <h2>{% translate "Tags" %}</h2>
-                    {% include "organizations/organization_tags.html" %}
-
-                </div>
-            </section>
+              </td>
+            </tr>
+            </tbody>
+          </table>
+        </div>
+        {% if perms.tools.add_indemnification %}
+          <a href="{% url "indemnification_add" organization.code %}"
+             class="button ghost">{% translate "Add indemnification" %}</a>
         {% endif %}
-    </main>
+        {% if perms.tools.can_recalculate_bits %}
+          <form method="post" class="inline">
+            {% csrf_token %}
+            <button type="submit"
+                    class="dropdown-button"
+                    name="action"
+                    value="recalculate">
+              {% translate "Rerun all bits" %}
+            </button>
+          </form>
+        {% endif %}
+      </div>
+    </section>
+    {% if organization.tags.all %}
+      <section>
+        <div>
+          <h2>{% translate "Tags" %}</h2>
+          {% include "organizations/organization_tags.html" %}
+
+        </div>
+      </section>
+    {% endif %}
+  </main>
 {% endblock content %}
 {% block html_at_end_body %}
-    {{ block.super }}
-    <script src="{% static "/js/checkboxToggler.js" %}" nonce="{{ request.csp_nonce }}"></script>
+  {{ block.super }}
+  <script src="{% static "/js/checkboxToggler.js" %}" nonce="{{ request.csp_nonce }}"></script>
 {% endblock html_at_end_body %}

rocky/rocky/views/organization_settings.py

@@ -1,11 +1,38 @@
+from datetime import datetime
+from enum import Enum
+
 from account.mixins import OrganizationPermissionRequiredMixin
+from django.contrib import messages
+from django.core.exceptions import PermissionDenied
+from django.http import HttpRequest, HttpResponse, HttpResponseBadRequest
+from django.utils.translation import gettext_lazy as _
 from django.views.generic import TemplateView
 from django_otp.decorators import otp_required
 from tools.view_helpers import OrganizationDetailBreadcrumbsMixin
 from two_factor.views.utils import class_view_decorator
 
 
+class PageActions(Enum):
+    RECALCULATE = "recalculate"
+
+
 @class_view_decorator(otp_required)
 class OrganizationSettingsView(OrganizationPermissionRequiredMixin, OrganizationDetailBreadcrumbsMixin, TemplateView):
     template_name = "organizations/organization_settings.html"
     permission_required = "tools.view_organization"
+
+    def post(self, request: HttpRequest, *args, **kwargs) -> HttpResponse:
+        """Perform actions based on action type"""
+        action = request.POST.get("action")
+        if not self.request.user.has_perm("tools.can_recalculate_bits"):
+            raise PermissionDenied()
+        if action == PageActions.RECALCULATE.value:
+            connector = self.octopoes_api_connector
+
+            start_time = datetime.now()
+            number_of_bits = connector.recalculate_bits()
+            duration = datetime.now() - start_time
+            messages.add_message(request, messages.INFO, _(f"Recalculated {number_of_bits} bits. Duration: {duration}"))
+            return self.get(request, *args, **kwargs)
+        else:
+            raise HttpResponseBadRequest("Unknown action")

rocky/tools/management/commands/setup_dev_account.py

@@ -48,6 +48,7 @@ def handle(self, *args, **options):
                 "change_organizationmember",
                 "can_delete_oois",
                 "add_indemnification",
+                "can_recalculate_bits",
             ]
         )
         self.group_admin.permissions.set(admin_permissions)

rocky/tools/migrations/0037_alter_organization_options.py

@@ -0,0 +1,25 @@
+# Generated by Django 3.2.19 on 2023-05-09 12:13
+
+from django.db import migrations
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("tools", "0036_merge_20230504_1629"),
+    ]
+
+    operations = [
+        migrations.AlterModelOptions(
+            name="organization",
+            options={
+                "permissions": (
+                    ("can_switch_organization", "Can switch organization"),
+                    ("can_scan_organization", "Can scan organization"),
+                    ("can_enable_disable_boefje", "Can enable or disable boefje"),
+                    ("can_set_clearance_level", "Can set clearance level"),
+                    ("can_delete_oois", "Can delete oois"),
+                    ("can_recalculate_bits", "Can recalculate bits"),
+                )
+            },
+        ),
+    ]

rocky/tools/models.py

@@ -78,6 +78,7 @@ class Meta:
             ("can_enable_disable_boefje", "Can enable or disable boefje"),
             ("can_set_clearance_level", "Can set clearance level"),
             ("can_delete_oois", "Can delete oois"),
+            ("can_recalculate_bits", "Can recalculate bits"),
         )
 
     def get_absolute_url(self):