From 56d266d4e6d36135129e72a4fbd82324ce83cb3b Mon Sep 17 00:00:00 2001 From: Burr Webb Date: Tue, 2 Jul 2019 15:41:06 -0600 Subject: [PATCH 1/5] adding API request handler. will account for API rate limiting when implemented. --- api_request_handler/__init__.py | 1 + api_request_handler/api_request_handler.py | 110 +++++++++++++++++++++ upload_to_platform.py | 80 ++++++--------- 3 files changed, 140 insertions(+), 51 deletions(-) create mode 100644 api_request_handler/__init__.py create mode 100644 api_request_handler/api_request_handler.py diff --git a/api_request_handler/__init__.py b/api_request_handler/__init__.py new file mode 100644 index 0000000..ae843b3 --- /dev/null +++ b/api_request_handler/__init__.py @@ -0,0 +1 @@ +from .api_request_handler import ApiRequestHandler diff --git a/api_request_handler/api_request_handler.py b/api_request_handler/api_request_handler.py new file mode 100644 index 0000000..7b263ad --- /dev/null +++ b/api_request_handler/api_request_handler.py @@ -0,0 +1,110 @@ +""" +****************************************************************** + +Name : api_request_handler.py +Module : api_request_handler +Description : RiskSense API Request Handler +Copyright : (c) RiskSense, Inc. +License : ???? + +****************************************************************** +""" + +import json +import time +import requests + + +class ApiRequestHandler: + + def __init__(self, api_key, user_agent=None, max_retries=5): + + """ + Initialize ApiRequestHandler + + :param api_key: RiskSense Platform API key + :type api_key: str + + :param user_agent: User-Agent + :type user_agent: str + + :param max_retries maximum number of retries for a request + :type max_retries int + """ + + self.api_key = api_key + self.retry_counter = 0 + + self.user_agent = user_agent + + self.max_retries = max_retries + + def make_api_request(self, method, url, body=None, files=None, retry=False): + + """ + + :param method: HTTP method to use for request (GET or POST) + :type method: str + + :param url: URL for API endpoint + :type url: str + + :param body: Body to be used in API request (if required) + :type body: dict + + :param files: Files to pass to API + :type files: dict + + :param retry: Indicate whether this is a retry attempt + :type retry: bool + + + :return: The requests module API response object if successfully sent. + A None is returned if max retries reached or invalid HTTP method provided. + """ + + if retry is True: + self.increment_retry_counter() + + if self.retry_counter == self.max_retries: + print(f"Max number of retries ({self.max_retries}) reached...") + return None + + header = { + "User-Agent": "upload_to_platform", + "x-api-key": self.api_key, + "content-type": "application/json" + } + + if method.lower() == "get": + + response = requests.get(url, headers=header) + + elif method.lower() == "post": + + response = requests.post(url, headers=header, data=json.dumps(body), files=files) + + else: + print(f"Unsupported method provided: {method}") + print("Please provide a supported HTTP method (GET or POST)") + return None + + if response and response.status_code == 503: + + time.sleep(1) + print(f"503 error returned, retrying (this was attempt number {self.retry_counter + 1})...") + new_response = self.make_api_request(method, url, body, files, retry=True) + + return new_response + + self.reset_retry_counter() + + return response + + def reset_retry_counter(self): + + self.retry_counter = 0 + + def increment_retry_counter(self): + + self.retry_counter += 1 diff --git a/upload_to_platform.py b/upload_to_platform.py index 7bd51de..c7230a7 100644 --- a/upload_to_platform.py +++ b/upload_to_platform.py @@ -17,8 +17,9 @@ import os import logging +from api_request_handler import ApiRequestHandler + import toml -import requests import progressbar @@ -38,14 +39,12 @@ def get_client_id(platform, key): :rtype: int """ + request_handler = ApiRequestHandler(key) + url = platform + "/api/v1/client?size=150" - header = { - 'x-api-key': key, - 'content-type': 'application/json' - } + raw_client_id_response = request_handler.make_api_request("GET", url) - raw_client_id_response = requests.get(url, headers=header) json_client_id_response = json.loads(raw_client_id_response.text) if raw_client_id_response.status_code == 200: @@ -107,16 +106,14 @@ def validate_client_id(client, platform, key): :rtype: bool """ + request_handler = ApiRequestHandler(key) + validity = False url = platform + "/api/v1/client/" + str(client) - header = { - 'x-api-key': key, - 'content-type': 'application/json' - } + raw_client_id_response = request_handler.make_api_request("GET", url) - raw_client_id_response = requests.get(url, headers=header) json_client_id_response = json.loads(raw_client_id_response.text) if raw_client_id_response.status_code == 200 and json_client_id_response['id'] == client: @@ -143,6 +140,8 @@ def find_network_id(platform, key, client): :return: The selected Network ID """ + request_handler = ApiRequestHandler(key) + network = 0 logging.info("Getting Network ID") @@ -162,12 +161,6 @@ def find_network_id(platform, key, client): url = platform + "/api/v1/client/" + str(client) + "/network/search" - header = { - 'x-api-key': key, - 'Content-Type': "application/json", - 'Cache-Control': "no-cache" - } - body = { "filters": [ { @@ -188,7 +181,8 @@ def find_network_id(platform, key, client): "size": 20 } - raw_network_search_response = requests.post(url, headers=header, data=json.dumps(body)) + raw_network_search_response = request_handler.make_api_request("POST", url, body=body) + json_network_search_response = json.loads(raw_network_search_response.text) if raw_network_search_response.status_code == 200 and \ @@ -254,6 +248,8 @@ def create_new_assessment(platform, key, client, name, start_date, notes): :rtype: int """ + request_handler = ApiRequestHandler(key) + created_id = 0 logging.info("Creating new assessment.") @@ -263,19 +259,14 @@ def create_new_assessment(platform, key, client, name, start_date, notes): url = platform + "/api/v1/client/" + str(client) + "/assessment" - header = { - 'x-api-key': key, - 'Content-Type': "application/json", - 'Cache-Control': "no-cache" - } - body = { "name": name, "startDate": start_date, "notes": notes } - raw_assessment_response = requests.post(url, headers=header, data=json.dumps(body)) + raw_assessment_response = request_handler.make_api_request("POST", url, body=body) + json_assessment_response = json.loads(raw_assessment_response.text) if raw_assessment_response.status_code == 201: @@ -310,6 +301,8 @@ def create_upload(platform, key, assessment, network, client): :rtype: int """ + request_handler = ApiRequestHandler(key) + today = datetime.date.today() current_time = time.time() @@ -319,19 +312,14 @@ def create_upload(platform, key, assessment, network, client): url = platform + "/api/v1/client/" + str(client) + "/upload" - header = { - 'x-api-key': key, - 'Content-Type': "application/json", - 'Cache-Control': "no-cache" - } - body = { 'assessmentId': assessment, 'networkId': network, 'name': upload_name } - raw_upload_response = requests.post(url, headers=header, data=json.dumps(body)) + raw_upload_response = request_handler.make_api_request("POST", url, body=body) + json_upload_json_response = json.loads(raw_upload_response.text) if raw_upload_response.status_code == 201: @@ -367,19 +355,16 @@ def add_file_to_upload(platform, key, client, upload, file_name, file_path): :type file_path: str """ + request_handler = ApiRequestHandler(key) + logging.info("Adding file to upload: %s", file_name) logging.debug("File Path: %s", file_path) url = platform + "/api/v1/client/" + str(client) + "/upload/" + str(upload) + "/file" - header = { - 'x-api-key': key, - 'Cache-Control': "no-cache" - } - upload_file = {'scanFile': (file_name, open(file_path + "/" + file_name, 'rb'))} - raw_add_file_response = requests.post(url, headers=header, files=upload_file) + raw_add_file_response = request_handler.make_api_request("POST", url, files=upload_file) if raw_add_file_response.status_code != 201: print(f"Error uploading file {file_name}. Status Code returned was {raw_add_file_response.status_code}") @@ -410,21 +395,17 @@ def begin_processing(platform, key, client, upload, run_urba): :type run_urba: bool """ + request_handler = ApiRequestHandler(key) + logging.info("Starting platform processing") url = platform + "/api/v1/client/" + str(client) + "/upload/" + str(upload) + "/start" - header = { - 'x-api-key': key, - 'Content-Type': "application/json", - 'Cache-Control': "no-cache" - } - body = { "autoUrba": run_urba } - raw_begin_processing_response = requests.post(url, headers=header, data=json.dumps(body)) + raw_begin_processing_response = request_handler.make_api_request("POST", url, body=body) if raw_begin_processing_response.status_code == 200: print("Uploaded file(s) now processing. This may take a while. Please wait...") @@ -456,17 +437,14 @@ def check_upload_state(platform, key, client, upload): :rtype: str """ + request_handler = ApiRequestHandler(key) + logging.info("Checking status of the upload processing") url = platform + "/api/v1/client/" + str(client) + "/upload/" + str(upload) - header = { - 'x-api-key': key, - 'Content-Type': "application/json", - 'Cache-Control': "no-cache" - } + raw_check_upload_state_response = request_handler.make_api_request("GET", url) - raw_check_upload_state_response = requests.get(url, headers=header) json_check_upload_state_response = json.loads(raw_check_upload_state_response.text) if raw_check_upload_state_response.status_code == 200: From 0e7144140c2e303d1964788137025aa0bd7ddd7a Mon Sep 17 00:00:00 2001 From: Burr Webb Date: Tue, 2 Jul 2019 16:21:54 -0600 Subject: [PATCH 2/5] updating with license info, user agent, versioning --- api_request_handler/api_request_handler.py | 26 +-- license.txt | 202 +++++++++++++++++++++ upload_to_platform.py | 30 +-- 3 files changed, 235 insertions(+), 23 deletions(-) create mode 100644 license.txt diff --git a/api_request_handler/api_request_handler.py b/api_request_handler/api_request_handler.py index 7b263ad..09bf250 100644 --- a/api_request_handler/api_request_handler.py +++ b/api_request_handler/api_request_handler.py @@ -5,7 +5,7 @@ Module : api_request_handler Description : RiskSense API Request Handler Copyright : (c) RiskSense, Inc. -License : ???? +License : Apache-2.0 ****************************************************************** """ @@ -17,27 +17,29 @@ class ApiRequestHandler: - def __init__(self, api_key, user_agent=None, max_retries=5): + def __init__(self, api_key, user_agent=None, max_retries=5, retry_wait_time=1): """ Initialize ApiRequestHandler - :param api_key: RiskSense Platform API key - :type api_key: str + :param api_key: RiskSense Platform API key + :type api_key: str - :param user_agent: User-Agent - :type user_agent: str + :param user_agent: User-Agent + :type user_agent: str - :param max_retries maximum number of retries for a request - :type max_retries int + :param max_retries: maximum number of retries for a request + :type max_retries: int + + :param retry_wait_time: Time to wait (in seconds) when 503 error encountered. + :type retry_wait_time: int """ self.api_key = api_key self.retry_counter = 0 - self.user_agent = user_agent - self.max_retries = max_retries + self.retry_wait_time = retry_wait_time def make_api_request(self, method, url, body=None, files=None, retry=False): @@ -71,7 +73,7 @@ def make_api_request(self, method, url, body=None, files=None, retry=False): return None header = { - "User-Agent": "upload_to_platform", + "User-Agent": self.user_agent, "x-api-key": self.api_key, "content-type": "application/json" } @@ -91,7 +93,7 @@ def make_api_request(self, method, url, body=None, files=None, retry=False): if response and response.status_code == 503: - time.sleep(1) + time.sleep(self.retry_wait_time) print(f"503 error returned, retrying (this was attempt number {self.retry_counter + 1})...") new_response = self.make_api_request(method, url, body, files, retry=True) diff --git a/license.txt b/license.txt new file mode 100644 index 0000000..7a4a3ea --- /dev/null +++ b/license.txt @@ -0,0 +1,202 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. \ No newline at end of file diff --git a/upload_to_platform.py b/upload_to_platform.py index c7230a7..d957eaf 100644 --- a/upload_to_platform.py +++ b/upload_to_platform.py @@ -5,7 +5,7 @@ Description : Uploads files to the RiskSense platform, and kicks off the processing of those files. Copyright : (c) RiskSense, Inc. -License : Proprietary +License : Apache-2.0 ****************************************************************** """ @@ -22,6 +22,10 @@ import toml import progressbar +__version__ = "0.5" + +USER_AGENT_STRING = "upload_to_platform_v" + __version__ + def get_client_id(platform, key): @@ -39,7 +43,7 @@ def get_client_id(platform, key): :rtype: int """ - request_handler = ApiRequestHandler(key) + request_handler = ApiRequestHandler(key, user_agent=USER_AGENT_STRING) url = platform + "/api/v1/client?size=150" @@ -106,7 +110,7 @@ def validate_client_id(client, platform, key): :rtype: bool """ - request_handler = ApiRequestHandler(key) + request_handler = ApiRequestHandler(key, user_agent=USER_AGENT_STRING) validity = False @@ -140,7 +144,7 @@ def find_network_id(platform, key, client): :return: The selected Network ID """ - request_handler = ApiRequestHandler(key) + request_handler = ApiRequestHandler(key, user_agent=USER_AGENT_STRING) network = 0 @@ -248,7 +252,7 @@ def create_new_assessment(platform, key, client, name, start_date, notes): :rtype: int """ - request_handler = ApiRequestHandler(key) + request_handler = ApiRequestHandler(key, user_agent=USER_AGENT_STRING) created_id = 0 @@ -301,7 +305,7 @@ def create_upload(platform, key, assessment, network, client): :rtype: int """ - request_handler = ApiRequestHandler(key) + request_handler = ApiRequestHandler(key, user_agent=USER_AGENT_STRING) today = datetime.date.today() current_time = time.time() @@ -355,7 +359,7 @@ def add_file_to_upload(platform, key, client, upload, file_name, file_path): :type file_path: str """ - request_handler = ApiRequestHandler(key) + request_handler = ApiRequestHandler(key, user_agent=USER_AGENT_STRING) logging.info("Adding file to upload: %s", file_name) logging.debug("File Path: %s", file_path) @@ -395,7 +399,7 @@ def begin_processing(platform, key, client, upload, run_urba): :type run_urba: bool """ - request_handler = ApiRequestHandler(key) + request_handler = ApiRequestHandler(key, user_agent=USER_AGENT_STRING) logging.info("Starting platform processing") @@ -437,7 +441,7 @@ def check_upload_state(platform, key, client, upload): :rtype: str """ - request_handler = ApiRequestHandler(key) + request_handler = ApiRequestHandler(key, user_agent=USER_AGENT_STRING) logging.info("Checking status of the upload processing") @@ -500,6 +504,10 @@ def main(): """ Main body of script """ + print() + print(f"RiskSense - Upload to Platform v{__version__}") + print() + conf_file = os.path.join(os.path.abspath(os.path.dirname(__file__)), 'conf', 'config.toml') config = read_config_file(conf_file) @@ -514,8 +522,8 @@ def main(): auto_urba = config["auto_urba"] if api_key == "": - print("No API Key configured. Please add your API Key to the configuration file.") - logging.info("No API Key configured. Please add your API Key to the configuration file.") + print("No API Key configured. Please add your API Key to the configuration file (conf/config.toml).") + logging.info("No API Key configured. Please add your API Key to the configuration file (conf/config.toml).") input("Please press ENTER to close.") exit(1) From 49c99371e7289635f05bea3ad9c9a86bcc12f277 Mon Sep 17 00:00:00 2001 From: Burr Webb Date: Fri, 19 Jul 2019 14:24:14 -0600 Subject: [PATCH 3/5] adding request handler with proper support for retries --- api_request_handler/api_request_handler.py | 138 +++++++++------ requirements.txt | 1 + upload_to_platform.py | 190 ++++++++++++++++----- 3 files changed, 229 insertions(+), 100 deletions(-) diff --git a/api_request_handler/api_request_handler.py b/api_request_handler/api_request_handler.py index 09bf250..3254609 100644 --- a/api_request_handler/api_request_handler.py +++ b/api_request_handler/api_request_handler.py @@ -1,5 +1,5 @@ """ -****************************************************************** +********************************************************************* Name : api_request_handler.py Module : api_request_handler @@ -7,20 +7,21 @@ Copyright : (c) RiskSense, Inc. License : Apache-2.0 -****************************************************************** +********************************************************************* """ import json -import time import requests +from requests.adapters import HTTPAdapter +from urllib3 import Retry class ApiRequestHandler: - def __init__(self, api_key, user_agent=None, max_retries=5, retry_wait_time=1): + def __init__(self, api_key, user_agent=None, max_retries=5): """ - Initialize ApiRequestHandler + Initialize ApiRequestHandler class. :param api_key: RiskSense Platform API key :type api_key: str @@ -30,83 +31,112 @@ def __init__(self, api_key, user_agent=None, max_retries=5, retry_wait_time=1): :param max_retries: maximum number of retries for a request :type max_retries: int - - :param retry_wait_time: Time to wait (in seconds) when 503 error encountered. - :type retry_wait_time: int """ self.api_key = api_key - self.retry_counter = 0 self.user_agent = user_agent self.max_retries = max_retries - self.retry_wait_time = retry_wait_time - - def make_api_request(self, method, url, body=None, files=None, retry=False): - """ + self.__retry_counter = 0 - :param method: HTTP method to use for request (GET or POST) - :type method: str + def make_request(self, http_method, url, body=None, files=None): - :param url: URL for API endpoint - :type url: str + """ - :param body: Body to be used in API request (if required) - :type body: dict + :param http_method: HTTP method to use for request (GET or POST) + :type http_method: str - :param files: Files to pass to API - :type files: dict + :param url: URL for API endpoint + :type url: str - :param retry: Indicate whether this is a retry attempt - :type retry: bool + :param body: Body to be used in API request (if required) + :type body: dict + :param files: Files to pass to API + :type files: dict - :return: The requests module API response object if successfully sent. - A None is returned if max retries reached or invalid HTTP method provided. + :return: The requests module API response object is returned if request is successfully sent. + A None is returned invalid HTTP method provided. + If max retries are reached, an exception is raised. """ - if retry is True: - self.increment_retry_counter() - - if self.retry_counter == self.max_retries: - print(f"Max number of retries ({self.max_retries}) reached...") - return None - header = { "User-Agent": self.user_agent, "x-api-key": self.api_key, "content-type": "application/json" } - if method.lower() == "get": - - response = requests.get(url, headers=header) - - elif method.lower() == "post": - - response = requests.post(url, headers=header, data=json.dumps(body), files=files) - + # If request is a GET... + if http_method.lower() == "get": + try: + response = self.__requests_retry_session().get(url, headers=header) + except requests.exceptions.RequestException: + raise + except Exception: + raise + + # If request is a POST... + elif http_method.lower() == "post": + + # If there are files involved for uploading... + if files is not None: + header.pop('content-type', None) + try: + response = self.__requests_retry_session().post(url, headers=header, files=files) + except requests.exceptions.RequestException: + raise + except Exception: + raise + + # If there aren't files involved for uploading, send a regular post request. + else: + try: + response = self.__requests_retry_session().post(url, headers=header, data=json.dumps(body)) + except requests.exceptions.RequestException: + raise + except Exception: + raise + + # If request is not a GET or a POST, exit and ask for a supported http method else: - print(f"Unsupported method provided: {method}") + print(f"Unsupported HTTP method provided: {http_method}") print("Please provide a supported HTTP method (GET or POST)") - return None + exit(1) - if response and response.status_code == 503: - - time.sleep(self.retry_wait_time) - print(f"503 error returned, retrying (this was attempt number {self.retry_counter + 1})...") - new_response = self.make_api_request(method, url, body, files, retry=True) + return response - return new_response + def __requests_retry_session(self, backoff_factor=0.5, status_forcelist=(429, 502, 503), session=None): + session = session or requests.Session() + retry = Retry( + total=self.max_retries, + read=self.max_retries, + connect=self.max_retries, + backoff_factor=backoff_factor, + status_forcelist=status_forcelist, + ) + adapter = HTTPAdapter(max_retries=retry) + session.mount('http://', adapter) + session.mount('https://', adapter) + return session + + @staticmethod + def valid_response(response, success_code): - self.reset_retry_counter() + """ + Check to see if API response is valid, and has contains the proper success code. - return response + :param response: Response object from Requests Module + :type response: object - def reset_retry_counter(self): + :param success_code: Expected success status code + :type success_code: int - self.retry_counter = 0 + :return: True/False indicating whether or not response is valid. + :rtype: bool + """ - def increment_retry_counter(self): + if response is not None and response.status_code == success_code: + return True - self.retry_counter += 1 + else: + return False diff --git a/requirements.txt b/requirements.txt index 07def5c..ee0d772 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,4 @@ progressbar2==3.39.2 +urllib3==1.24.1 requests==2.21.0 toml==0.10.0 \ No newline at end of file diff --git a/upload_to_platform.py b/upload_to_platform.py index d957eaf..799c454 100644 --- a/upload_to_platform.py +++ b/upload_to_platform.py @@ -1,4 +1,5 @@ -""" ****************************************************************** +""" +****************************************************************** Name : upload_to_platform.py Project : Upload to Platform @@ -7,7 +8,8 @@ Copyright : (c) RiskSense, Inc. License : Apache-2.0 -****************************************************************** """ +****************************************************************** +""" import json import time @@ -21,6 +23,7 @@ import toml import progressbar +import requests __version__ = "0.5" @@ -47,11 +50,23 @@ def get_client_id(platform, key): url = platform + "/api/v1/client?size=150" - raw_client_id_response = request_handler.make_api_request("GET", url) + raw_client_id_response = None - json_client_id_response = json.loads(raw_client_id_response.text) + try: + raw_client_id_response = request_handler.make_request("GET", url) + + except Exception as ex: + print("ERROR. There was a problem trying to get a list of available client IDs.") + print(ex) + + logging.critical("ERROR. There was a problem trying to get a list of available client IDs.") + logging.critical(ex) + + exit(1) + + if request_handler.valid_response(raw_client_id_response, 200): - if raw_client_id_response.status_code == 200: + json_client_id_response = json.loads(raw_client_id_response.text) if json_client_id_response['page']['totalElements'] == 1: found_id = json_client_id_response['_embedded']['clients'][0]['id'] @@ -116,12 +131,25 @@ def validate_client_id(client, platform, key): url = platform + "/api/v1/client/" + str(client) - raw_client_id_response = request_handler.make_api_request("GET", url) + raw_client_id_response = None + + try: + raw_client_id_response = request_handler.make_request("GET", url) + + except Exception as ex: + print("ERROR. There was a problem validating the client ID.") + print(ex) + + logging.critical("ERROR. There was a problem validating the client ID. (Client ID: %s)", client) + logging.critical("Exception: \n %s", ex) + + exit(1) - json_client_id_response = json.loads(raw_client_id_response.text) + if request_handler.valid_response(raw_client_id_response, 200): + json_client_id_response = json.loads(raw_client_id_response.text) - if raw_client_id_response.status_code == 200 and json_client_id_response['id'] == client: - validity = True + if json_client_id_response['id'] == client: + validity = True return validity @@ -142,6 +170,7 @@ def find_network_id(platform, key, client): :type client: int :return: The selected Network ID + :rtype: int """ request_handler = ApiRequestHandler(key, user_agent=USER_AGENT_STRING) @@ -185,37 +214,46 @@ def find_network_id(platform, key, client): "size": 20 } - raw_network_search_response = request_handler.make_api_request("POST", url, body=body) + raw_network_search_response = None - json_network_search_response = json.loads(raw_network_search_response.text) + try: + raw_network_search_response = request_handler.make_request("POST", url, body=body) - if raw_network_search_response.status_code == 200 and \ - json_network_search_response['page']['totalElements'] != 0: + except Exception as ex: + print("ERROR. There was a problem getting a list of available networks from the platform.") + print(ex) - z = 0 - network_list = [] - while z < len(json_network_search_response['_embedded']['networks']): - if json_network_search_response['_embedded']['networks'][z]['clientId'] == client: - network_list.append([json_network_search_response['_embedded']['networks'][z]['id'], - json_network_search_response['_embedded']['networks'][z]['name']]) - z += 1 + logging.critical("ERROR. There was a problem getting a list of available networks from the platform.") + logging.critical("Exception: \n %s", ex) - y = 0 - while y < len(network_list): - print(f"{y} - {network_list[y][1]}") - y += 1 + exit(1) - list_id = input("Please enter the number that corresponds with your network: ") - network = network_list[int(list_id)][0] + if request_handler.valid_response(raw_network_search_response, 200): + json_network_search_response = json.loads(raw_network_search_response.text) - elif raw_network_search_response.status_code == 200 and \ - json_network_search_response['page']['totalElements'] == 0: + if json_network_search_response['page']['totalElements'] != 0: + z = 0 + network_list = [] + while z < len(json_network_search_response['_embedded']['networks']): + if json_network_search_response['_embedded']['networks'][z]['clientId'] == client: + network_list.append([json_network_search_response['_embedded']['networks'][z]['id'], + json_network_search_response['_embedded']['networks'][z]['name']]) + z += 1 - print() - print("No such network found.") - input("Press ENTER to close.") - print() - exit(1) + y = 0 + while y < len(network_list): + print(f"{y} - {network_list[y][1]}") + y += 1 + + list_id = input("Please enter the number that corresponds with your network: ") + network = network_list[int(list_id)][0] + + else: + print() + print("No such network found.") + input("Press ENTER to close.") + print() + exit(1) else: print(f"An error occurred during the search for your network. Status code " @@ -269,12 +307,24 @@ def create_new_assessment(platform, key, client, name, start_date, notes): "notes": notes } - raw_assessment_response = request_handler.make_api_request("POST", url, body=body) + raw_assessment_response = None - json_assessment_response = json.loads(raw_assessment_response.text) + try: + raw_assessment_response = request_handler.make_request("POST", url, body=body) + + except Exception as ex: + print("ERROR. Unable to create a new assessment.") + print(ex) + + logging.critical("ERROR. There was a problem creating a new assessment.") + logging.critical("Exception: \n %s", ex) + + exit(1) - if raw_assessment_response.status_code == 201: + if request_handler.valid_response(raw_assessment_response, 201): + json_assessment_response = json.loads(raw_assessment_response.text) created_id = json_assessment_response['id'] + else: print(f"Error Creating New Assessment. Status Code returned was {raw_assessment_response.status_code}") @@ -322,12 +372,24 @@ def create_upload(platform, key, assessment, network, client): 'name': upload_name } - raw_upload_response = request_handler.make_api_request("POST", url, body=body) + raw_upload_response = None + + try: + raw_upload_response = request_handler.make_request("POST", url, body=body) + + except Exception as ex: + print("ERROR. There was a problem creating a new upload.") + print(ex) - json_upload_json_response = json.loads(raw_upload_response.text) + logging.critical("ERROR. There was a problem creating a new upload.") + logging.critical("Exception: \n %s", ex) + + exit(1) - if raw_upload_response.status_code == 201: + if request_handler.valid_response(raw_upload_response, 201): + json_upload_json_response = json.loads(raw_upload_response.text) new_upload_id = json_upload_json_response['id'] + else: print(f"Error creating new upload. Status Code returned was {raw_upload_response.status_code}") return @@ -368,9 +430,21 @@ def add_file_to_upload(platform, key, client, upload, file_name, file_path): upload_file = {'scanFile': (file_name, open(file_path + "/" + file_name, 'rb'))} - raw_add_file_response = request_handler.make_api_request("POST", url, files=upload_file) + raw_add_file_response = None + + try: + raw_add_file_response = request_handler.make_request("POST", url, files=upload_file) + + except Exception as ex: + print(f"ERROR. There was a problem adding a file ({file_name})to the upload.") + print(ex) - if raw_add_file_response.status_code != 201: + logging.critical("ERROR. There was a problem adding a file (%s)to the upload.", file_name) + logging.critical("Exception: \n %s", ex) + + exit(1) + + if not request_handler.valid_response(raw_add_file_response, 201): print(f"Error uploading file {file_name}. Status Code returned was {raw_add_file_response.status_code}") print(raw_add_file_response.text) logging.info("Error uploading file " + file_name + ". " @@ -409,14 +483,27 @@ def begin_processing(platform, key, client, upload, run_urba): "autoUrba": run_urba } - raw_begin_processing_response = request_handler.make_api_request("POST", url, body=body) + raw_begin_processing_response = None + + try: + raw_begin_processing_response = request_handler.make_request("POST", url, body=body) + + except Exception as ex: + print("ERROR. There was a problem starting platform processing of the uploaded file(s).") + print(ex) - if raw_begin_processing_response.status_code == 200: + logging.critical("ERROR. There was a problem starting platform processing of the uploaded file(s).") + logging.critical("Exception: \n %s", ex) + + exit(1) + + if request_handler.valid_response(raw_begin_processing_response, 200): print("Uploaded file(s) now processing. This may take a while. Please wait...") else: print("An error has occurred when trying to start processing of your upload(s).") print(raw_begin_processing_response.text) + logging.info(raw_begin_processing_response.text) @@ -447,11 +534,21 @@ def check_upload_state(platform, key, client, upload): url = platform + "/api/v1/client/" + str(client) + "/upload/" + str(upload) - raw_check_upload_state_response = request_handler.make_api_request("GET", url) + raw_check_upload_state_response = None - json_check_upload_state_response = json.loads(raw_check_upload_state_response.text) + try: + raw_check_upload_state_response = request_handler.make_request("GET", url) - if raw_check_upload_state_response.status_code == 200: + except Exception as ex: + print("There was an exception while checking the state of the upload.") + logging.critical("There was an exception while checking the state of the upload.") + logging.critical(ex) + + state = "EXCEPTION" + return state + + if request_handler.valid_response(raw_check_upload_state_response, 200): + json_check_upload_state_response = json.loads(raw_check_upload_state_response.text) state = json_check_upload_state_response['state'] logging.debug("State: %s", state) @@ -557,6 +654,7 @@ def main(): logging.info("No files found to process.") print() input("Please press ENTER to close.") + exit(1) logging.info(" *** Configuration read. Files to process identified. Starting Script. ***") print("*** Configuration read. Files to process identified. Starting Script. ***") From 73efb199aa769bcc868a71a6a81baab081f0b833 Mon Sep 17 00:00:00 2001 From: Burr Webb Date: Fri, 19 Jul 2019 14:24:59 -0600 Subject: [PATCH 4/5] slight refactor to request handler --- api_request_handler/api_request_handler.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/api_request_handler/api_request_handler.py b/api_request_handler/api_request_handler.py index 3254609..ef98aa2 100644 --- a/api_request_handler/api_request_handler.py +++ b/api_request_handler/api_request_handler.py @@ -60,6 +60,8 @@ def make_request(self, http_method, url, body=None, files=None): If max retries are reached, an exception is raised. """ + response = None + header = { "User-Agent": self.user_agent, "x-api-key": self.api_key, From 42ffc563df10baeb310187ed8e38811fa39723fb Mon Sep 17 00:00:00 2001 From: Burr Webb Date: Fri, 19 Jul 2019 14:55:13 -0600 Subject: [PATCH 5/5] update to version information, and refactoring of imports --- upload_to_platform.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/upload_to_platform.py b/upload_to_platform.py index 799c454..1a95b29 100644 --- a/upload_to_platform.py +++ b/upload_to_platform.py @@ -23,9 +23,8 @@ import toml import progressbar -import requests -__version__ = "0.5" +__version__ = "0.5.0" USER_AGENT_STRING = "upload_to_platform_v" + __version__