prelude-sdk 2.6.38__py3-none-any.whl → 2.6.40__py3-none-any.whl

prelude_sdk/controllers/build_controller.py

@@ -14,10 +14,7 @@ class BuildController(HttpController):
     def clone_test(self, source_test_id):
         """Clone a test"""
         res = self.post(
-            f"{self.account.hq}/build/tests",
-            json=dict(source_test_id=source_test_id),
-            headers=self.account.headers,
-            timeout=10,
+            f"{self.account.hq}/build/tests", json=dict(source_test_id=source_test_id)
         )
         return res.json()
 
@@ -30,12 +27,7 @@ class BuildController(HttpController):
         if test_id:
             body["id"] = test_id
 
-        res = self.post(
-            f"{self.account.hq}/build/tests",
-            json=body,
-            headers=self.account.headers,
-            timeout=10,
-        )
+        res = self.post(f"{self.account.hq}/build/tests", json=body)
         return res.json()
 
     @verify_credentials
@@ -58,33 +50,21 @@ class BuildController(HttpController):
         if technique is not None:
             body["technique"] = technique
 
-        res = self.post(
-            f"{self.account.hq}/build/tests/{test_id}",
-            json=body,
-            headers=self.account.headers,
-            timeout=10,
-        )
+        res = self.post(f"{self.account.hq}/build/tests/{test_id}", json=body)
         return res.json()
 
     @verify_credentials
     def delete_test(self, test_id, purge):
         """Delete an existing test"""
         res = self.delete(
-            f"{self.account.hq}/build/tests/{test_id}",
-            json=dict(purge=purge),
-            headers=self.account.headers,
-            timeout=10,
+            f"{self.account.hq}/build/tests/{test_id}", json=dict(purge=purge)
         )
         return res.json()
 
     @verify_credentials
     def undelete_test(self, test_id):
         """Undelete a tombstoned test"""
-        res = self.post(
-            f"{self.account.hq}/build/tests/{test_id}/undelete",
-            headers=self.account.headers,
-            timeout=10,
-        )
+        res = self.post(f"{self.account.hq}/build/tests/{test_id}/undelete")
         return res.json()
 
     @verify_credentials
@@ -97,11 +77,11 @@ class BuildController(HttpController):
         query_params = ""
         if skip_compile:
             query_params = "?" + urllib.parse.urlencode(dict(skip_compile=True))
+
         res = self.post(
             f"{self.account.hq}/build/tests/{test_id}/{filename}{query_params}",
             data=data,
             headers=h,
-            timeout=10,
         )
         return res.json()
 
@@ -111,18 +91,12 @@ class BuildController(HttpController):
         res = self.post(
             f"{self.account.hq}/build/compile",
             json=dict(code=code, source_test_id=source_test_id),
-            headers=self.account.headers,
-            timeout=10,
         )
         return res.json()
 
     @verify_credentials
     def get_compile_status(self, job_id):
-        res = self.get(
-            f"{self.account.hq}/build/compile/{job_id}",
-            headers=self.account.headers,
-            timeout=10,
-        )
+        res = self.get(f"{self.account.hq}/build/compile/{job_id}")
         return res.json()
 
     @verify_credentials
@@ -140,12 +114,7 @@ class BuildController(HttpController):
         if tests:
             body["tests"] = tests
 
-        res = self.post(
-            f"{self.account.hq}/build/threats",
-            json=body,
-            headers=self.account.headers,
-            timeout=10,
-        )
+        res = self.post(f"{self.account.hq}/build/threats", json=body)
         return res.json()
 
     @verify_credentials
@@ -171,12 +140,7 @@ class BuildController(HttpController):
         if tests is not None:
             body["tests"] = tests
 
-        res = self.post(
-            f"{self.account.hq}/build/threats/{threat_id}",
-            json=body,
-            headers=self.account.headers,
-            timeout=10,
-        )
+        res = self.post(f"{self.account.hq}/build/threats/{threat_id}", json=body)
         return res.json()
 
     @verify_credentials
