diff --git a/threatstack/base.py b/threatstack/base.py index f5f7474..1f9c5ae 100644 --- a/threatstack/base.py +++ b/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, diff --git a/threatstack/v1/client.py b/threatstack/v1/client.py index 5f140a8..d42ac12 100644 --- a/threatstack/v1/client.py +++ b/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 diff --git a/threatstack/v2/client.py b/threatstack/v2/client.py index 5f67938..2299f61 100644 --- a/threatstack/v2/client.py +++ b/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 diff --git a/threatstack/v2/resources.py b/threatstack/v2/resources.py index 313c100..d5e0445 100644 --- a/threatstack/v2/resources.py +++ b/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):