[output] Adding support for assign priority to incidents

stream_alert/alert_processor/outputs/pagerduty.py

@@ -176,6 +176,7 @@ class PagerDutyIncidentOutput(OutputDispatcher):
     USERS_ENDPOINT = 'users'
     POLICIES_ENDPOINT = 'escalation_policies'
     SERVICES_ENDPOINT = 'services'
+    PRIORITIES_ENDPOINT = 'priorities'
 
     def __init__(self, *args, **kwargs):
         OutputDispatcher.__init__(self, *args, **kwargs)
@@ -336,6 +337,45 @@ def _item_verify(self, item_str, item_key, item_type, get_id=True):
 
         return item_id
 
+    def _priority_verify(self, context):
+        """Method to verify the existance of a incident priority with the API
+
+        Args:
+            context (dict): Context provided in the alert record
+
+        Returns:
+            dict: JSON object be used in the API call, containing the priority id
+                  and the priority reference, empty if it fails or it does not exist
+        """
+        priority_name = context.get('incident_priority', False)
+        if not priority_name:
+            return dict()
+
+        priorities_url = self._get_endpoint(self._base_url, self.PRIORITIES_ENDPOINT)
+        resp = self._get_request(priorities_url, {}, self._headers, False)
+
+        if not self._check_http_response(resp):
+            return dict()
+
+        response = resp.json()
+        if not response:
+            return dict()
+
+        priorities = response.get('priorities', [])
+
+        if not priorities:
+            return dict()
+
+        # If the requested priority is in the list, get the id
+        priority_id = next(
+            (item for item in priorities if item["name"] == priority_name), {}).get('id', False)
+
+        # If the priority id is found, compose the JSON
+        if priority_id:
+            return {'id': priority_id, 'type': 'priority_reference'}
+
+        return dict()
+
     def _incident_assignment(self, context):
         """Method to determine if the incident gets assigned to a user or an escalation policy
 
@@ -401,6 +441,11 @@ def dispatch(self, **kwargs):
         if rule_context:
             rule_context = rule_context.get(self.__service__, {})
 
+        # Use the priority provided in the context, use it or the incident will be low priority
+        incident_priority = {}
+        if rule_context:
+            incident_priority = self._priority_verify(rule_context)
+
         # Incident assignment goes in this order:
         #  Provided user -> provided policy -> default policy
         assigned_key, assigned_value = self._incident_assignment(rule_context)
@@ -418,6 +463,7 @@ def dispatch(self, **kwargs):
                 'type': 'incident',
                 'title': incident_title,
                 'service': incident_service,
+                'priority': incident_priority,
                 'body': incident_body,
                 assigned_key: assigned_value
             }

tests/unit/stream_alert_alert_processor/test_outputs/test_pagerduty.py

@@ -143,6 +143,7 @@ def test_dispatch_bad_descriptor(self, log_mock):
         log_mock.assert_called_with('Failed to send alert to %s', self.SERVICE)
 
 
+#pylint: disable=too-many-public-methods
 @mock_s3
 @mock_kms
 class TestPagerDutyIncidentOutput(object):
@@ -315,6 +316,71 @@ def test_item_verify_no_get_id_success(self, get_mock):
 
         assert_true(self._dispatcher._item_verify('valid_item', 'items', 'item_reference', False))
 
+    @patch('requests.get')
+    def test_priority_verify_success(self, get_mock):
+        """PagerDutyIncidentOutput - Priority Verify Success"""
+        priority_name = 'priority_name'
+        # /priorities
+        get_mock.return_value.status_code = 200
+        json_check = {'priorities': [{'id': 'verified_priority_id', 'name': priority_name}]}
+        get_mock.return_value.json.return_value = json_check
+
+        context = {'incident_priority': priority_name}
+
+        priority_verified = self._dispatcher._priority_verify(context)
+        assert_equal(priority_verified['id'], 'verified_priority_id')
+        assert_equal(priority_verified['type'], 'priority_reference')
+
+    @patch('requests.get')
+    def test_priority_verify_fail(self, get_mock):
+        """PagerDutyIncidentOutput - Priority Verify Fail"""
+        # /priorities
+        get_mock.return_value.status_code = 404
+
+        context = {'incident_priority': 'priority_name'}
+
+        priority_not_verified = self._dispatcher._priority_verify(context)
+        assert_equal(priority_not_verified, dict())
+
+    @patch('requests.get')
+    def test_priority_verify_empty(self, get_mock):
+        """PagerDutyIncidentOutput - Priority Verify Empty"""
+        # /priorities
+        get_mock.return_value.status_code = 200
+        json_check = {}
+        get_mock.return_value.json.return_value = json_check
+
+        context = {'incident_priority': 'priority_name'}
+
+        priority_not_verified = self._dispatcher._priority_verify(context)
+        assert_equal(priority_not_verified, dict())
+
+    @patch('requests.get')
+    def test_priority_verify_not_found(self, get_mock):
+        """PagerDutyIncidentOutput - Priority Verify Not Found"""
+        # /priorities
+        get_mock.return_value.status_code = 200
+        json_check = {'priorities': [{'id': 'verified_priority_id', 'name': 'not_priority_name'}]}
+        get_mock.return_value.json.return_value = json_check
+
+        context = {'incident_priority': 'priority_name'}
+
+        priority_not_verified = self._dispatcher._priority_verify(context)
+        assert_equal(priority_not_verified, dict())
+
+    @patch('requests.get')
+    def test_priority_verify_invalid(self, get_mock):
+        """PagerDutyIncidentOutput - Priority Verify Invalid"""
+        # /priorities
+        get_mock.return_value.status_code = 200
+        json_check = {'not_priorities': [{'id': 'verified_priority_id', 'name': 'priority_name'}]}
+        get_mock.return_value.json.return_value = json_check
+
+        context = {'incident_priority': 'priority_name'}
+
+        priority_not_verified = self._dispatcher._priority_verify(context)
+        assert_equal(priority_not_verified, dict())
+
     @patch('requests.get')
     def test_incident_assignment_user(self, get_mock):
         """PagerDutyIncidentOutput - Incident Assignment User"""
@@ -415,6 +481,36 @@ def test_dispatch_success_good_policy(self, get_mock, post_mock, log_mock):
 
         log_mock.assert_called_with('Successfully sent alert to %s', self.SERVICE)
 
+    @patch('logging.Logger.info')
+    @patch('requests.post')
+    @patch('requests.get')
+    def test_dispatch_success_with_priority(self, get_mock, post_mock, log_mock):
+        """PagerDutyIncidentOutput - Dispatch Success With Priority"""
+         # /priorities, /users, /escalation_policies, /services
+        type(get_mock.return_value).status_code = PropertyMock(side_effect=[200, 200, 200, 200])
+        json_user = {'users': [{'id': 'user_id'}]}
+        json_priority = {'priorities': [{'id': 'priority_id', 'name': 'priority_name'}]}
+        json_policy = {'escalation_policies': [{'id': 'policy_id'}]}
+        json_service = {'services': [{'id': 'service_id'}]}
+        get_mock.return_value.json.side_effect = [json_user, json_priority,
+                                                  json_policy, json_service]
+
+        # /incidents
+        post_mock.return_value.status_code = 200
+
+        ctx = {
+            'pagerduty-incident': {
+                'assigned_policy': 'valid_policy',
+                'incident_priority': 'priority_name'
+            }
+        }
+
+        assert_true(self._dispatcher.dispatch(descriptor=self.DESCRIPTOR,
+                                              rule_name='rule_name',
+                                              alert=get_alert(context=ctx)))
+
+        log_mock.assert_called_with('Successfully sent alert to %s', self.SERVICE)
+
     @patch('logging.Logger.info')
     @patch('requests.post')
     @patch('requests.get')