@@ -185,19 +149,13 @@ class BuildController(HttpController):
         res = self.delete(
             f"{self.account.hq}/build/threats/{threat_id}",
             json=dict(purge=purge),
-            headers=self.account.headers,
-            timeout=10,
         )
         return res.json()
 
     @verify_credentials
     def undelete_threat(self, threat_id):
         """Undelete a tombstoned threat"""
-        res = self.post(
-            f"{self.account.hq}/build/threats/{threat_id}/undelete",
-            headers=self.account.headers,
-            timeout=10,
-        )
+        res = self.post(f"{self.account.hq}/build/threats/{threat_id}/undelete")
         return res.json()
 
     @verify_credentials
@@ -211,12 +169,7 @@ class BuildController(HttpController):
         if rule_id:
             body["rule_id"] = rule_id
 
-        res = self.post(
-            f"{self.account.hq}/build/detections",
-            json=body,
-            headers=self.account.headers,
-            timeout=10,
-        )
+        res = self.post(f"{self.account.hq}/build/detections", json=body)
         return res.json()
 
     @verify_credentials
@@ -228,22 +181,13 @@ class BuildController(HttpController):
         if test_id:
             body["test_id"] = test_id
 
-        res = self.post(
-            f"{self.account.hq}/build/detections/{detection_id}",
-            json=body,
-            headers=self.account.headers,
-            timeout=10,
-        )
+        res = self.post(f"{self.account.hq}/build/detections/{detection_id}", json=body)
         return res.json()
 
     @verify_credentials
     def delete_detection(self, detection_id: str):
         """Delete an existing detection"""
-        res = self.delete(
-            f"{self.account.hq}/build/detections/{detection_id}",
-            headers=self.account.headers,
-            timeout=10,
-        )
+        res = self.delete(f"{self.account.hq}/build/detections/{detection_id}")
         return res.json()
 
     @verify_credentials
@@ -265,12 +209,7 @@ class BuildController(HttpController):
         if threat_hunt_id:
             body["id"] = threat_hunt_id
 
-        res = self.post(
-            f"{self.account.hq}/build/threat_hunts",
-            json=body,
-            headers=self.account.headers,
-            timeout=10,
-        )
+        res = self.post(f"{self.account.hq}/build/threat_hunts", json=body)
         threat_hunt = res.json()
         if self.account.resolve_enums:
             self.resolve_enums(threat_hunt, [(Control, "control")])
@@ -294,10 +233,7 @@ class BuildController(HttpController):
             body["test_id"] = test_id
 
         res = self.post(
-            f"{self.account.hq}/build/threat_hunts/{threat_hunt_id}",
-            json=body,
-            headers=self.account.headers,
-            timeout=10,
+            f"{self.account.hq}/build/threat_hunts/{threat_hunt_id}", json=body
         )
         threat_hunt = res.json()
         if self.account.resolve_enums:
@@ -307,9 +243,5 @@ class BuildController(HttpController):
     @verify_credentials
     def delete_threat_hunt(self, threat_hunt_id: str):
         """Delete an existing threat hunt"""
-        res = self.delete(
-            f"{self.account.hq}/build/threat_hunts/{threat_hunt_id}",
-            headers=self.account.headers,
-            timeout=10,
-        )
+        res = self.delete(f"{self.account.hq}/build/threat_hunts/{threat_hunt_id}")
         return res.json()

prelude_sdk/controllers/detect_controller.py

@@ -19,7 +19,6 @@ class DetectController(HttpController):
             f"{self.account.hq}/detect/endpoint",
             headers=dict(account=account, token=token, _product="py-sdk"),
             json=body,
-            timeout=10,
         )
         return res.text
 
@@ -30,36 +29,23 @@ class DetectController(HttpController):
         if tags is not None:
             body["tags"] = tags
 
-        res = self.post(
-            f"{self.account.hq}/detect/endpoint/{endpoint_id}",
-            headers=self.account.headers,
-            json=body,
-            timeout=10,
-        )
+        res = self.post(f"{self.account.hq}/detect/endpoint/{endpoint_id}", json=body)
         return res.json()
 
     @verify_credentials
     def delete_endpoint(self, ident: str):
         """Delete an endpoint from your account"""
         params = dict(id=ident)
