Add risk rating mapping for tsunami.

Ostorlab · Feb 26, 2024 · 0156641 · 0156641
1 parent e4d8361
commit 0156641
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 2 deletions.
diff --git a/agent/tsunami_agent.py b/agent/tsunami_agent.py
@@ -31,6 +31,16 @@
 )
 logger = logging.getLogger(__name__)
 
+# severity mapping defined in https://github.com/google/tsunami-security-scanner/blob/master/proto/vulnerability.proto
+RISK_MAPPING = {
+    "CRITICAL": "CRITICAL",
+    "HIGH": "HIGH",
+    "MEDIUM": "MEDIUM",
+    "LOW": "LOW",
+    "MINIMAL": "INFO",
+    "SEVERITY_UNSPECIFIED": "INFO",
+}
+
 
 class AgentTsunami(
     agent.Agent,
@@ -175,7 +185,7 @@ def _report_vulnerability(
         vuln_location: agent_report_vulnerability_mixin.VulnerabilityLocation,
     ) -> None:
         # risk_rating will be HIGH for all detected vulnerabilities
-        risk_rating = "HIGH"
+        risk_rating = RISK_MAPPING[vulnerability["vulnerability"]["severity"]]
         self.report_vulnerability(
             entry=kb.Entry(
                 title=vulnerability["vulnerability"]["title"],

diff --git a/tests/agent_test.py b/tests/agent_test.py
@@ -66,11 +66,12 @@ def testTsunamiAgent_WhenTsunamiScanHasVulnerabilities_ShouldReportVulnerabiliti
                 "vulnerability": {
                     "title": "Ostorlab Platform",
                     "description": "Ostorlab is not password protected",
+                    "severity": "CRITICAL",
                 }
             }
         ],
     }
-    risk_rating = "HIGH"
+    risk_rating = "CRITICAL"
     description = "Ostorlab is not password protected"
     kb_entry = kb.Entry(
         title="Ostorlab Platform",
@@ -123,6 +124,7 @@ def testTsunamiAgent_WhenLinkAssetAndTsunamiScanHasVulnerabilities_ShouldReportV
                 "vulnerability": {
                     "title": "Ostorlab Platform",
                     "description": "Ostorlab is not password protected",
+                    "severity": "HIGH",
                 }
             }
         ],
@@ -185,6 +187,7 @@ def testTsunamiAgent_WhenServiceAssetAndTsunamiScanHasVulnerabilities_ShouldRepo
                 "vulnerability": {
                     "title": "Ostorlab Platform",
                     "description": "Ostorlab is not password protected",
+                    "severity": "HIGH",
                 }
             }
         ],
@@ -247,6 +250,7 @@ def testTsunamiAgent_WhenDomainNameAssetAndTsunamiScanHasVulnerabilities_ShouldR
                 "vulnerability": {
                     "title": "Ostorlab Platform",
                     "description": "Ostorlab is not password protected",
+                    "severity": "HIGH",
                 }
             }
         ],
@@ -347,6 +351,7 @@ def testAgentTsunami_whenIpRangeScanned_emitsExactIpWhereVulnWasFound(
                 "vulnerability": {
                     "title": "Ostorlab Platform",
                     "description": "Ostorlab is not password protected",
+                    "severity": "HIGH",
                 }
             }
         ],