Error handling for Generate Report (#2274)

Co-authored-by: Noam Blitz <[email protected]> Co-authored-by: noamblitz <[email protected]> Co-authored-by: Jeroen Dekkers <[email protected]>
minvws · Jan 23, 2024 · 4a35dac · 4a35dac
1 parent fbb9a85
commit 4a35dac
Show file tree

Hide file tree

Showing 17 changed files with 212 additions and 89 deletions.
diff --git a/rocky/reports/report_types/aggregate_organisation_report/report.py b/rocky/reports/report_types/aggregate_organisation_report/report.py
@@ -447,18 +447,21 @@ def aggregate_reports(
 ):
     aggregate_report = AggregateOrganisationReport(connector)
     report_data = {}
+    error_oois = []
 
     for ooi in input_ooi_references:
         report_data[ooi] = {}
-        for options, report_types in aggregate_report.reports.items():
-            for report_type in report_types:
-                if Reference.from_str(ooi).class_type in report_type.input_ooi_types and report_type.id in [
-                    report["id"] for report in selected_report_types
-                ]:
-                    report = report_type(connector)
-                    data = report.generate_data(ooi, valid_time=valid_time)
-                    report_data[ooi][report_type.id] = data
-
+        try:
+            for options, report_types in aggregate_report.reports.items():
+                for report_type in report_types:
+                    if Reference.from_str(ooi).class_type in report_type.input_ooi_types and report_type.id in [
+                        report.id for report in selected_report_types
+                    ]:
+                        report = report_type(connector)
+                        data = report.generate_data(ooi, valid_time=valid_time)
+                        report_data[ooi][report_type.id] = data
+        except Exception:
+            error_oois.append(ooi)
     post_processed_data = aggregate_report.post_process_data(report_data, valid_time=valid_time)
 
-    return aggregate_report, post_processed_data, report_data
+    return aggregate_report, post_processed_data, report_data, error_oois
diff --git a/rocky/reports/report_types/definitions.py b/rocky/reports/report_types/definitions.py
@@ -35,6 +35,17 @@ def __init__(self, octopoes_api_connector: OctopoesAPIConnector):
     def generate_data(self, input_ooi: str, valid_time: datetime) -> Dict[str, Any]:
         raise NotImplementedError
 
+    @classmethod
+    def class_attributes(cls) -> Dict[str, any]:
+        return {
+            "id": cls.id,
+            "name": cls.name,
+            "description": cls.description,
+            "plugins": cls.plugins,
+            "input_ooi_types": cls.input_ooi_types,
+            "template_path": cls.template_path,
+        }
+
 
 class AggregateReport(ABC):
     id: str

diff --git a/rocky/reports/report_types/helpers.py b/rocky/reports/report_types/helpers.py
@@ -7,6 +7,7 @@
 from reports.report_types.dns_report.report import DNSReport
 from reports.report_types.ipv6_report.report import IPv6Report
 from reports.report_types.mail_report.report import MailReport
+from reports.report_types.multi_organization_report.report import MultiOrganizationReport
 from reports.report_types.name_server_report.report import NameServerSystemReport
 from reports.report_types.open_ports_report.report import OpenPortsReport
 from reports.report_types.rpki_report.report import RPKIReport
@@ -32,6 +33,8 @@
 ]
 AGGREGATE_REPORTS = [AggregateOrganisationReport]
 
+MULTI_REPORTS = [MultiOrganizationReport]
+
 
 def get_ooi_types_with_report() -> Set[Type[OOI]]:
     """
@@ -60,7 +63,7 @@ def get_report_by_id(report_id: str) -> Type[Report]:
     """
     Get report type by id
     """
-    for report in REPORTS:
+    for report in REPORTS + MULTI_REPORTS:
         if report.id == report_id:
             return report
     raise ValueError(f"Report with id {report_id} not found")

diff --git a/rocky/reports/report_types/ipv6_report/report.py b/rocky/reports/report_types/ipv6_report/report.py
@@ -6,6 +6,7 @@
 from django.utils.translation import gettext_lazy as _
 
 from octopoes.models import Reference
+from octopoes.models.exception import ObjectNotFoundException
 from octopoes.models.ooi.dns.zone import Hostname
 from octopoes.models.ooi.network import IPAddressV4, IPAddressV6
 from octopoes.models.path import Path
@@ -33,21 +34,25 @@ def generate_data(self, input_ooi: str, valid_time: datetime) -> Dict[str, Any]:
         For hostnames, check whether they point to ipv6 addresses.
         For ip addresses, check all hostnames that point to them, and check whether they point to ipv6 addresses.
         """
