From 48db77b521a91ded220e84dcd360250a5e52e45c Mon Sep 17 00:00:00 2001
From: Joseph Schorr <josephschorr@users.noreply.github.com>
Date: Fri, 10 Mar 2017 11:25:55 -0500
Subject: [PATCH] Fix bug in QSS notifications

---
 endpoints/notificationevent.py | 17 ++++++++++++-----
 test/test_secscan.py           |  1 +
 2 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/endpoints/notificationevent.py b/endpoints/notificationevent.py
index cbba64fd0..94e795b64 100644
--- a/endpoints/notificationevent.py
+++ b/endpoints/notificationevent.py
@@ -112,6 +112,7 @@ def _build_summary(event_data):
 
 class VulnerabilityFoundEvent(NotificationEvent):
   CONFIG_LEVEL = 'level'
+  PRIORITY_KEY = 'priority'
   VULNERABILITY_KEY = 'vulnerability'
   MULTIPLE_VULNERABILITY_KEY = 'vulnerabilities'
 
@@ -120,7 +121,8 @@ class VulnerabilityFoundEvent(NotificationEvent):
     return 'vulnerability_found'
 
   def get_level(self, event_data, notification_data):
-    priority = event_data[VulnerabilityFoundEvent.CONFIG_LEVEL]['priority']
+    vuln_data = event_data[VulnerabilityFoundEvent.VULNERABILITY_KEY]
+    priority = vuln_data[VulnerabilityFoundEvent.PRIORITY_KEY]
     if priority == 'Defcon1' or priority == 'Critical':
       return 'error'
 
@@ -138,6 +140,8 @@ class VulnerabilityFoundEvent(NotificationEvent):
       'namespace_name': notification.repository.namespace_user.username,
       'name': notification.repository.name,
     })
+
+    level = event_config.get(VulnerabilityFoundEvent.CONFIG_LEVEL, 'Critical')
     return build_event_data(repo, {
       'tags': ['latest', 'prod', 'foo', 'bar', 'baz'],
       'image': 'some-image-id',
@@ -145,7 +149,7 @@ class VulnerabilityFoundEvent(NotificationEvent):
         'id': 'CVE-FAKE-CVE',
         'description': 'A futurist vulnerability',
         'link': 'https://security-tracker.debian.org/tracker/CVE-FAKE-CVE',
-        'priority': get_priority_for_index(event_config[VulnerabilityFoundEvent.CONFIG_LEVEL])
+        'priority': get_priority_for_index(level)
       },
     })
 
@@ -167,10 +171,13 @@ class VulnerabilityFoundEvent(NotificationEvent):
     return actual_level_index <= filter_level_index
 
   def get_summary(self, event_data, notification_data):
+    vuln_key = VulnerabilityFoundEvent.VULNERABILITY_KEY
+    priority_key = VulnerabilityFoundEvent.PRIORITY_KEY
+
     multiple_vulns = event_data.get(VulnerabilityFoundEvent.MULTIPLE_VULNERABILITY_KEY)
     if multiple_vulns is not None:
-      top_priority = multiple_vulns[0].get('priority', 'Unknown')
-      matching = [v for v in multiple_vulns if v.get('priority', 'Unknown') == top_priority]
+      top_priority = multiple_vulns[0].get(priority_key, 'Unknown')
+      matching = [v for v in multiple_vulns if v.get(priority_key, 'Unknown') == top_priority]
 
       msg = '%s %s' % (len(matching), top_priority)
       if len(matching) < len(multiple_vulns):
@@ -180,7 +187,7 @@ class VulnerabilityFoundEvent(NotificationEvent):
       return msg % (event_data['repository'], len(event_data['tags']))
     else:
       msg = '%s vulnerability detected in repository %s in %s tags'
-      return msg % (event_data['vulnerability']['priority'], event_data['repository'],
+      return msg % (event_data[vuln_key][priority_key], event_data['repository'],
                     len(event_data['tags']))
 
 
diff --git a/test/test_secscan.py b/test/test_secscan.py
index cc29e5d74..b5da8e2b6 100644
--- a/test/test_secscan.py
+++ b/test/test_secscan.py
@@ -355,6 +355,7 @@ class TestSecurityScanner(unittest.TestCase):
       event = VulnerabilityFoundEvent()
       msg = '1 Low and 1 more vulnerabilities were detected in repository devtable/simple in 2 tags'
       self.assertEquals(msg, event.get_summary(body['event_data'], {}))
+      self.assertEquals('info', event.get_level(body['event_data'], {}))
     else:
       self.assertIsNone(queue_item)