-        res = self.delete(
-            f"{self.account.hq}/detect/endpoint",
-            headers=self.account.headers,
-            json=params,
-            timeout=10,
-        )
+
+        res = self.delete(f"{self.account.hq}/detect/endpoint", json=params)
         return res.json()
 
     @verify_credentials
     def list_endpoints(self, days: int = 90):
         """List all endpoints on your account"""
         params = dict(days=days)
-        res = self.get(
-            f"{self.account.hq}/detect/endpoint",
-            headers=self.account.headers,
-            params=params,
-            timeout=10,
-        )
+
+        res = self.get(f"{self.account.hq}/detect/endpoint", params=params)
         endpoints = res.json()
         if self.account.resolve_enums:
             self.resolve_enums(endpoints, [(Control, "control")])
@@ -69,57 +55,38 @@ class DetectController(HttpController):
     def describe_activity(self, filters: dict, view: str = "protected"):
         """Get report for an Account"""
         params = dict(view=view, **filters)
-        res = self.get(
-            f"{self.account.hq}/detect/activity",
-            headers=self.account.headers,
-            params=params,
-            timeout=10,
-        )
+
+        res = self.get(f"{self.account.hq}/detect/activity", params=params)
         return res.json()
 
     @verify_credentials
     def threat_hunt_activity(self, threat_hunt_id=None, test_id=None, threat_id=None):
         """Get threat hunt activity"""