-        ref = Reference.from_str(input_ooi)
-        if ref.class_type == IPAddressV4 or ref.class_type == IPAddressV6:
+        try:
+            ooi = self.octopoes_api_connector.get(Reference.from_str(input_ooi), valid_time)
+        except ObjectNotFoundException as e:
+            logger.error("No data found for OOI '%s' on date %s.", str(e), str(valid_time))
+            raise ObjectNotFoundException(e)
+
+        if ooi.reference.class_type == IPAddressV4 or ooi.reference.class_type == IPAddressV6:
             path = Path.parse("IPAddress.<address [is ResolvedHostname].hostname")
-            hostnames = self.octopoes_api_connector.query(path=path, source=ref, valid_time=valid_time)
-            hostnames = [h.reference for h in hostnames]
+            hostnames = self.octopoes_api_connector.query(path=path, source=ooi.reference, valid_time=valid_time)
         else:
-            hostnames = [ref]
+            hostnames = [ooi]
 
         results = {}
         for hostname in hostnames:
             path = Path.parse("Hostname.<hostname [is ResolvedHostname].address")
-            ips = self.octopoes_api_connector.query(path=path, source=hostname, valid_time=valid_time)
+            ips = self.octopoes_api_connector.query(path=path, source=hostname.reference, valid_time=valid_time)
 
             results = {
-                hostname.tokenized.name: {"enabled": any(ip.reference.class_type == IPAddressV6 for ip in ips)}
+                hostname.name: {"enabled": any(ip.reference.class_type == IPAddressV6 for ip in ips)}
                 for hostname in hostnames
             }
 

diff --git a/rocky/reports/report_types/mail_report/report.py b/rocky/reports/report_types/mail_report/report.py
@@ -5,6 +5,7 @@
 from django.utils.translation import gettext_lazy as _
 
 from octopoes.models import OOI, Reference
+from octopoes.models.exception import ObjectNotFoundException
 from octopoes.models.ooi.dns.zone import Hostname
 from octopoes.models.ooi.network import IPAddressV4, IPAddressV6
 from reports.report_types.definitions import Report
@@ -23,17 +24,22 @@ class MailReport(Report):
     template_path = "mail_report/report.html"
 
     def generate_data(self, input_ooi: str, valid_time: datetime) -> Dict[str, Any]:
-        reference = Reference.from_str(input_ooi)
         hostnames = []
         mail_security_measures = {}
 
-        if reference.class_type == Hostname:
-            hostnames = [reference]
-        elif reference.class_type in (IPAddressV4, IPAddressV6):
+        try:
+            ooi = self.octopoes_api_connector.get(Reference.from_str(input_ooi), valid_time)
+        except ObjectNotFoundException as e:
+            logger.error("No data found for OOI '%s' on date %s.", str(e), str(valid_time))
+            raise ObjectNotFoundException(e)
+
+        if ooi.reference.class_type == Hostname:
+            hostnames = [ooi]
+        elif ooi.reference.class_type in (IPAddressV4, IPAddressV6):
             hostnames = [
                 x.reference
                 for x in self.octopoes_api_connector.query(
-                    "IPAddress.<address[is ResolvedHostname].hostname", valid_time, reference
+                    "IPAddress.<address[is ResolvedHostname].hostname", valid_time, ooi.reference
                 )
             ]
 

