Refactor/173 use proper retry

prismacloud/api/cwpp/cwpp.py

@@ -4,6 +4,8 @@
 import time
 
 import requests
+from requests.adapters import HTTPAdapter, Retry
+
 
 class PrismaCloudAPICWPPMixin():
     """ Requests and Output """
@@ -16,7 +18,8 @@ def login_compute(self):
             # Login via CWP.
             self.login('https://%s/api/v1/authenticate' % self.api_compute)
         else:
-            self.error_and_exit(418, "Specify a Prisma Cloud URL or Prisma Cloud Compute URL")
+            self.error_and_exit(
+                418, "Specify a Prisma Cloud URL or Prisma Cloud Compute URL")
         self.debug_print('New API Token: %s' % self.token)
 
     def extend_login_compute(self):
@@ -43,78 +46,91 @@ def execute_compute(self, action, endpoint, query_params=None, body_params=None,
         limit = 50
         more = False
         results = []
-        while offset == 0 or more is True:
-            if int(time.time() - self.token_timer) > self.token_limit:
-                self.extend_login_compute()
-            if paginated:
-                url = 'https://%s/%s?limit=%s&offset=%s' % (self.api_compute, endpoint, limit, offset)
-            else:
-                url = 'https://%s/%s' % (self.api_compute, endpoint)
-            if self.token:
-                if self.api:
-                    # Authenticate via CSPM
-                    request_headers['x-redlock-auth'] = self.token
+
+        retries = Retry(total=self.retry_number,
+                        status_forcelist=self.retry_status_codes, raise_on_status=False)
+
+        with requests.Session() as session:
+            session.mount('https://%s/%s' % (self.api_compute,
+                          endpoint), adapter=HTTPAdapter(max_retries=retries))
+
+            while offset == 0 or more is True:
+                if int(time.time() - self.token_timer) > self.token_limit:
+                    self.extend_login_compute()
+                if paginated:
+                    url = 'https://%s/%s?limit=%s&offset=%s' % (
+                        self.api_compute, endpoint, limit, offset)
                 else:
-                    # Authenticate via CWP
-                    request_headers['Authorization'] = "Bearer %s" % self.token
-            self.debug_print('API URL: %s' % url)
-            self.debug_print('API Request Headers: (%s)' % request_headers)
-            self.debug_print('API Query Params: %s' % query_params)
-            self.debug_print('API Body Params: %s' % body_params_json)
-            # Add User-Agent to the headers
-            request_headers['User-Agent'] = self.user_agent
-            api_response = requests.request(action, url, headers=request_headers, params=query_params, data=body_params_json, verify=self.verify, timeout=self.timeout)
-            self.debug_print('API Response Status Code: (%s)' % api_response.status_code)
-            self.debug_print('API Response Headers: (%s)' % api_response.headers)
-            if api_response.status_code in self.retry_status_codes:
-                for exponential_wait in self.retry_waits:
-                    time.sleep(exponential_wait)
-                    api_response = requests.request(action, url, headers=request_headers, params=query_params, data=body_params_json, verify=self.verify, timeout=self.timeout)
-                    if api_response.ok:
-                        break # retry loop
-            if api_response.ok:
-                if not api_response.content:
-                    return None
-                if api_response.headers.get('Content-Type') == 'application/x-gzip':
-                    return api_response.content
-                if api_response.headers.get('Content-Type') == 'text/csv':
-                    return api_response.content.decode('utf-8')
-                try:
-                    result = json.loads(api_response.content)
-                    #if result is None:
-                    #    self.logger.error('JSON returned None, API: (%s) with query params: (%s) and body params: (%s) parsing response: (%s)' % (url, query_params, body_params, api_response.content))
-                    #    if force:
-                    #        return results # or continue
-                    #    self.error_and_exit(api_response.status_code, 'JSON returned None, API: (%s) with query params: (%s) and body params: (%s) parsing response: (%s)' % (url, query_params, body_params, api_response.content))
-                except ValueError:
-                    self.logger.error('JSON raised ValueError, API: (%s) with query params: (%s) and body params: (%s) parsing response: (%s)' % (url, query_params, body_params, api_response.content))
-                    if force:
-                        return results # or continue
-                    self.error_and_exit(api_response.status_code, 'JSON raised ValueError, API: (%s) with query params: (%s) and body params: (%s) parsing response: (%s)' % (url, query_params, body_params, api_response.content))
-                if 'Total-Count' in api_response.headers:
-                    self.debug_print('Retrieving Next Page of Results: Offset/Total Count: %s/%s' % (offset, api_response.headers['Total-Count']))
-                    total_count = int(api_response.headers['Total-Count'])
-                    if total_count > 0:
-                        results.extend(result)
-                    offset += limit
-                    more = bool(offset < total_count)
+                    url = 'https://%s/%s' % (self.api_compute, endpoint)
+                if self.token:
+                    if self.api:
+                        # Authenticate via CSPM
+                        request_headers['x-redlock-auth'] = self.token
+                    else:
+                        # Authenticate via CWP
+                        request_headers['Authorization'] = "Bearer %s" % self.token
+                self.debug_print('API URL: %s' % url)
+                self.debug_print('API Request Headers: (%s)' % request_headers)
+                self.debug_print('API Query Params: %s' % query_params)
+                self.debug_print('API Body Params: %s' % body_params_json)
+                # Add User-Agent to the headers
+                request_headers['User-Agent'] = self.user_agent
+                api_response = session.request(action, url, headers=request_headers, params=query_params,
+                                               data=body_params_json, verify=self.verify, timeout=self.timeout)
+                self.debug_print('API Response Status Code: (%s)' %
+                                 api_response.status_code)
+                self.debug_print('API Response Headers: (%s)' %
+                                 api_response.headers)
+                if api_response.ok:
+                    if not api_response.content:
+                        return None
+                    if api_response.headers.get('Content-Type') == 'application/x-gzip':
+                        return api_response.content
+                    if api_response.headers.get('Content-Type') == 'text/csv':
+                        return api_response.content.decode('utf-8')
+                    try:
+                        result = json.loads(api_response.content)
+                        # if result is None:
+                        #    self.logger.error('JSON returned None, API: (%s) with query params: (%s) and body params: (%s) parsing response: (%s)' % (url, query_params, body_params, api_response.content))
+                        #    if force:
+                        #        return results # or continue
+                        #    self.error_and_exit(api_response.status_code, 'JSON returned None, API: (%s) with query params: (%s) and body params: (%s) parsing response: (%s)' % (url, query_params, body_params, api_response.content))
+                    except ValueError:
+                        self.logger.error('JSON raised ValueError, API: (%s) with query params: (%s) and body params: (%s) parsing response: (%s)' % (
+                            url, query_params, body_params, api_response.content))
+                        if force:
+                            return results  # or continue
+                        self.error_and_exit(api_response.status_code, 'JSON raised ValueError, API: (%s) with query params: (%s) and body params: (%s) parsing response: (%s)' % (
+                            url, query_params, body_params, api_response.content))
+                    if 'Total-Count' in api_response.headers:
+                        self.debug_print('Retrieving Next Page of Results: Offset/Total Count: %s/%s' % (
+                            offset, api_response.headers['Total-Count']))
+                        total_count = int(api_response.headers['Total-Count'])
+                        if total_count > 0:
+                            results.extend(result)
+                        offset += limit
+                        more = bool(offset < total_count)
+                    else:
+                        return result
                 else:
-                    return result
-            else:
-                self.logger.error('API: (%s) responded with a status of: (%s), with query: (%s) and body params: (%s)' % (url, api_response.status_code, query_params, body_params))
-                if force:
-                    return results
-                self.error_and_exit(api_response.status_code, 'API: (%s) with query params: (%s) and body params: (%s) responded with an error and this response:\n%s' % (url, query_params, body_params, api_response.text))
-        return results
+                    self.logger.error('API: (%s) responded with a status of: (%s), with query: (%s) and body params: (%s)' % (
+                        url, api_response.status_code, query_params, body_params))
+                    if force:
+                        return results
+                    self.error_and_exit(api_response.status_code, 'API: (%s) with query params: (%s) and body params: (%s) responded with an error and this response:\n%s' % (
+                        url, query_params, body_params, api_response.text))
+            return results
 
     # The Compute API setting is optional.
 
     def validate_api_compute(self):
         if not self.api_compute:
-            self.error_and_exit(500, 'Please specify a Prisma Cloud Compute API URL.')
+            self.error_and_exit(
+                500, 'Please specify a Prisma Cloud Compute API URL.')
 
     # Exit handler (Error).
 
     @classmethod
     def error_and_exit(cls, error_code, error_message='', system_message=''):
-        raise SystemExit('\n\nStatus Code: %s\n%s\n%s\n' % (error_code, error_message, system_message))
+        raise SystemExit('\n\nStatus Code: %s\n%s\n%s\n' %
+                         (error_code, error_message, system_message))

prismacloud/api/pc_lib_api.py

@@ -43,6 +43,7 @@ def __init__(self):
         self.token_limit        = 590 # aka 9 minutes
         self.retry_status_codes = [425, 429, 500, 502, 503, 504]
         self.retry_waits        = [1, 2, 4, 8, 16, 32]
+        self.retry_number       = 6
         self.max_workers        = 8
         #
         self.error_log          = 'error.log'

pyproject.toml

@@ -1,3 +1,6 @@
 [build-system]
 requires = ["setuptools"]
-build-backend = "setuptools.build_meta"
+build-backend = "setuptools.build_meta"
+
+[project.optional-dependencies]
+test = ["coverage==7.6.10", "responses==0.25.3"]

requirements_test.txt

@@ -0,0 +1,2 @@
+coverage==7.6.10
+responses==0.25.3

setup.py

@@ -34,5 +34,8 @@
         'requests',
         'update_checker'
     ],
+    extras_require={
+        'test': ['coverage==7.6.10', 'responses==0.25.3']
+    },
     python_requires='>=3.6'
 )