diff --git a/roboflow/adapters/rfapi.py b/roboflow/adapters/rfapi.py index 2b8cac04..fc1d57c4 100644 --- a/roboflow/adapters/rfapi.py +++ b/roboflow/adapters/rfapi.py @@ -10,6 +10,7 @@ from requests_toolbelt.multipart.encoder import MultipartEncoder from roboflow.config import API_URL, DEFAULT_BATCH_NAME, DEFAULT_JOB_NAME +from roboflow.util.redact import redact_api_key from roboflow.util.search_utils import normalize_media_types @@ -21,23 +22,26 @@ class RoboflowError(Exception): callers can branch on auth (401) vs not-found (404) without string matching the message. Existing call sites that pass only a message still work; the attribute defaults to `None`. + + API keys in the message (for example in the URL of a wrapped `requests` + exception) are replaced with `***`. """ def __init__(self, message, status_code=None): - super().__init__(message) + super().__init__(redact_api_key(message)) self.status_code = status_code class ImageUploadError(RoboflowError): def __init__(self, message, status_code=None): - self.message = message + self.message = redact_api_key(message) self.retries = 0 super().__init__(self.message, status_code=status_code) class AnnotationSaveError(RoboflowError): def __init__(self, message, status_code=None): - self.message = message + self.message = redact_api_key(message) self.retries = 0 super().__init__(self.message, status_code=status_code) diff --git a/roboflow/cli/__init__.py b/roboflow/cli/__init__.py index d8588ad5..da954644 100644 --- a/roboflow/cli/__init__.py +++ b/roboflow/cli/__init__.py @@ -429,8 +429,20 @@ def main() -> None: except Exception as exc: import json as _json - payload = {"error": {"message": str(exc)}} + from roboflow.util.redact import redact_api_key + + payload = {"error": {"message": redact_api_key(str(exc))}} print(_json.dumps(payload), file=sys.stderr) sys.exit(1) else: - app() + try: + app() + except Exception as exc: + # Print the traceback as Python would, but without the API key that + # request URLs in `requests` exception messages may contain. + import traceback + + from roboflow.util.redact import redact_api_key + + sys.stderr.write(redact_api_key("".join(traceback.format_exception(exc)))) + sys.exit(1) diff --git a/roboflow/cli/_output.py b/roboflow/cli/_output.py index 70ba2c6c..91e8b60a 100644 --- a/roboflow/cli/_output.py +++ b/roboflow/cli/_output.py @@ -88,12 +88,9 @@ def _translate_api_hints(message: str) -> str: def _sanitize_credentials(text: str) -> str: """Strip API keys from URLs and other sensitive patterns in error messages.""" - import re + from roboflow.util.redact import redact_api_key - # Match api_key=... up to the next whitespace, query separator, quote, or backslash. - # Older patterns missed keys containing '-' or other URL-safe characters and would - # echo them to the terminal when an exception bubbled up from `requests`. - return re.sub(r"api_key=[^\s&\"'\\<>]+", "api_key=***", text) + return redact_api_key(text) def _parse_error_message(raw: str) -> tuple[Optional[dict[str, Any]], str]: diff --git a/roboflow/core/version.py b/roboflow/core/version.py index a4633a3b..3d6fa8ac 100644 --- a/roboflow/core/version.py +++ b/roboflow/core/version.py @@ -37,6 +37,7 @@ from roboflow.util.annotations import amend_data_yaml from roboflow.util.general import extract_zip, write_line from roboflow.util.model_processor import package_custom_weights_interactive, validate_model_type_for_project +from roboflow.util.redact import redact_api_key from roboflow.util.train_recipe import fold_epochs_into_recipe from roboflow.util.versions import get_model_format, get_wrong_dependencies_versions @@ -685,7 +686,7 @@ def _upload_zip(self, model_type: str, model_path: str, model_file_name: str): else: res.raise_for_status() except Exception as e: - print(f"An error occured when getting the model upload URL: {e}") + print(f"An error occured when getting the model upload URL: {redact_api_key(str(e))}") return res = requests.put( @@ -710,7 +711,7 @@ def _upload_zip(self, model_type: str, model_path: str, model_file_name: str): ) except Exception as e: - print(f"An error occured when uploading the model: {e}") + print(f"An error occured when uploading the model: {redact_api_key(str(e))}") def __download_zip(self, link, location, format): """ diff --git a/roboflow/core/workspace.py b/roboflow/core/workspace.py index 081e9ac6..b16b2c2f 100644 --- a/roboflow/core/workspace.py +++ b/roboflow/core/workspace.py @@ -17,6 +17,7 @@ from roboflow.adapters import rfapi, vision_events_api from roboflow.adapters.rfapi import AnnotationSaveError, ImageUploadError, RoboflowError from roboflow.config import API_URL, APP_URL, DEMO_KEYS +from roboflow.util.redact import redact_api_key if TYPE_CHECKING: from roboflow.core.device import Device @@ -713,7 +714,7 @@ def _upload(imagedesc): annotation_msg = f"annotations = ERR: {e.message}" print(f"{image_msg} / {annotation_msg}") except Exception as e: - print(f"[ERR] {image_path} ({e})") + print(f"[ERR] {image_path} ({redact_api_key(str(e))})") with concurrent.futures.ThreadPoolExecutor(max_workers=num_workers) as executor: list(executor.map(_upload, images)) @@ -946,7 +947,7 @@ def _upload_zip( try: res.raise_for_status() except Exception as e: - error_message = str(e) + error_message = redact_api_key(str(e)) status_code = str(res.status_code) print("\n\033[91m❌ ERROR\033[0m: Failed to get model deployment URL") @@ -970,7 +971,7 @@ def _upload_zip( ) except Exception as e: - print(f"An error occured when uploading the model: {e}") + print(f"An error occured when uploading the model: {redact_api_key(str(e))}") def search( self, diff --git a/roboflow/util/redact.py b/roboflow/util/redact.py new file mode 100644 index 00000000..1b0fdcc7 --- /dev/null +++ b/roboflow/util/redact.py @@ -0,0 +1,19 @@ +"""Keep API keys out of messages that are printed or end up in logs.""" + +import re +from typing import Any + +# Match api_key=... up to the next whitespace, query separator, quote, or backslash. +_API_KEY_PARAM = re.compile(r"api_key=[^\s&\"'\\<>]+") + + +def redact_api_key(text: Any) -> Any: + """Replace the value of every ``api_key=`` URL parameter in *text* with ``***``. + + Most SDK requests still pass the API key as a query parameter, and the + messages of ``requests`` exceptions contain the full request URL. + Values other than strings are returned unchanged. + """ + if not isinstance(text, str): + return text + return _API_KEY_PARAM.sub("api_key=***", text) diff --git a/tests/cli/test_redact_errors.py b/tests/cli/test_redact_errors.py new file mode 100644 index 00000000..22fb0f66 --- /dev/null +++ b/tests/cli/test_redact_errors.py @@ -0,0 +1,53 @@ +"""Unhandled errors in the CLI entry point must not print the API key.""" + +import contextlib +import io +import json +import unittest +from unittest.mock import patch + +import requests + +from roboflow.cli import main + +API_KEY = "abc-DEF_123" +CONNECTION_ERROR = requests.ConnectionError( + f"HTTPSConnectionPool(host='api.roboflow.com', port=443): Max retries exceeded with url: /ws/proj?api_key={API_KEY}" +) + + +def _run_main(argv: list) -> tuple: + stderr = io.StringIO() + with ( + patch("sys.argv", ["roboflow", *argv]), + patch("roboflow.adapters.rfapi.get_project", side_effect=CONNECTION_ERROR), + contextlib.redirect_stderr(stderr), + contextlib.redirect_stdout(io.StringIO()), + ): + try: + main() + except SystemExit as exc: + return exc.code, stderr.getvalue() + return 0, stderr.getvalue() + + +class TestUnhandledErrorsAreRedacted(unittest.TestCase): + def test_json_mode(self) -> None: + code, stderr = _run_main(["--json", "-k", API_KEY, "-w", "ws", "project", "get", "proj"]) + + self.assertEqual(code, 1) + self.assertNotIn(API_KEY, stderr) + self.assertIn("api_key=***", json.loads(stderr)["error"]["message"]) + + def test_text_mode_prints_redacted_traceback(self) -> None: + code, stderr = _run_main(["-k", API_KEY, "-w", "ws", "project", "get", "proj"]) + + self.assertEqual(code, 1) + self.assertNotIn(API_KEY, stderr) + self.assertIn("Traceback (most recent call last)", stderr) + self.assertIn("requests.exceptions.ConnectionError", stderr) + self.assertIn("api_key=***", stderr) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/util/test_redact.py b/tests/util/test_redact.py new file mode 100644 index 00000000..ecffbdee --- /dev/null +++ b/tests/util/test_redact.py @@ -0,0 +1,64 @@ +import contextlib +import io +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import responses + +from roboflow.adapters.rfapi import AnnotationSaveError, ImageUploadError, RoboflowError +from roboflow.util.redact import redact_api_key +from tests.helpers import get_version + +API_KEY = "abc-DEF_123" + + +class TestRedactApiKey(unittest.TestCase): + def test_redacts_query_parameter(self) -> None: + text = f"Max retries exceeded with url: /ws/proj?api_key={API_KEY}&nocache=true (Caused by ...)" + self.assertEqual( + redact_api_key(text), + "Max retries exceeded with url: /ws/proj?api_key=***&nocache=true (Caused by ...)", + ) + + def test_redacts_every_occurrence(self) -> None: + text = f"'https://a/b?api_key={API_KEY}' and \"https://c/d?x=1&api_key={API_KEY}\"" + self.assertNotIn(API_KEY, redact_api_key(text)) + self.assertEqual(redact_api_key(text).count("api_key=***"), 2) + + def test_leaves_other_values_unchanged(self) -> None: + self.assertEqual(redact_api_key("no key here"), "no key here") + self.assertEqual(redact_api_key({"api_key": API_KEY}), {"api_key": API_KEY}) + + +class TestErrorMessages(unittest.TestCase): + def test_exceptions_redact_their_message(self) -> None: + message = f"500 Server Error for url: https://api.roboflow.com/ws/proj?api_key={API_KEY}" + for error_class in (RoboflowError, ImageUploadError, AnnotationSaveError): + error = error_class(message) + self.assertNotIn(API_KEY, str(error)) + self.assertIn("api_key=***", str(error)) + self.assertNotIn(API_KEY, ImageUploadError(message).message) + self.assertNotIn(API_KEY, AnnotationSaveError(message).message) + + def test_version_deploy_error_is_printed_without_key(self) -> None: + bundle = SimpleNamespace(model_type="yolov8n", archive_path=SimpleNamespace(name="roboflow_deploy.zip")) + stdout = io.StringIO() + with ( + responses.RequestsMock() as rsps, + patch("roboflow.core.version.package_custom_weights_interactive", return_value=bundle), + contextlib.redirect_stdout(stdout), + ): + rsps.add( + responses.GET, + "https://api.roboflow.com/Test Workspace Name/Test Project Name/1/uploadModel", + status=500, + ) + get_version(api_key=API_KEY).deploy("yolov8n", "weights-dir", "weights/best.pt") + + self.assertIn("api_key=***", stdout.getvalue()) + self.assertNotIn(API_KEY, stdout.getvalue()) + + +if __name__ == "__main__": + unittest.main()