Updates for Lacework v4.20 (#24)

* fix: removed `policy_ui` parameters for compatibility
* fix: cleanup for HTTP 204 response on search endpoints
* refactor: readability improvements for tests

Author: alannix-lw

laceworksdk/api/alert_channels.py

@@ -140,7 +140,10 @@ def search(self,
 
         response = self._session.post(api_uri, data=query_data, org=org)
 
-        return response.json()
+        if response.status_code == 204:
+            return {"data": []}
+        else:
+            return response.json()
 
     def update(self,
                guid,

laceworksdk/api/alert_rules.py

@@ -129,7 +129,10 @@ def search(self,
 
         response = self._session.post(api_uri, data=query_data, org=org)
 
-        return response.json()
+        if response.status_code == 204:
+            return {"data": []}
+        else:
+            return response.json()
 
     def update(self,
                guid,

laceworksdk/api/audit_logs.py

@@ -72,4 +72,7 @@ def search(self,
 
         response = self._session.post(api_uri, data=query_data, org=org)
 
-        return response.json()
+        if response.status_code == 204:
+            return {"data": []}
+        else:
+            return response.json()

laceworksdk/api/policies.py

@@ -36,7 +36,6 @@ def create(self,
                evaluator_id,
                limit=None,
                eval_frequency=None,
-               policy_ui=None,
                org=False):
         """
         A method to create a new Lacework Query Language (LQL) policy.
@@ -55,7 +54,6 @@ def create(self,
         :param limit: An integer representing the number of results to return.
         :param eval_frequency: A string representing the frequency in which to evaluate the policy.
             ("Hourly", "Daily")
-        :param policy_ui: A dictionary which specifies the policyUi domain and subdomain.
         :param org: A boolean representing whether the request should be performed
             at the Organization level
 
@@ -67,9 +65,6 @@ def create(self,
         # Build the Policies request URI
         api_uri = "/api/v2/Policies"
 
-        if policy_ui is None:
-            policy_ui = {}
-
         data = {
             "policyType": policy_type,
             "queryId": query_id,
@@ -80,7 +75,6 @@ def create(self,
             "severity": severity,
             "alertEnabled": int(bool(alert_enabled)),
             "alertProfile": alert_profile,
-            "policyUi": policy_ui,
             "evaluatorId": evaluator_id
         }
 
@@ -146,7 +140,6 @@ def update(self,
                alert_profile=None,
                limit=None,
                eval_frequency=None,
-               policy_ui=None,
                org=False):
         """
         A method to update a Lacework Query Language (LQL) policy.
@@ -165,7 +158,6 @@ def update(self,
         :param limit: An integer representing the number of results to return.
         :param eval_frequency: A string representing the frequency in which to evaluate the policy.
             ("Hourly", "Daily")
-        :param policy_ui: A dictionary which specifies the policyUi domain and subdomain.
         :param org: A boolean representing whether the request should be performed
             at the Organization level
 
@@ -201,8 +193,6 @@ def update(self,
             data["limit"] = limit
         if eval_frequency:
             data["evalFrequency"] = eval_frequency
-        if policy_ui:
-            data["policyUi"] = policy_ui
 
         response = self._session.patch(api_uri, org=org, data=data)
 

laceworksdk/api/report_rules.py

@@ -129,7 +129,10 @@ def search(self,
 
         response = self._session.post(api_uri, data=query_data, org=org)
 
-        return response.json()
+        if response.status_code == 204:
+            return {"data": []}
+        else:
+            return response.json()
 
     def update(self,
                guid,

tests/api/v2/test_policies.py

@@ -17,7 +17,7 @@
 end_time = current_time.strftime("%Y-%m-%dT%H:%M:%S.000Z")
 
 POLICY_ID = None
-RANDOM_TEXT = "".join(random.choice(string.ascii_uppercase) for _ in range(8))
+RANDOM_TEXT = "".join(random.choices(string.ascii_uppercase, k=8))
 
 # Tests
 

tests/api/v2/test_queries.py

@@ -16,7 +16,7 @@
 start_time = start_time.strftime("%Y-%m-%dT%H:%M:%S.000Z")
 end_time = current_time.strftime("%Y-%m-%dT%H:%M:%S.000Z")
 
-RANDOM_TEXT = "".join(random.choice(string.ascii_uppercase) for _ in range(8))
+RANDOM_TEXT = "".join(random.choices(string.ascii_uppercase, k=8))
 
 # Tests