diff --git a/rocky/reports/report_types/name_server_report/report.py b/rocky/reports/report_types/name_server_report/report.py
@@ -6,6 +6,7 @@
 from django.utils.translation import gettext_lazy as _
 
 from octopoes.models import Reference
+from octopoes.models.exception import ObjectNotFoundException
 from octopoes.models.ooi.dns.zone import Hostname
 from octopoes.models.ooi.findings import RiskLevelSeverity
 from octopoes.models.ooi.network import IPAddressV4, IPAddressV6
@@ -66,15 +67,20 @@ class NameServerSystemReport(Report):
     template_path = "name_server_report/report.html"
 
     def generate_data(self, input_ooi: str, valid_time: datetime) -> Dict[str, Any]:
-        reference = Reference.from_str(input_ooi)
         hostnames = []
 
-        if reference.class_type == Hostname:
-            hostnames = [self.octopoes_api_connector.get(reference)]
+        try:
+            ooi = self.octopoes_api_connector.get(Reference.from_str(input_ooi), valid_time)
+        except ObjectNotFoundException as e:
+            logger.error("No data found for OOI '%s' on date %s.", str(e), str(valid_time))
+            raise ObjectNotFoundException(e)
 
-        elif reference.class_type in (IPAddressV4, IPAddressV6):
+        if ooi.reference.class_type == Hostname:
+            hostnames = [ooi]
+
+        elif ooi.reference.class_type in (IPAddressV4, IPAddressV6):
             hostnames = self.octopoes_api_connector.query(
-                "IPAddress.<address[is ResolvedHostname].hostname", valid_time, reference
+                "IPAddress.<address[is ResolvedHostname].hostname", valid_time, ooi.reference
             )
 
         name_server_checks = NameServerChecks()

diff --git a/rocky/reports/report_types/open_ports_report/report.py b/rocky/reports/report_types/open_ports_report/report.py
@@ -5,6 +5,7 @@
 from django.utils.translation import gettext_lazy as _
 
 from octopoes.models import Reference
+from octopoes.models.exception import ObjectNotFoundException
 from octopoes.models.ooi.dns.zone import Hostname
 from octopoes.models.ooi.network import IPAddressV4, IPAddressV6
 from octopoes.models.path import Path
@@ -22,15 +23,20 @@ class OpenPortsReport(Report):
     template_path = "open_ports_report/report.html"
 
     def generate_data(self, input_ooi: str, valid_time: datetime) -> Dict[str, Any]:
-        ref = Reference.from_str(input_ooi)
-        if ref.class_type == Hostname:
+        try:
+            ooi = self.octopoes_api_connector.get(Reference.from_str(input_ooi), valid_time)
+        except ObjectNotFoundException as e:
+            logger.error("No data found for OOI '%s' on date %s.", str(e), str(valid_time))
+            raise ObjectNotFoundException(e)
+
+        if ooi.reference.class_type == Hostname:
             path = Path.parse("Hostname.<hostname [is ResolvedHostname].address")
-            ips = self.octopoes_api_connector.query(path=path, source=ref, valid_time=valid_time)
+            ips = self.octopoes_api_connector.query(path=path, source=ooi.reference, valid_time=valid_time)
             if not ips:
                 return {}
             references = [ip.reference for ip in ips]
         else:
-            references = [ref]
+            references = [ooi]
 
         results = {}
         for ref in references:

diff --git a/rocky/reports/report_types/rpki_report/report.py b/rocky/reports/report_types/rpki_report/report.py
@@ -5,6 +5,7 @@
 from django.utils.translation import gettext_lazy as _
 
 from octopoes.models import Reference
+from octopoes.models.exception import ObjectNotFoundException
 from octopoes.models.ooi.dns.zone import Hostname
 from octopoes.models.ooi.network import IPAddressV4, IPAddressV6
 from reports.report_types.definitions import Report
@@ -24,13 +25,18 @@ class RPKIReport(Report):
     template_path = "rpki_report/report.html"
 
     def generate_data(self, input_ooi: str, valid_time: datetime) -> Dict[str, Any]:
-        reference = Reference.from_str(input_ooi)
-        if reference.class_type == Hostname:
+        try:
+            ooi = self.octopoes_api_connector.get(Reference.from_str(input_ooi), valid_time)
+        except ObjectNotFoundException as e:
+            logger.error("No data found for OOI '%s' on date %s.", str(e), str(valid_time))
+            raise ObjectNotFoundException(e)
+
+        if ooi.reference.class_type == Hostname:
             ips = self.octopoes_api_connector.query(
-                "Hostname.<hostname[is ResolvedHostname].address", valid_time, reference
+                "Hostname.<hostname[is ResolvedHostname].address", valid_time, ooi.reference
             )
         else:
-            ips = [self.octopoes_api_connector.get(reference)]
+            ips = [ooi]
 
         rpki_ips = {}
         number_of_ips = len(ips)

diff --git a/rocky/reports/report_types/safe_connections_report/report.py b/rocky/reports/report_types/safe_connections_report/report.py
@@ -1,15 +1,19 @@
 from datetime import datetime
+from logging import getLogger
 from typing import Any, Dict
 
 from django.utils.translation import gettext_lazy as _
 
 from octopoes.models import Reference
+from octopoes.models.exception import ObjectNotFoundException
 from octopoes.models.ooi.dns.zone import Hostname
 from octopoes.models.ooi.network import IPAddressV4, IPAddressV6
 from reports.report_types.definitions import Report
 
 CIPHER_FINDINGS = ["KAT-RECOMMENDATION-BAD-CIPHER", "KAT-MEDIUM-BAD-CIPHER", "KAT-CRITICAL-BAD-CIPHER"]
 
+logger = getLogger(__name__)
+
 
 class SafeConnectionsReport(Report):
     id = "safe-connections-report"
@@ -20,14 +24,18 @@ class SafeConnectionsReport(Report):
     template_path = "safe_connections_report/report.html"
 
     def generate_data(self, input_ooi: str, valid_time: datetime) -> Dict[str, Any]:
-        reference = Reference.from_str(input_ooi)
+        try:
+            ooi = self.octopoes_api_connector.get(Reference.from_str(input_ooi), valid_time)
+        except ObjectNotFoundException as e:
+            logger.error("No data found for OOI '%s' on date %s.", str(e), str(valid_time))
+            raise ObjectNotFoundException(e)
 
-        if reference.class_type == Hostname:
+        if ooi.reference.class_type == Hostname:
             ips = self.octopoes_api_connector.query(
-                "Hostname.<hostname[is ResolvedHostname].address", valid_time, reference
+                "Hostname.<hostname[is ResolvedHostname].address", valid_time, ooi.reference
             )
         else:
-            ips = [self.octopoes_api_connector.get(reference)]
+            ips = [ooi]
 
         sc_ips = {}
         number_of_ips = len(ips)

diff --git a/rocky/reports/report_types/systems_report/report.py b/rocky/reports/report_types/systems_report/report.py
@@ -7,6 +7,7 @@
 from strenum import StrEnum
 
 from octopoes.models import Reference
+from octopoes.models.exception import ObjectNotFoundException
 from octopoes.models.ooi.dns.zone import Hostname
 from octopoes.models.ooi.network import IPAddressV4, IPAddressV6
 from reports.report_types.definitions import Report
@@ -37,15 +38,20 @@ class SystemReport(Report):
     template_path = "systems_report/report.html"
 
     def generate_data(self, input_ooi: str, valid_time: datetime) -> Dict[str, Any]:
-        reference = Reference.from_str(input_ooi)
         ips = []
 
-        if reference.class_type == Hostname:
+        try:
+            ooi = self.octopoes_api_connector.get(Reference.from_str(input_ooi), valid_time)
+        except ObjectNotFoundException as e:
+            logger.error("No data found for OOI '%s' on date %s.", str(e), str(valid_time))
+            raise ObjectNotFoundException(e)
+
+        if ooi.reference.class_type == Hostname:
             ips = self.octopoes_api_connector.query(
-                "Hostname.<hostname[is ResolvedHostname].address", valid_time, reference
+                "Hostname.<hostname[is ResolvedHostname].address", valid_time, ooi.reference
             )
