From ba56b902c1074b1f30efcac5affe90a495651310 Mon Sep 17 00:00:00 2001 From: AlexButeau Date: Wed, 9 Apr 2025 08:59:57 +0200 Subject: [PATCH 1/2] feat: add possibility for custom certificate in requests --- tofupilot/client.py | 35 +++++++++++++++++++++++++++++++---- tofupilot/openhtf/upload.py | 11 ++++++++++- tofupilot/utils/files.py | 14 ++++++++++++-- 3 files changed, 53 insertions(+), 7 deletions(-) diff --git a/tofupilot/client.py b/tofupilot/client.py index 56e7427..5bda516 100644 --- a/tofupilot/client.py +++ b/tofupilot/client.py @@ -36,7 +36,12 @@ class TofuPilotClient: """Wrapper for TofuPilot's API that provides additional support for handling attachments.""" - def __init__(self, api_key: Optional[str] = None, url: Optional[str] = None): + def __init__( + self, + api_key: Optional[str] = None, + url: Optional[str] = None, + custom_certificate_path: Optional[str] = None, + ): self._current_version = version("tofupilot") print_version_banner(self._current_version) self._logger = setup_logger(logging.INFO) @@ -52,6 +57,7 @@ def __init__(self, api_key: Optional[str] = None, url: Optional[str] = None): "Content-Type": "application/json", "Authorization": f"Bearer {self._api_key}", } + self._custom_certificate_path = custom_certificate_path self._max_attachments = CLIENT_MAX_ATTACHMENTS self._max_file_size = FILE_MAX_SIZE check_latest_version(self._logger, self._current_version, "tofupilot") @@ -159,6 +165,7 @@ def create_run( # pylint: disable=too-many-arguments,too-many-locals json=payload, headers=self._headers, timeout=SECONDS_BEFORE_TIMEOUT, + cert=self._custom_certificate_path, ) response.raise_for_status() result = handle_response(self._logger, response) @@ -166,7 +173,12 @@ def create_run( # pylint: disable=too-many-arguments,too-many-locals run_id = result.get("id") if run_id and attachments: upload_attachments( - self._logger, self._headers, self._url, attachments, run_id + self._logger, + self._headers, + self._url, + attachments, + run_id, + self._custom_certificate_path, ) return result @@ -229,6 +241,7 @@ def create_run_from_openhtf_report(self, file_path: str): initialize_url, data=json.dumps(payload), headers=self._headers, + cert=self._custom_certificate_path, timeout=SECONDS_BEFORE_TIMEOUT, ) @@ -249,7 +262,13 @@ def create_run_from_openhtf_report(self, file_path: str): timeout=SECONDS_BEFORE_TIMEOUT, ) - notify_server(self._headers, self._url, upload_id, run_id) + notify_server( + self._headers, + self._url, + upload_id, + run_id, + self._custom_certificate_path, + ) self._logger.success( "Attachment %s successfully uploaded and linked to run.", @@ -293,6 +312,7 @@ def get_runs(self, serial_number: str) -> dict: response = requests.get( f"{self._url}/runs", headers=self._headers, + cert=self._custom_certificate_path, params=params, timeout=SECONDS_BEFORE_TIMEOUT, ) @@ -326,6 +346,7 @@ def delete_run(self, run_id: str) -> dict: response = requests.delete( f"{self._url}/runs/{run_id}", headers=self._headers, + cert=self._custom_certificate_path, timeout=SECONDS_BEFORE_TIMEOUT, ) response.raise_for_status() @@ -366,6 +387,7 @@ def update_unit( f"{self._url}/units/{serial_number}", json=payload, headers=self._headers, + cert=self._custom_certificate_path, timeout=SECONDS_BEFORE_TIMEOUT, ) response.raise_for_status() @@ -399,6 +421,7 @@ def delete_unit(self, serial_number: str) -> dict: response = requests.delete( f"{self._url}/units/{serial_number}", headers=self._headers, + cert=self._custom_certificate_path, timeout=SECONDS_BEFORE_TIMEOUT, ) response.raise_for_status() @@ -430,7 +453,9 @@ def upload_and_create_from_openhtf_report( # Upload report try: - upload_id = upload_file(self._headers, self._url, file_path) + upload_id = upload_file( + self._headers, self._url, file_path, self._custom_certificate_path + ) except requests.exceptions.HTTPError as http_err: return handle_http_error(self._logger, http_err) except requests.RequestException as e: @@ -451,6 +476,7 @@ def upload_and_create_from_openhtf_report( f"{self._url}/import", json=payload, headers=self._headers, + cert=self._custom_certificate_path, timeout=SECONDS_BEFORE_TIMEOUT, ) response.raise_for_status() @@ -478,6 +504,7 @@ def get_websocket_url(self) -> dict: response = requests.get( f"{self._url}/rooms", headers=self._headers, + cert=self._custom_certificate_path, timeout=SECONDS_BEFORE_TIMEOUT, ) response.raise_for_status() diff --git a/tofupilot/openhtf/upload.py b/tofupilot/openhtf/upload.py index 42a3a93..db3081d 100644 --- a/tofupilot/openhtf/upload.py +++ b/tofupilot/openhtf/upload.py @@ -48,12 +48,14 @@ def __init__( allow_nan: Optional[bool] = False, url: Optional[str] = None, client: Optional[TofuPilotClient] = None, + custom_certificate_path: Optional[str] = None, ): self.allow_nan = allow_nan self.client = client or TofuPilotClient(api_key=api_key, url=url) self._logger = self.client._logger self._url = self.client._url self._headers = self.client._headers + self._custom_certificate_path = custom_certificate_path self._max_attachments = self.client._max_attachments self._max_file_size = self.client._max_file_size @@ -133,6 +135,7 @@ def __call__(self, test_record: TestRecord): initialize_url, data=json.dumps(payload), headers=self._headers, + cert=self._custom_certificate_path, timeout=SECONDS_BEFORE_TIMEOUT, ) @@ -148,7 +151,13 @@ def __call__(self, test_record: TestRecord): timeout=SECONDS_BEFORE_TIMEOUT, ) - notify_server(self._headers, self._url, upload_id, run_id) + notify_server( + self._headers, + self._url, + upload_id, + run_id, + self._custom_certificate_path, + ) self._logger.success( "Attachment %s successfully uploaded and linked to run.", diff --git a/tofupilot/utils/files.py b/tofupilot/utils/files.py index e8d3cf1..cb641ec 100644 --- a/tofupilot/utils/files.py +++ b/tofupilot/utils/files.py @@ -46,6 +46,7 @@ def upload_file( headers: dict, url: str, file_path: str, + custom_certificate_path: Optional[str] = None, ) -> bool: """Initializes an upload and stores file in it""" # Upload initialization @@ -58,6 +59,7 @@ def upload_file( data=json.dumps(payload), headers=headers, timeout=SECONDS_BEFORE_TIMEOUT, + cert=custom_certificate_path, ) response.raise_for_status() @@ -78,7 +80,13 @@ def upload_file( return upload_id -def notify_server(headers: dict, url: str, upload_id: str, run_id: str) -> bool: +def notify_server( + headers: dict, + url: str, + upload_id: str, + run_id: str, + custom_certificate_path: Optional[str] = None, +) -> bool: """Tells TP server to sync upload with newly created run""" sync_url = f"{url}/uploads/sync" sync_payload = {"upload_id": upload_id, "run_id": run_id} @@ -86,6 +94,7 @@ def notify_server(headers: dict, url: str, upload_id: str, run_id: str) -> bool: response = requests.post( sync_url, data=json.dumps(sync_payload), + cert=custom_certificate_path, headers=headers, timeout=SECONDS_BEFORE_TIMEOUT, ) @@ -99,13 +108,14 @@ def upload_attachments( url: str, paths: List[Dict[str, Optional[str]]], run_id: str, + custom_certificate_path: Optional[str] = None, ): """Creates one upload per file and stores them into TofuPilot""" for file_path in paths: logger.info("Uploading %s...", file_path) upload_id = upload_file(headers, url, file_path) - notify_server(headers, url, upload_id, run_id) + notify_server(headers, url, upload_id, run_id, custom_certificate_path) logger.success( f"Attachment {file_path} successfully uploaded and linked to run." From 9982e4837345c6c69eb2ec4ec33c8cc4a5d54c35 Mon Sep 17 00:00:00 2001 From: Julien Buteau Date: Wed, 16 Apr 2025 10:13:51 +0200 Subject: [PATCH 2/2] feat: adds custom certificate verification Allows users to specify a custom certificate path for verifying the TofuPilot server's certificate. This is useful for connecting to instances with custom or self-signed certificates, enhancing security when dealing with non-standard deployments. The 'custom_certificate_path' parameter is renamed to 'verify' for clarity and consistency with requests library's naming convention. --- tofupilot/client.py | 37 ++++++++++++++++----------- tofupilot/openhtf/upload.py | 18 +++++++++---- tofupilot/utils/files.py | 50 +++++++++++++++++++++++++++++-------- 3 files changed, 76 insertions(+), 29 deletions(-) diff --git a/tofupilot/client.py b/tofupilot/client.py index 5bda516..088880f 100644 --- a/tofupilot/client.py +++ b/tofupilot/client.py @@ -34,13 +34,22 @@ class TofuPilotClient: - """Wrapper for TofuPilot's API that provides additional support for handling attachments.""" + """Wrapper for TofuPilot's API that provides additional support for handling attachments. + + Args: + api_key (Optional[str]): API key for authentication with TofuPilot's API. + If not provided, the TOFUPILOT_API_KEY environment variable will be used. + url (Optional[str]): Base URL for TofuPilot's API. + If not provided, the TOFUPILOT_URL environment variable or the default endpoint will be used. + verify (Optional[str]): Path to a CA bundle file to verify TofuPilot's server certificate. + Useful for connecting to instances with custom/self-signed certificates. + """ def __init__( self, api_key: Optional[str] = None, url: Optional[str] = None, - custom_certificate_path: Optional[str] = None, + verify: Optional[str] = None, ): self._current_version = version("tofupilot") print_version_banner(self._current_version) @@ -57,7 +66,7 @@ def __init__( "Content-Type": "application/json", "Authorization": f"Bearer {self._api_key}", } - self._custom_certificate_path = custom_certificate_path + self._verify = verify self._max_attachments = CLIENT_MAX_ATTACHMENTS self._max_file_size = FILE_MAX_SIZE check_latest_version(self._logger, self._current_version, "tofupilot") @@ -165,7 +174,7 @@ def create_run( # pylint: disable=too-many-arguments,too-many-locals json=payload, headers=self._headers, timeout=SECONDS_BEFORE_TIMEOUT, - cert=self._custom_certificate_path, + verify=self._verify, ) response.raise_for_status() result = handle_response(self._logger, response) @@ -178,7 +187,7 @@ def create_run( # pylint: disable=too-many-arguments,too-many-locals self._url, attachments, run_id, - self._custom_certificate_path, + self._verify, ) return result @@ -241,7 +250,7 @@ def create_run_from_openhtf_report(self, file_path: str): initialize_url, data=json.dumps(payload), headers=self._headers, - cert=self._custom_certificate_path, + verify=self._verify, timeout=SECONDS_BEFORE_TIMEOUT, ) @@ -267,7 +276,7 @@ def create_run_from_openhtf_report(self, file_path: str): self._url, upload_id, run_id, - self._custom_certificate_path, + self._verify, ) self._logger.success( @@ -312,7 +321,7 @@ def get_runs(self, serial_number: str) -> dict: response = requests.get( f"{self._url}/runs", headers=self._headers, - cert=self._custom_certificate_path, + verify=self._verify, params=params, timeout=SECONDS_BEFORE_TIMEOUT, ) @@ -346,7 +355,7 @@ def delete_run(self, run_id: str) -> dict: response = requests.delete( f"{self._url}/runs/{run_id}", headers=self._headers, - cert=self._custom_certificate_path, + verify=self._verify, timeout=SECONDS_BEFORE_TIMEOUT, ) response.raise_for_status() @@ -387,7 +396,7 @@ def update_unit( f"{self._url}/units/{serial_number}", json=payload, headers=self._headers, - cert=self._custom_certificate_path, + verify=self._verify, timeout=SECONDS_BEFORE_TIMEOUT, ) response.raise_for_status() @@ -421,7 +430,7 @@ def delete_unit(self, serial_number: str) -> dict: response = requests.delete( f"{self._url}/units/{serial_number}", headers=self._headers, - cert=self._custom_certificate_path, + verify=self._verify, timeout=SECONDS_BEFORE_TIMEOUT, ) response.raise_for_status() @@ -454,7 +463,7 @@ def upload_and_create_from_openhtf_report( # Upload report try: upload_id = upload_file( - self._headers, self._url, file_path, self._custom_certificate_path + self._headers, self._url, file_path, self._verify ) except requests.exceptions.HTTPError as http_err: return handle_http_error(self._logger, http_err) @@ -476,7 +485,7 @@ def upload_and_create_from_openhtf_report( f"{self._url}/import", json=payload, headers=self._headers, - cert=self._custom_certificate_path, + verify=self._verify, timeout=SECONDS_BEFORE_TIMEOUT, ) response.raise_for_status() @@ -504,7 +513,7 @@ def get_websocket_url(self) -> dict: response = requests.get( f"{self._url}/rooms", headers=self._headers, - cert=self._custom_certificate_path, + verify=self._verify, timeout=SECONDS_BEFORE_TIMEOUT, ) response.raise_for_status() diff --git a/tofupilot/openhtf/upload.py b/tofupilot/openhtf/upload.py index db3081d..56737b6 100644 --- a/tofupilot/openhtf/upload.py +++ b/tofupilot/openhtf/upload.py @@ -24,6 +24,14 @@ class upload: # pylint: disable=invalid-name This function behaves similarly to manually parsing the OpenHTF JSON test report and calling `TofuPilotClient().create_run()` with the parsed data. + Args: + api_key (Optional[str]): API key for authentication with TofuPilot's API. + allow_nan (Optional[bool]): Whether to allow NaN values in JSON serialization. + url (Optional[str]): Base URL for TofuPilot's API. + client (Optional[TofuPilotClient]): An existing TofuPilot client instance to use. + verify (Optional[str]): Path to a CA bundle file to verify TofuPilot's server certificate. + Useful for connecting to instances with custom/self-signed certificates. + ### Usage Example: ```python @@ -48,14 +56,14 @@ def __init__( allow_nan: Optional[bool] = False, url: Optional[str] = None, client: Optional[TofuPilotClient] = None, - custom_certificate_path: Optional[str] = None, + verify: Optional[str] = None, ): self.allow_nan = allow_nan - self.client = client or TofuPilotClient(api_key=api_key, url=url) + self.client = client or TofuPilotClient(api_key=api_key, url=url, verify=verify) self._logger = self.client._logger self._url = self.client._url self._headers = self.client._headers - self._custom_certificate_path = custom_certificate_path + self._verify = verify self._max_attachments = self.client._max_attachments self._max_file_size = self.client._max_file_size @@ -135,7 +143,7 @@ def __call__(self, test_record: TestRecord): initialize_url, data=json.dumps(payload), headers=self._headers, - cert=self._custom_certificate_path, + verify=self._verify, timeout=SECONDS_BEFORE_TIMEOUT, ) @@ -156,7 +164,7 @@ def __call__(self, test_record: TestRecord): self._url, upload_id, run_id, - self._custom_certificate_path, + self._verify, ) self._logger.success( diff --git a/tofupilot/utils/files.py b/tofupilot/utils/files.py index cb641ec..5dd3616 100644 --- a/tofupilot/utils/files.py +++ b/tofupilot/utils/files.py @@ -46,9 +46,19 @@ def upload_file( headers: dict, url: str, file_path: str, - custom_certificate_path: Optional[str] = None, + verify: Optional[str] = None, ) -> bool: - """Initializes an upload and stores file in it""" + """Initializes an upload and stores file in it + + Args: + headers (dict): Request headers including authorization + url (str): Base API URL + file_path (str): Path to the file to upload + verify (Optional[str]): Path to a CA bundle file to verify the server certificate + + Returns: + str: The ID of the created upload + """ # Upload initialization initialize_url = f"{url}/uploads/initialize" file_name = os.path.basename(file_path) @@ -59,7 +69,7 @@ def upload_file( data=json.dumps(payload), headers=headers, timeout=SECONDS_BEFORE_TIMEOUT, - cert=custom_certificate_path, + verify=verify, ) response.raise_for_status() @@ -85,16 +95,27 @@ def notify_server( url: str, upload_id: str, run_id: str, - custom_certificate_path: Optional[str] = None, + verify: Optional[str] = None, ) -> bool: - """Tells TP server to sync upload with newly created run""" + """Tells TP server to sync upload with newly created run + + Args: + headers (dict): Request headers including authorization + url (str): Base API URL + upload_id (str): ID of the upload to link + run_id (str): ID of the run to link to + verify (Optional[str]): Path to a CA bundle file to verify the server certificate + + Returns: + bool: True if successful + """ sync_url = f"{url}/uploads/sync" sync_payload = {"upload_id": upload_id, "run_id": run_id} response = requests.post( sync_url, data=json.dumps(sync_payload), - cert=custom_certificate_path, + verify=verify, headers=headers, timeout=SECONDS_BEFORE_TIMEOUT, ) @@ -108,14 +129,23 @@ def upload_attachments( url: str, paths: List[Dict[str, Optional[str]]], run_id: str, - custom_certificate_path: Optional[str] = None, + verify: Optional[str] = None, ): - """Creates one upload per file and stores them into TofuPilot""" + """Creates one upload per file and stores them into TofuPilot + + Args: + logger (Logger): Logger instance + headers (dict): Request headers including authorization + url (str): Base API URL + paths (List[Dict[str, Optional[str]]]): List of file paths to upload + run_id (str): ID of the run to link files to + verify (Optional[str]): Path to a CA bundle file to verify the server certificate + """ for file_path in paths: logger.info("Uploading %s...", file_path) - upload_id = upload_file(headers, url, file_path) - notify_server(headers, url, upload_id, run_id, custom_certificate_path) + upload_id = upload_file(headers, url, file_path, verify) + notify_server(headers, url, upload_id, run_id, verify) logger.success( f"Attachment {file_path} successfully uploaded and linked to run."