-        filters = dict(
+        params = dict(
             threat_hunt_id=threat_hunt_id, test_id=test_id, threat_id=threat_id
         )
-        res = self.get(
-            f"{self.account.hq}/detect/threat_hunt_activity",
-            headers=self.account.headers,
-            params=filters,
-            timeout=10,
-        )
+
+        res = self.get(f"{self.account.hq}/detect/threat_hunt_activity", params=params)
         return res.json()
 
     @verify_credentials
     def list_tests(self, filters: dict = None):
         """List all tests available to an account"""
         res = self.get(
-            f"{self.account.hq}/detect/tests",
-            headers=self.account.headers,
-            params=filters if filters else {},
-            timeout=10,
+            f"{self.account.hq}/detect/tests", params=filters if filters else {}
         )
         return res.json()
 
     @verify_credentials
     def get_test(self, test_id):
         """Get properties of an existing test"""
-        res = self.get(
-            f"{self.account.hq}/detect/tests/{test_id}",
-            headers=self.account.headers,
-            timeout=10,
-        )
+        res = self.get(f"{self.account.hq}/detect/tests/{test_id}")
         return res.json()
 
     @verify_credentials
     def list_techniques(self):
         """List techniques"""
-        res = self.get(
-            f"{self.account.hq}/detect/techniques",
-            headers=self.account.headers,
-            timeout=10,
-        )
+        res = self.get(f"{self.account.hq}/detect/techniques")
         return res.json()
 
     @verify_credentials
@@ -127,51 +94,32 @@ class DetectController(HttpController):
         """List threats"""
         res = self.get(
             f"{self.account.hq}/detect/threats",
-            headers=self.account.headers,
-            params={},
-            timeout=10,
         )
         return res.json()
 
     @verify_credentials
     def get_threat(self, threat_id):
         """Get properties of an existing threat"""
-        res = self.get(
-            f"{self.account.hq}/detect/threats/{threat_id}",
-            headers=self.account.headers,
-            timeout=10,
-        )
+        res = self.get(f"{self.account.hq}/detect/threats/{threat_id}")
         return res.json()
 
     @verify_credentials
     def list_detections(self):
         """List detections"""
-        res = self.get(
-            f"{self.account.hq}/detect/detections",
-            headers=self.account.headers,
-            params={},
-            timeout=10,
-        )
+        res = self.get(f"{self.account.hq}/detect/detections")
         return res.json()
 
     @verify_credentials
     def get_detection(self, detection_id):
         """Get properties of an existing detection"""
-        res = self.get(
-            f"{self.account.hq}/detect/detections/{detection_id}",
-            headers=self.account.headers,
-            timeout=10,
-        )
+        res = self.get(f"{self.account.hq}/detect/detections/{detection_id}")
         return res.json()
 
     @verify_credentials
     def list_threat_hunts(self, filters: dict = None):
         """List threat hunts"""
         res = self.get(
-            f"{self.account.hq}/detect/threat_hunts",
-            headers=self.account.headers,
-            params=filters if filters else {},
-            timeout=10,
+            f"{self.account.hq}/detect/threat_hunts", params=filters if filters else {}
         )
         threat_hunts = res.json()
         if self.account.resolve_enums:
@@ -181,11 +129,7 @@ class DetectController(HttpController):
     @verify_credentials
     def get_threat_hunt(self, threat_hunt_id):
         """Get properties of an existing threat hunt"""
-        res = self.get(
-            f"{self.account.hq}/detect/threat_hunts/{threat_hunt_id}",
-            headers=self.account.headers,
-            timeout=10,
-        )
+        res = self.get(f"{self.account.hq}/detect/threat_hunts/{threat_hunt_id}")
         threat_hunt = res.json()
         if self.account.resolve_enums:
             self.resolve_enums(threat_hunt, [(Control, "control")])
@@ -195,35 +139,27 @@ class DetectController(HttpController):
     def do_threat_hunt(self, threat_hunt_id):
         """Run a threat hunt"""
         res = self.post(
-            f"{self.account.hq}/detect/threat_hunts/{threat_hunt_id}",
-            headers=self.account.headers,
-            timeout=30,
+            f"{self.account.hq}/detect/threat_hunts/{threat_hunt_id}", timeout=30
         )
         return res.json()
 
     @verify_credentials
     def download(self, test_id, filename):
         """Clone a test file or attachment"""
-        res = self.get(
-            f"{self.account.hq}/detect/tests/{test_id}/{filename}",
-            headers=self.account.headers,
-            timeout=10,
-        )
+        res = self.get(f"{self.account.hq}/detect/tests/{test_id}/{filename}")
         return res.content
 
     @verify_credentials
     def schedule(self, items: list):
-        """Schedule tests and threats so endpoints will start running them
+        """
+        Schedule tests and threats so endpoints will start running them
 
         Example: items=[dict(run_code='DAILY', tags='grp-1,grp2', test_id='123-123-123'),
                         dict(run_code='DAILY', tags='grp-1', threat_id='abc-def-ghi')]
         """
-        res = self.post(
-            url=f"{self.account.hq}/detect/queue",
-            headers=self.account.headers,
-            json=dict(items=items),
-            timeout=10,
-        )
+        body = dict(items=items)
+
+        res = self.post(url=f"{self.account.hq}/detect/queue", json=body)
         schedule = res.json()
         if self.account.resolve_enums:
             self.resolve_enums(schedule, [(RunCode, "run_code")])
@@ -231,26 +167,21 @@ class DetectController(HttpController):
 
     @verify_credentials
     def unschedule(self, items: list):
-        """Unschedule tests and threats so endpoints will stop running them
+        """
+        Unschedule tests and threats so endpoints will stop running them
 
         Example: items=[dict(tags='grp-1,grp2', test_id='123-123-123'),
                         dict(tags='grp-1', threat_id='abc-def-ghi')]
         """
-        res = self.delete(
-            f"{self.account.hq}/detect/queue",
-            headers=self.account.headers,
-            json=dict(items=items),
-            timeout=10,
-        )
+        body = dict(items=items)
+
+        res = self.delete(f"{self.account.hq}/detect/queue", json=body)
         return res.json()
 
     @verify_credentials
     def accept_terms(self, name, version):
         """Accept terms and conditions"""
-        res = self.post(
-            f"{self.account.hq}/iam/terms",
-            headers=self.account.headers,
-            json=dict(name=name, version=version),
-            timeout=10,
-        )
+        body = dict(name=name, version=version)
+
+        res = self.post(f"{self.account.hq}/iam/terms", json=body)
         return res.json()

prelude_sdk/controllers/export_controller.py

@@ -17,15 +17,11 @@ class ExportController(HttpController):
         top: int = None,
     ):
         """Download partner data as a CSV"""
