Merge pull request #10 from CybercentreCanada/hotfix/retries

Hotfix/retries
CybercentreCanada · May 5, 2021 · 315bb23 · 315bb23
2 parents 1c92625 + ee991cc
commit 315bb23
Showing 1 changed file with 36 additions and 28 deletions.
diff --git a/assemblyline_client/__init__.py b/assemblyline_client/__init__.py
@@ -172,36 +172,44 @@ def request(self, func, path, process, **kw):
             while self.max_retries < 1 or retries <= self.max_retries:
                 if retries:
                     time.sleep(min(2, 2 ** (retries - 7)))
-                response = func('/'.join((self.server, path)), **kw)
-                if 'XSRF-TOKEN' in response.cookies:
-                    self.session.headers.update({'X-XSRF-TOKEN': response.cookies['XSRF-TOKEN']})
-                if response.ok:
-                    return process(response)
-                elif response.status_code == 401:
-                    try:
-                        resp_data = response.json()
-                        if resp_data["api_error_message"] == "Authentication required":
-                            self._authenticate()
-                        else:
+                    stream = kw.get('files', {}).get('bin', None)
+                    if stream and 'seek' in dir(stream):
+                        stream.seek(0)
+
+                try:
+                    response = func('/'.join((self.server, path)), **kw)
+                    if 'XSRF-TOKEN' in response.cookies:
+                        self.session.headers.update({'X-XSRF-TOKEN': response.cookies['XSRF-TOKEN']})
+                    if response.ok:
+                        return process(response)
+                    elif response.status_code == 401:
+                        try:
+                            resp_data = response.json()
+                            if resp_data["api_error_message"] == "Authentication required":
+                                self._authenticate()
+                            else:
+                                raise ClientError(resp_data["api_error_message"], response.status_code,
+                                                  api_version=resp_data["api_server_version"],
+                                                  api_response=resp_data["api_response"])
+                        except Exception as e:
+                            if isinstance(e, ClientError):
+                                raise
+
+                            raise ClientError(response.content, response.status_code)
+
+                    elif response.status_code not in (502, 503, 504):
+                        try:
+                            resp_data = response.json()
                             raise ClientError(resp_data["api_error_message"], response.status_code,
                                               api_version=resp_data["api_server_version"],
                                               api_response=resp_data["api_response"])
-                    except Exception as e:
-                        if isinstance(e, ClientError):
-                            raise
-
-                        raise ClientError(response.content, response.status_code)
-
-                elif response.status_code not in (502, 503, 504):
-                    try:
-                        resp_data = response.json()
-                        raise ClientError(resp_data["api_error_message"], response.status_code,
-                                          api_version=resp_data["api_server_version"],
-                                          api_response=resp_data["api_response"])
-                    except Exception as e:
-                        if isinstance(e, ClientError):
-                            raise
-
-                        raise ClientError(response.content, response.status_code)
+                        except Exception as e:
+                            if isinstance(e, ClientError):
+                                raise
+
+                            raise ClientError(response.content, response.status_code)
+                except OSError as e:
+                    if "Connection aborted" not in str(e):
+                        raise
 
                 retries += 1