-        elif reference.class_type in (IPAddressV4, IPAddressV6):
-            ips = [self.octopoes_api_connector.get(reference)]
+        elif ooi.reference.class_type in (IPAddressV4, IPAddressV6):
+            ips = [ooi]
 
         ip_services = {}
 

diff --git a/rocky/reports/report_types/vulnerability_report/report.py b/rocky/reports/report_types/vulnerability_report/report.py
@@ -6,6 +6,7 @@
 from django.utils.translation import gettext_lazy as _
 
 from octopoes.models import Reference
+from octopoes.models.exception import ObjectNotFoundException
 from octopoes.models.ooi.dns.zone import Hostname
 from octopoes.models.ooi.findings import Finding, KATFindingType, RiskLevelSeverity
 from octopoes.models.ooi.network import IPAddressV4, IPAddressV6
@@ -31,18 +32,22 @@ def get_finding_valid_time_history(self, reference: str) -> List[datetime]:
         return valid_time_history
 
     def get_findings(self, input_ooi: str, valid_time: datetime) -> Dict[str, Dict[str, List[Union[str, Finding]]]]:
-        reference = Reference.from_str(input_ooi)
-
         aggregated_findings = []
         finding_types = []
         findings_data = {}
 
-        if reference.class_type == Hostname:
+        try:
+            ooi = self.octopoes_api_connector.get(Reference.from_str(input_ooi), valid_time)
+        except ObjectNotFoundException as e:
+            logger.error("No data found for OOI '%s' on date %s.", str(e), str(valid_time))
+            raise ObjectNotFoundException(e)
+
+        if ooi.reference.class_type == Hostname:
             ips = self.octopoes_api_connector.query(
-                "Hostname.<hostname[is ResolvedHostname].address", valid_time, reference
+                "Hostname.<hostname[is ResolvedHostname].address", valid_time, ooi.reference
             )
         else:
-            ips = [self.octopoes_api_connector.get(reference)]
+            ips = [ooi]
 
         for ip in ips:
             findings_ip_port = self.octopoes_api_connector.query(

diff --git a/rocky/reports/report_types/web_system_report/report.py b/rocky/reports/report_types/web_system_report/report.py
@@ -6,6 +6,7 @@
 from django.utils.translation import gettext_lazy as _
 
 from octopoes.models import Reference
+from octopoes.models.exception import ObjectNotFoundException
 from octopoes.models.ooi.dns.zone import Hostname
 from octopoes.models.ooi.findings import KATFindingType, RiskLevelSeverity
 from octopoes.models.ooi.network import IPAddressV4, IPAddressV6
@@ -110,15 +111,20 @@ class WebSystemReport(Report):
     template_path = "web_system_report/report.html"
 
     def generate_data(self, input_ooi: str, valid_time: datetime) -> Dict[str, Any]:
-        reference = Reference.from_str(input_ooi)
         hostnames = []
 
-        if reference.class_type == Hostname:
-            hostnames = [self.octopoes_api_connector.get(reference)]
+        try:
+            ooi = self.octopoes_api_connector.get(Reference.from_str(input_ooi), valid_time)
+        except ObjectNotFoundException as e:
+            logger.error("No data found for OOI '%s' on date %s.", str(e), str(valid_time))
+            raise ObjectNotFoundException(e)
 
-        elif reference.class_type in (IPAddressV4, IPAddressV6):
+        if ooi.reference.class_type == Hostname:
+            hostnames = [ooi]
+
+        elif ooi.reference.class_type in (IPAddressV4, IPAddressV6):
             hostnames = self.octopoes_api_connector.query(
-                "IPAddress.<address[is ResolvedHostname].hostname", valid_time, reference
+                "IPAddress.<address[is ResolvedHostname].hostname", valid_time, ooi.reference
             )
 
         web_checks = WebChecks()