Implement methods for dismiss and updated authentication

threatstack/base.py

@@ -47,11 +47,11 @@ def api_key(self, k):
             raise errors.ThreatStackClientError("api_key is required")
         self._api_key = k
 
-    def request_headers(self, method, full_url):
+    def request_headers(self, method, full_url, data, content_type):
         raise NotImplementedError("Please Implement this method")
 
     @retry(**RETRY_OPTS)
-    def http_request(self, method, path, data=None, params=None):
+    def http_request(self, method, path, data=None, params=None, content_type=None):
         """ Wraps HTTP calls to ThrestStack API """
 
         s = Session()
@@ -62,7 +62,7 @@ def http_request(self, method, path, data=None, params=None):
         except:
             pass
 
-        headers = self.request_headers(method, full_url)
+        headers = self.request_headers(method, full_url, data, content_type)
 
         req = Request(
             method,

threatstack/v1/client.py

@@ -21,7 +21,7 @@ def __init__(self, api_key=None, org_id=None, user_id=None, timeout=None):
         self.organizations = resources.Organizations(self)
         self.policies = resources.Policies(self)
 
-    def request_headers(self, _method, _url):
+    def request_headers(self, _method, _url, _data, _content_type):
         headers = { "Authorization": self.api_key }
         if self.org_id:
             headers["Organization"] = self.org_id

threatstack/v2/client.py

@@ -42,14 +42,18 @@ def user_id(self, k):
             raise ThreatStackClientError("user_id is required.")
         self._user_id = k
 
-    def request_headers(self, method, url):
+    def request_headers(self, method, url, data, content_type):
         credentials = {
             'id': self.user_id,
             'key': self.api_key,
             'algorithm': 'sha256'
         }
+        if method == 'POST':
+            sender = Sender(credentials, url, method, always_hash_content=False, ext=self.org_id, content=data,
+                            content_type=content_type)
+            return {'Authorization': sender.request_header, 'Content-Type': content_type}
         sender = Sender(credentials, url, method, always_hash_content=False, ext=self.org_id)
-        return { 'Authorization': sender.request_header }
+        return {'Authorization': sender.request_header}
 
     def handle_response(self, resp):
         # ThreatStack can return various things

threatstack/v2/resources.py

@@ -95,6 +95,56 @@ def events(self, alert_id=""):
         resp = self.client.http_request("GET", path)
         return resp
 
+    def dismiss_by_id(self, ids, dismiss_reason, dismiss_reason_text=None):
+        """
+        :param ids: A list of valid alertIds. minItems:1 maxItems: 512
+        :param dismiss_reason: The reason the alert was dismissed. Allowed Values: business_op, company_policy,
+                               maintenance, none, auto_dismiss, OTHER
+        :param dismiss_reason_text: Reason the alert was dismissed if reason is other.
+        :return:
+        """
+        path = "{}/dismiss".format(self.name)
+        data = dict()
+        data["ids"] = ids
+        data["dismissReason"] = dismiss_reason
+        if dismiss_reason_text is not None:
+            data["dismissReasonText"] = dismiss_reason_text
+        resp = self.client.http_request("POST", path, data=json.dumps(data), content_type='application/json')
+        return resp
+
+    def dismiss_by_time_range(self, alerts_from, alerts_until, dismiss_reason, severity=None, rule_id=None,
+                              agent_id=None, dismiss_reason_text=None):
+        """
+        :param alerts_from: The first date and time from which to dismiss alerts.
+                            With alerts_until forms a timeframe within which alerts will be dismissed.
+                            Format: ISO-8601 date and time
+        :param alerts_until: The last date and time from which to dismiss alerts.
+                             With alters_from, forms a timeframe within which alerts will be dismissed.
+                             Format: ISO-8601 date and time
+        :param dismiss_reason: The reason the alert was dismissed. Allowed Values: business_op, company_policy,
+                               maintenance, none, auto_dismiss, OTHER
+        :param severity: Severity level of the alert. Allowed Values: 1, 2, 3
+        :param rule_id: Unique id of the rule that generated the alert.
+        :param agent_id: Unique id of the Agent on which the alert occurred.
+        :param dismiss_reason_text: Reason the alert was dismissed if reason is other.
+        :return:
+        """
+        path = "{}/dismiss".format(self.name)
+        data = dict()
+        data["from"] = alerts_from
+        data["until"] = alerts_until
+        if severity is not None:
+            data["severity"] = severity
+        if rule_id is not None:
+            data["ruleID"] = rule_id
+        if agent_id is not None:
+            data["agentID"] = agent_id
+        data["dismissReason"] = dismiss_reason
+        if dismiss_reason_text is not None:
+            data["dismissReasonText"] = dismiss_reason_text
+        resp = self.client.http_request("POST", path, data=json.dumps(data), content_type='application/json')
+        return resp
+
 
 class Vulnerabilities(Resource):
     def list(self, active=None, **kwargs):