perf(crisis_room): optimize loading of aggregates

octopoes/octopoes/api/router.py

@@ -2,7 +2,7 @@
 from datetime import datetime, timezone
 from http import HTTPStatus
 from logging import getLogger
-from typing import List, Optional, Set, Type
+from typing import Dict, List, Optional, Set, Type
 
 from fastapi import APIRouter, Depends, HTTPException, Path, Query, status
 from requests import HTTPError, RequestException
@@ -275,6 +275,14 @@ def get_scan_profile_inheritance(
     return octopoes.get_scan_profile_inheritance(reference, valid_time, [start])
 
 
+@router.get("/finding_types/count")
+def get_finding_type_count(
+    octopoes: OctopoesService = Depends(octopoes_service),
+    valid_time: datetime = Depends(extract_valid_time),
+) -> Dict[str, int]:
+    return octopoes.ooi_repository.get_finding_type_count(valid_time)
+
+
 @router.post("/node")
 def create_node(
     client: str = Depends(extract_client),

octopoes/octopoes/connector/octopoes.py

@@ -1,6 +1,6 @@
 import json
 from datetime import datetime
-from typing import List, Optional, Set, Type, Union
+from typing import Dict, List, Optional, Set, Type, Union
 
 import requests
 from pydantic.tools import parse_obj_as
@@ -162,3 +162,8 @@ def get_scan_profile_inheritance(
         params = {"reference": str(reference), "valid_time": valid_time}
         res = self.session.get(f"/{self.client}/scan_profiles/inheritance", params=params)
         return parse_obj_as(List[InheritanceSection], res.json())
+
+    def get_finding_type_count(self, valid_time: Optional[datetime] = None) -> Dict[str, int]:
+        params = {"valid_time": valid_time}
+        res = self.session.get(f"/{self.client}/finding_types/count", params=params)
+        return res.json()

octopoes/octopoes/repositories/ooi_repository.py

@@ -103,6 +103,9 @@ def get_tree(
     def list_oois_without_scan_profile(self, valid_time: datetime) -> Set[Reference]:
         raise NotImplementedError
 
+    def get_finding_type_count(self, valid_time: datetime) -> Dict[str, int]:
+        raise NotImplementedError
+
 
 class XTDBReferenceNode(BaseModel):
     __root__: Dict[str, Union[str, List[XTDBReferenceNode], XTDBReferenceNode]]
@@ -553,3 +556,12 @@ def list_oois_without_scan_profile(self, valid_time: datetime) -> Set[Reference]
         """
         response = self.session.client.query(query, valid_time=valid_time)
         return {Reference.from_str(row[0]) for row in response}
+
+    def get_finding_type_count(self, valid_time: datetime) -> Dict[str, int]:
+        query = """
+                    {:query {
+                     :find [?finding_type (count ?finding)]
+                     :where [[?finding :Finding/finding_type ?finding_type]] }}
+                """
+        response = self.session.client.query(query, valid_time=valid_time)
+        return {finding_type: count for finding_type, count in response}

rocky/crisis_room/views.py

@@ -1,6 +1,7 @@
 import logging
+from dataclasses import dataclass
 from datetime import datetime, timezone
-from typing import List, Union
+from typing import Dict, List, Set
 
 from account.models import KATUser
 from django.conf import settings
@@ -10,19 +11,66 @@
 from django.views.generic import TemplateView
 from django_otp.decorators import otp_required
 from tools.forms.base import ObservedAtForm
-from tools.models import Organization
+from tools.models import OOIInformation, Organization
+from tools.ooi_helpers import (
+    RiskLevelSeverity,
+    get_risk_level_score,
+    get_risk_level_score_for_cve,
+    get_risk_level_score_for_retirejs,
+    get_risk_level_score_for_snyk,
+)
 from tools.view_helpers import BreadcrumbsMixin, convert_date_to_datetime
 from two_factor.views.utils import class_view_decorator
 
 from octopoes.connector import ConnectorException
 from octopoes.connector.octopoes import OctopoesAPIConnector
-from octopoes.models.ooi.findings import Finding
-from rocky.views.ooi_report import build_findings_list_from_store
 from rocky.views.ooi_view import ConnectorFormMixin
 
 logger = logging.getLogger(__name__)
 
 
+# dataclass to store finding type counts
+@dataclass
+class OrganizationFindingTypeCount:
+    name: str
+    code: str
+    finding_count_per_severity: Dict[str, int]
+
+    @property
+    def total(self) -> int:
+        return sum(self.finding_count_per_severity.values())
+
+    @property
+    def total_critical(self) -> int:
+        return self.finding_count_per_severity[RiskLevelSeverity.CRITICAL.value]
+
+
+def load_finding_type_risks(finding_types: Set[str]) -> Dict[str, str]:
+    """Load finding type from db and map to uniform dataclass"""
+
+    # initialize default entry for each finding_type
+    severities = {ft: RiskLevelSeverity.CRITICAL.value for ft in finding_types}
+
+    # query for data in ooi information table and override for each hit
+    information_records = list(OOIInformation.objects.filter(id__in=finding_types))
+    for ooi_info in information_records:
+        finding_type_type = ooi_info.pk.split("|")[0]
+        # hook into the old code to re-use the risk level calculation logic
+        # FIXME: refactor this duplicated logic in future PR
+        if finding_type_type == "CVEFindingType":
+            risk_level = get_risk_level_score_for_cve(ooi_info.data)
+        elif finding_type_type == "RetireJSFindingType":
+            risk_level = get_risk_level_score_for_retirejs(ooi_info.data)
+        elif finding_type_type == "SnykFindingType":
+            risk_level = get_risk_level_score_for_snyk(ooi_info.data)
+        else:
+            risk_level = get_risk_level_score(ooi_info.data)
+
+        severities[ooi_info.pk] = risk_level["risk_level_severity"]
+
+    return severities
+
+
 class CrisisRoomBreadcrumbsMixin(BreadcrumbsMixin):
     breadcrumbs = [
         {"url": "", "text": "Crisis Room"},
@@ -31,32 +79,31 @@ class CrisisRoomBreadcrumbsMixin(BreadcrumbsMixin):
 
 @class_view_decorator(otp_required)
 class CrisisRoomView(CrisisRoomBreadcrumbsMixin, ConnectorFormMixin, TemplateView):
-    ooi_types = {Finding}
     template_name = "crisis_room/crisis_room.html"
     connector_form_class = ObservedAtForm
 
-    def sort_finding_list_by_total(self, finding_list) -> List:
+    def sort_finding_list_by_total(
+        self, org_finding_type_counts: List[OrganizationFindingTypeCount]
+    ) -> List[OrganizationFindingTypeCount]:
         is_desc = self.request.GET.get("sort_total_by", "desc") != "asc"
-        _finding_list = finding_list.copy()
-        _finding_list.sort(key=lambda x: x["meta"]["total"], reverse=is_desc)
-        return _finding_list
+        return sorted(org_finding_type_counts, key=lambda x: x.total, reverse=is_desc)
 
-    def sort_finding_list_by_critical(self, finding_list) -> List:
+    def sort_finding_list_by_critical(
+        self, org_finding_type_counts: List[OrganizationFindingTypeCount]
+    ) -> List[OrganizationFindingTypeCount]:
         is_desc = self.request.GET.get("sort_critical_by", "desc") != "asc"
-        finding_list.sort(key=lambda x: x["meta"]["total_by_severity"]["critical"], reverse=is_desc)
-        return finding_list
+        return sorted(org_finding_type_counts, key=lambda x: x.total_critical, reverse=is_desc)
 
-    def get_list_for_org(self, organization: Organization) -> Union[List, None]:
+    def get_finding_type_count(self, organization: Organization) -> Dict[str, int]:
         try:
             api_connector = OctopoesAPIConnector(settings.OCTOPOES_API, organization.code)
-
-            return api_connector.list(self.ooi_types, valid_time=self.get_observed_at()).items
+            return api_connector.get_finding_type_count(valid_time=self.get_observed_at())
         except ConnectorException:
             messages.add_message(
                 self.request, messages.ERROR, _("Failed to get list of findings, check server logs for more details.")
             )
             logger.exception("Failed to get list of findings for organization %s", organization.code)
-            return []
+            return {}
 
     def get_observed_at(self) -> datetime:
         if "observed_at" not in self.request.GET:
@@ -73,22 +120,36 @@ def get_context_data(self, **kwargs):
 
         user: KATUser = self.request.user
 
-        findings_per_org = []
+        # query each organization's finding type count
+        finding_type_counts = {}
         for org in user.organizations:
-            findings = self.get_list_for_org(org)
-            findings_store = {finding.primary_key: finding for finding in findings}
+            finding_type_counts[(org.code, org.name)] = self.get_finding_type_count(org)
+
+        # find all unique finding types and query their severity
+        unique_finding_types = set()
+        for org_code, finding_type_count in finding_type_counts.items():
+            unique_finding_types.update(finding_type_count.keys())
 
-            findings_ = build_findings_list_from_store(findings_store)
-            findings_["organization"] = org
-            findings_per_org.append(findings_)
+        severities = load_finding_type_risks(unique_finding_types)
+
+        # aggregate per organization, per severity
+        org_finding_type_counts = []
+        for (org_code, org_name), finding_type_count in finding_type_counts.items():
+            org_severity_count: Dict[str, int] = {severity_level.value: 0 for severity_level in RiskLevelSeverity}
+            for finding_type, count in finding_type_count.items():
+                finding_type_severity = severities[finding_type]
+                org_severity_count[finding_type_severity] += count
+            org_finding_type_counts.append(OrganizationFindingTypeCount(org_name, org_code, org_severity_count))
 
         context["breadcrumb_list"] = [
             {"url": reverse("crisis_room"), "text": "CRISIS ROOM"},
         ]
 
         context["organizations"] = user.organizations
-        context["findings_per_org_total"] = self.sort_finding_list_by_total(findings_per_org)
-        context["findings_per_org_critical"] = self.sort_finding_list_by_critical(findings_per_org)
+
+        context["org_finding_type_counts"] = self.sort_finding_list_by_total(org_finding_type_counts)
+        context["org_finding_type_counts_critical"] = self.sort_finding_list_by_critical(org_finding_type_counts)
+
         context["observed_at_form"] = self.get_connector_form()
         context["observed_at"] = self.get_observed_at().date()
 

rocky/rocky/templates/crisis_room/crisis_room_findings_block.html

@@ -14,12 +14,12 @@ <h1>{% translate "Total findings" %}</h1>
                         </tr>
                     </thead>
                     <tbody>
-                        {% for findings in findings_per_org_total %}
+                        {% for org_finding_type_count in org_finding_type_counts %}
                             <tr>
                                 <td>
-                                    <a href="{% url "organization_crisis_room" findings.organization.code %}">{{ findings.organization.name }}</a>
+                                    <a href="{% url "organization_crisis_room" org_finding_type_count.code %}">{{ org_finding_type_count.name }}</a>
                                 </td>
-                                <td class="number">{{ findings|get_item:"meta"|get_item:"total" }}</td>
+                                <td class="number">{{ org_finding_type_count.total }}</td>
                                 <td>
                                     <button class="expando-button"
                                             data-icon-open-class="icon ti-chevron-down"
@@ -31,27 +31,27 @@ <h1>{% translate "Total findings" %}</h1>
                             </tr>
                             <tr class="expando-row">
                                 <td colspan="3">
-                                    <span class="sr-only">{{ org }} {% translate " Finding Details" %}</span>
+                                    <span class="sr-only">{{ org_finding_type_count.name }} {% translate " Finding Details" %}</span>
                                     <table>
                                         <thead>
                                             <th>{% translate "Severity" %}</th>
                                             <th>{% translate "Occurrences" %}</th>
                                         </thead>
                                         <tbody>
-                                            {% for key, value in findings.meta.total_by_severity.items %}
+                                            {% for severity, count in org_finding_type_count.finding_count_per_severity.items %}
                                                 <tr>
                                                     <td>
-                                                        {% if value != 0 %}
-                                                            <a href="{% url "finding_list" organization_code=findings.organization.code %}?severity={{ key }}"><span class="{{ key }}">{{ key|title }}</span></a>
+                                                        {% if count != 0 %}
+                                                            <a href="{% url "finding_list" organization_code=org_finding_type_count.code %}?severity={{ severity }}"><span class="{{ severity }}">{{ severity|title }}</span></a>
                                                         {% else %}
-                                                            <span class="{{ key }}">{{ key|title }}</span>
+                                                            <span class="{{ severity }}">{{ severity|title }}</span>
                                                         {% endif %}
                                                     </td>
                                                     <td class="number">
-                                                        {% if value != 0 %}
-                                                            <a href="{% url "finding_list" organization_code=findings.organization.code %}?severity={{ key }}">{{ value }}</a>
+                                                        {% if count != 0 %}
+                                                            <a href="{% url "finding_list" organization_code=org_finding_type_count.code %}?severity={{ severity }}">{{ count }}</a>
                                                         {% else %}
-                                                            {{ value }}
+                                                            {{ count }}
                                                         {% endif %}
                                                     </td>
                                                 </tr>
@@ -83,12 +83,12 @@ <h1>{% translate "Critical findings" %}</h1>
                         </tr>
                     </thead>
                     <tbody>
-                        {% for findings in findings_per_org_critical %}
+                        {% for org_finding_type_count in org_finding_type_counts_critical %}
                             <tr>
                                 <td>
-                                    <a href="{% url "organization_crisis_room" findings.organization.code %}">{{ findings.organization.name }}</a>
+                                    <a href="{% url "organization_crisis_room" org_finding_type_count.code %}">{{ org_finding_type_count.name }}</a>
                                 </td>
-                                <td class="number">{{ findings|get_item:"meta"|get_item:"total_by_severity"|get_item:"critical" }}</td>
+                                <td class="number">{{ org_finding_type_count.total_critical }}</td>
                                 <td>
                                     <button class="expando-button"
                                             data-icon-open-class="icon ti-chevron-down"
@@ -100,27 +100,27 @@ <h1>{% translate "Critical findings" %}</h1>
                             </tr>
                             <tr class="expando-row">
                                 <td colspan="3">
-                                    <span class="sr-only">{{ org }} {% translate " Finding Details" %}</span>
+                                    <span class="sr-only">{{ org_finding_type_count.name }} {% translate " Finding Details" %}</span>
                                     <table>
                                         <thead>
                                             <th>{% translate "Severity" %}</th>
                                             <th>{% translate "Occurrences" %}</th>
                                         </thead>
                                         <tbody>
-                                            {% for key, value in findings.meta.total_by_severity.items %}
+                                            {% for severity, count in org_finding_type_count.finding_count_per_severity.items %}
                                                 <tr>
                                                     <td>
-                                                        {% if value != 0 %}
-                                                            <a href="{% url "finding_list" organization_code=findings.organization.code %}?severity={{ key }}"><span class="{{ key }}">{{ key|title }}</span></a>
+                                                        {% if count != 0 %}
+                                                            <a href="{% url "finding_list" organization_code=org_finding_type_count.code %}?severity={{ severity }}"><span class="{{ severity }}">{{ severity|title }}</span></a>
                                                         {% else %}
-                                                            <span class="{{ key }}">{{ key|title }}</span>
+                                                            <span class="{{ severity }}">{{ severity|title }}</span>
                                                         {% endif %}
                                                     </td>
                                                     <td class="number">
-                                                        {% if value != 0 %}
-                                                            <a href="{% url "finding_list" organization_code=findings.organization.code %}?severity={{ key }}">{{ value }}</a>
+                                                        {% if count != 0 %}
+                                                            <a href="{% url "finding_list" organization_code=org_finding_type_count.code %}?severity={{ severity }}">{{ count }}</a>
                                                         {% else %}
-                                                            {{ value }}
+                                                            {{ count }}
                                                         {% endif %}
                                                     </td>
                                                 </tr>

rocky/tests/conftest.py

@@ -1,5 +1,6 @@
 import binascii
 import json
+import logging
 from os import urandom
 from pathlib import Path
 from unittest.mock import MagicMock, patch
@@ -26,6 +27,9 @@
 from octopoes.models.ooi.network import Network
 from rocky.scheduler import Task
 
+# Quiet faker locale messages down in tests.
+logging.getLogger("faker").setLevel(logging.INFO)
+
 
 def create_user(django_user_model, email, password, name, device_name, superuser=False):
     user = django_user_model.objects.create_user(email=email, password=password)