diff --git a/stream_alert/alert_processor/outputs/pagerduty.py b/stream_alert/alert_processor/outputs/pagerduty.py index 004cc7c3f..03642dfc3 100644 --- a/stream_alert/alert_processor/outputs/pagerduty.py +++ b/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,48 @@ 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 + """ + if not context: + return dict() + + 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 +444,9 @@ 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 = 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 +464,7 @@ def dispatch(self, **kwargs): 'type': 'incident', 'title': incident_title, 'service': incident_service, + 'priority': incident_priority, 'body': incident_body, assigned_key: assigned_value } diff --git a/tests/unit/stream_alert_alert_processor/test_outputs/test_pagerduty.py b/tests/unit/stream_alert_alert_processor/test_outputs/test_pagerduty.py index cd120d104..1f427e5e5 100644 --- a/tests/unit/stream_alert_alert_processor/test_outputs/test_pagerduty.py +++ b/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')