[fix] Top 5 failing (#640)

someengineering · Oct 29, 2024 · 36e3109 · 36e3109
1 parent c5717bf
commit 36e3109
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 17 deletions.
diff --git a/fixbackend/inventory/inventory_service.py b/fixbackend/inventory/inventory_service.py
@@ -589,12 +589,9 @@ async def issues_since(
                 )
 
             async def benchmark_summary(
-                benmark_account_summaries: Dict[BenchmarkId, Dict[CloudAccountId, BenchmarkAccountSummary]]
-            ) -> Tuple[
-                BenchmarkById, Dict[ReportSeverity, Set[SecurityCheckId]], Dict[SecurityCheckId, Set[BenchmarkId]]
-            ]:
+                bench_account_summaries: Dict[BenchmarkId, Dict[CloudAccountId, BenchmarkAccountSummary]]
+            ) -> Tuple[BenchmarkById, Dict[SecurityCheckId, Set[BenchmarkId]]]:
                 summaries: BenchmarkById = {}
-                failed_checks_by_severity: Dict[ReportSeverity, Set[SecurityCheckId]] = defaultdict(set)
                 benchmark_by_check_id: Dict[SecurityCheckId, Set[BenchmarkId]] = defaultdict(set)
                 benchmarks = await self.client.benchmarks(db, short=True, with_checks=True)
                 for b in benchmarks:
@@ -607,15 +604,13 @@ async def benchmark_summary(
                         clouds=b["clouds"],
                         description=b["description"],
                         nr_of_checks=len(b["report_checks"]),
-                        account_summary=benmark_account_summaries.get(benchmark_id, {}),
+                        account_summary=bench_account_summaries.get(benchmark_id, {}),
                     )
                     summaries[summary.id] = summary
                     for check in b["report_checks"]:
                         check_id = SecurityCheckId(check["id"])
-                        severity = ReportSeverity(check["severity"])
-                        failed_checks_by_severity[severity].add(check_id)
                         benchmark_by_check_id[check_id].add(benchmark_id)
-                return summaries, failed_checks_by_severity, benchmark_by_check_id
+                return summaries, benchmark_by_check_id
 
             async def timeseries_infected() -> TimeSeries:
                 start = now - timedelta(days=62 if is_free else 14)
@@ -628,13 +623,21 @@ async def timeseries_infected() -> TimeSeries:
                 return TimeSeries(name="infected_resources", start=start, end=now, granularity=granularity, data=data)
 
             async def top_issues(
-                checks_by_severity: Dict[ReportSeverity, Set[SecurityCheckId]],
                 benchmark_by_check_id: Dict[SecurityCheckId, Set[BenchmarkId]],
                 benchmarks: Dict[BenchmarkId, BenchmarkSummary],
                 num: int,
             ) -> List[Json]:
-                check_ids = dict_values_by(checks_by_severity, lambda x: ReportSeverityPriority[x])
-                top = list(islice(check_ids, num))
+                query = (
+                    "aggregate(/security.issues[*].check, /security.issues[*].severity: sum(1) as count): "
+                    "/security.has_issues==true"
+                )
+                async with self.client.aggregate(db, query) as ctx:
+                    all_failing = sorted(
+                        [e async for e in ctx],
+                        key=lambda x: (ReportSeverityPriority[x["group"]["severity"]], x["count"]),
+                        reverse=True,
+                    )
+                top = list(islice((a["group"]["check"] for a in all_failing), num))
                 checks = await self.client.checks(db, check_ids=top)
                 for check in checks:
                     check["benchmarks"] = [
@@ -657,7 +660,7 @@ def overall_score(
                 return total_score // total_accounts if total_accounts > 0 else 100
 
             (
-                (benchmarks, failed_checks_by_severity, benchmark_by_check_id),
+                (benchmarks, benchmark_by_check_id),
                 vulnerable_changed,
                 compliant_changed,
                 infected_resources_ts,
@@ -669,7 +672,7 @@ def overall_score(
             )
 
             # get issues for the top 5 issue_ids
-            tops = await top_issues(failed_checks_by_severity, benchmark_by_check_id, benchmarks, num=5)
+            tops = await top_issues(benchmark_by_check_id, benchmarks, num=5)
 
             # sort top changed account by score
             vulnerable_changed.accounts_selection.sort(

diff --git a/tests/fixbackend/inventory/inventory_service_test.py b/tests/fixbackend/inventory/inventory_service_test.py
@@ -130,10 +130,10 @@ async def mock(request: Request) -> Response:
             )
         elif request.url.path == "/graph/fix/search/list" and content.startswith("search is(account) and /metadata.exported_at>="):  # fmt: skip
             return nd_json_response(accounts_json)
-        elif request.url.path == "/graph/fix/search/aggregate" and content.startswith("search /security.has_issues==true"):  # fmt: skip
+        elif request.url.path == "/graph/fix/search/aggregate" and "/security.has_issues==true" in content:  # fmt: skip
             return nd_json_response(
-                [{"group": {"check_id": "aws_c1", "severity": "low", "account_id": "123", "account_name": "t1", "cloud": "aws"}, "count": 8},  # fmt: skip
-                 {"group": {"check_id": "gcp_c2", "severity": "critical", "account_id": "234", "account_name": "t2", "cloud": "gcp"}, "count": 2}]  # fmt: skip
+                [{"group": {"check": "aws_c1", "severity": "low"}, "count": 8},  # fmt: skip
+                 {"group": {"check": "gcp_c2", "severity": "critical"}, "count": 2}]  # fmt: skip
             )
         elif request.url.path == "/graph/fix/node/some_node_id":
             return json_response(azure_virtual_machine_resource_json)