-        params = {
+        body = {
             "$filter": filter,
             "$orderby": orderby,
             "$top": top,
         }
-        res = self.post(
-            f"{self.account.hq}/export/scm/{export_type.name}",
-            params=params,
-            headers=self.account.headers,
-            timeout=10,
-        )
+
+        res = self.post(f"{self.account.hq}/export/scm/{export_type.name}", json=body)
         return res.json()

prelude_sdk/controllers/generate_controller.py

@@ -11,6 +11,7 @@ class GenerateController(HttpController):
     def upload_threat_intel(self, file: str):
         with open(file, "rb") as f:
             body = f.read()
+
         res = self.post(
             f"{self.account.hq}/generate/threat-intel",
             data=body,
@@ -21,19 +22,15 @@ class GenerateController(HttpController):
 
     @verify_credentials
     def get_threat_intel(self, job_id: str):
-        res = self.get(
-            f"{self.account.hq}/generate/threat-intel/{job_id}",
-            headers=self.account.headers,
-            timeout=10,
-        )
+        res = self.get(f"{self.account.hq}/generate/threat-intel/{job_id}")
         return res.json()
 
     @verify_credentials
     def generate_from_partner_advisory(self, partner: Control, advisory_id: str):
         params = dict(advisory_id=advisory_id)
+
         res = self.post(
             f"{self.account.hq}/generate/partner-advisories/{partner.name}",
-            headers=self.account.headers,
             json=params,
             timeout=30,
         )

prelude_sdk/controllers/http_controller.py

@@ -42,42 +42,58 @@ class HttpController(object):
                 elif isinstance(v, dict) or isinstance(v, list):
                     self._resolve_enum(v, enum_class, key)
 
-    def get(self, url, retry=True, **kwargs):
-        res = self._session.get(url, **kwargs)
+    def get(self, url, retry=True, timeout=10, headers=None, **kwargs):
+        res = self._session.get(
+            url, timeout=timeout, headers=headers or self.account.headers, **kwargs
+        )
         if res.status_code == 200:
             return res
         if res.status_code == 401 and retry and self.account.token_location:
             self.account.refresh_tokens()
             self.account.update_auth_header()
-            return self.get(url, retry=False, **kwargs)
+            return self.get(
+                url, retry=False, timeout=timeout, headers=headers, **kwargs
+            )
         raise Exception(res.text)
 
-    def post(self, url, retry=True, **kwargs):
-        res = self._session.post(url, **kwargs)
+    def post(self, url, retry=True, timeout=10, headers=None, **kwargs):
+        res = self._session.post(
+            url, timeout=timeout, headers=headers or self.account.headers, **kwargs
+        )
         if res.status_code == 200:
             return res
         if res.status_code == 401 and retry and self.account.token_location:
             self.account.refresh_tokens()
             self.account.update_auth_header()
-            return self.post(url, retry=False, **kwargs)
+            return self.post(
+                url, retry=False, timeout=timeout, headers=headers, **kwargs
+            )
         raise Exception(res.text)
 
-    def delete(self, url, retry=True, **kwargs):
-        res = self._session.delete(url, **kwargs)
+    def delete(self, url, retry=True, timeout=10, headers=None, **kwargs):
+        res = self._session.delete(
+            url, timeout=timeout, headers=headers or self.account.headers, **kwargs
+        )
         if res.status_code == 200:
             return res
         if res.status_code == 401 and retry and self.account.token_location:
             self.account.refresh_tokens()
             self.account.update_auth_header()
-            return self.delete(url, retry=False, **kwargs)
+            return self.delete(
+                url, retry=False, timeout=timeout, headers=headers, **kwargs
+            )
         raise Exception(res.text)
 
-    def put(self, url, retry=True, **kwargs):
-        res = self._session.put(url, **kwargs)
+    def put(self, url, retry=True, timeout=10, headers=None, **kwargs):
+        res = self._session.put(
+            url, timeout=timeout, headers=headers or self.account.headers, **kwargs
+        )
         if res.status_code == 200:
             return res
         if res.status_code == 401 and retry and self.account.token_location:
             self.account.refresh_tokens()
             self.account.update_auth_header()
-            return self.put(url, retry=False, **kwargs)
+            return self.put(
+                url, retry=False, timeout=timeout, headers=headers, **kwargs
+            )
         raise Exception(res.text)