From 2dc12defd8aacf3fcac9ef3651aaad84507b9408 Mon Sep 17 00:00:00 2001 From: Punit Naik Date: Fri, 20 Mar 2026 14:43:21 +0530 Subject: [PATCH] feat(encryption): Encrypt manifest credentials via chuck-api Add encrypt_credential method to AmperityAPIClient that calls the chuck-api /api/encrypt endpoint. Snowflake manifest setup now encrypts passwords and private keys before embedding them, prefixed with ENC:. Jira ticket: [CATALYST-129](https://amperity.atlassian.net/browse/CATALYST-129) --- chuck_data/clients/amperity.py | 53 +++++++++++++++++++++ chuck_data/commands/setup_stitch.py | 46 ++++++++++++++++++- tests/unit/clients/test_amperity.py | 71 +++++++++++++++++++++++++++++ 3 files changed, 168 insertions(+), 2 deletions(-) diff --git a/chuck_data/clients/amperity.py b/chuck_data/clients/amperity.py index cd77c03..c4b9fa2 100644 --- a/chuck_data/clients/amperity.py +++ b/chuck_data/clients/amperity.py @@ -319,6 +319,59 @@ def fetch_amperity_job_init( logging.debug(f"Connection error: {e}") raise ConnectionError(f"Connection error occurred: {e}") + def encrypt_credential(self, plaintext: str, token: str) -> str: + """Encrypt a credential for safe storage in a manifest. + + Calls the chuck-api /api/encrypt endpoint which uses a derived + AES-256 key to encrypt the plaintext. + + Args: + plaintext: The credential value to encrypt (e.g. password, PAT). + token: Amperity CLI authentication token. + + Returns: + The encrypted credential string with ENC: prefix + (e.g. "ENC:base64ciphertext..."). + + Raises: + ValueError: If the API returns an error. + ConnectionError: If connection to Amperity API fails. + """ + try: + url = f"https://{self.base_url}/api/encrypt" + headers = { + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + } + payload = {"text": plaintext} + + response = requests.post( + url, + headers=headers, + data=json.dumps(payload), + timeout=30, + ) + response.raise_for_status() + data = response.json() + encrypted = data.get("encrypted") + if not encrypted: + raise ValueError("API returned empty encrypted value") + return f"ENC:{encrypted}" + + except requests.exceptions.HTTPError as e: + resp = e.response + resp_text = resp.text if resp else "" + logging.error( + "Encrypt credential HTTP error: %s, Response: %s", + e, + resp_text, + ) + raise ValueError(f"Failed to encrypt credential: {resp_text}") + + except requests.RequestException as e: + logging.error("Encrypt credential connection error: %s", e) + raise ConnectionError(f"Connection error during encryption: {e}") + def record_job_submission( self, databricks_run_id: str, token: str, job_id: str ) -> bool: diff --git a/chuck_data/commands/setup_stitch.py b/chuck_data/commands/setup_stitch.py index 8e8192b..dde75e2 100644 --- a/chuck_data/commands/setup_stitch.py +++ b/chuck_data/commands/setup_stitch.py @@ -2311,13 +2311,55 @@ def _snowflake_prepare_manifest( # Step 3: Generate manifest console.print("\nStep 3: Generating manifest...") + + # Encrypt credentials via chuck-api before embedding in manifest + from chuck_data.clients.amperity import AmperityAPIClient + + raw_password = getattr(client, "_password", None) + raw_private_key = getattr(client, "_pem_private_key", None) + encrypted_password = None + encrypted_private_key = None + + amperity_token = get_amperity_token() + if amperity_token and (raw_password or raw_private_key): + amp_client = AmperityAPIClient() + try: + if raw_password: + encrypted_password = amp_client.encrypt_credential( + raw_password, amperity_token + ) + console.print( + f"[{SUCCESS_STYLE}]Encrypted password for manifest" + f" storage[/{SUCCESS_STYLE}]" + ) + if raw_private_key: + encrypted_private_key = amp_client.encrypt_credential( + raw_private_key, amperity_token + ) + console.print( + f"[{SUCCESS_STYLE}]Encrypted private key for manifest" + f" storage[/{SUCCESS_STYLE}]" + ) + except Exception as e: + logging.warning("Could not encrypt credentials: %s", e) + console.print( + f"[{WARNING}]Could not encrypt credentials via API, " + f"storing in plaintext: {e}[/{WARNING}]" + ) + encrypted_password = raw_password + encrypted_private_key = raw_private_key + else: + # No amperity token — fall back to plaintext + encrypted_password = raw_password + encrypted_private_key = raw_private_key + snowflake_conn = { "account": getattr(client, "account", get_snowflake_account() or ""), "user": getattr(client, "user", get_snowflake_user() or ""), "warehouse": getattr(client, "warehouse", get_snowflake_warehouse() or ""), "role": getattr(client, "role", get_snowflake_role() or ""), - "password": getattr(client, "_password", None), - "private_key": getattr(client, "_pem_private_key", None), + "password": encrypted_password, + "private_key": encrypted_private_key, } manifest = generate_snowflake_manifest( diff --git a/tests/unit/clients/test_amperity.py b/tests/unit/clients/test_amperity.py index 9c3de89..dda1ef0 100644 --- a/tests/unit/clients/test_amperity.py +++ b/tests/unit/clients/test_amperity.py @@ -220,3 +220,74 @@ def test_record_job_submission_payload_format(mock_post): payload = json.loads(call_args[1]["data"]) assert "databricks-run-id" in payload # kebab-case, not snake_case assert "job-id" in payload # kebab-case, not snake_case + + +# --- encrypt_credential tests --- + + +@patch("chuck_data.clients.amperity.requests.post") +def test_encrypt_credential_success(mock_post): + """Test successful credential encryption.""" + client = AmperityAPIClient() + + mock_response = Mock() + mock_response.status_code = 200 + mock_response.json.return_value = {"encrypted": "abc123encrypted=="} + mock_response.raise_for_status = Mock() + mock_post.return_value = mock_response + + result = client.encrypt_credential("my-secret-password", "test-token") + + assert result == "ENC:abc123encrypted==" + + # Verify the request was made correctly + mock_post.assert_called_once() + call_args = mock_post.call_args + assert "api/encrypt" in call_args[0][0] + assert call_args[1]["headers"]["Authorization"] == "Bearer test-token" + + payload = json.loads(call_args[1]["data"]) + assert payload["text"] == "my-secret-password" + + +@patch("chuck_data.clients.amperity.requests.post") +def test_encrypt_credential_http_error(mock_post): + """Test encrypt_credential with HTTP error.""" + client = AmperityAPIClient() + + mock_response = Mock() + mock_response.status_code = 500 + mock_response.text = "Internal Server Error" + mock_response.raise_for_status.side_effect = requests.exceptions.HTTPError( + response=mock_response + ) + mock_post.return_value = mock_response + + with pytest.raises(ValueError, match="Failed to encrypt"): + client.encrypt_credential("secret", "test-token") + + +@patch("chuck_data.clients.amperity.requests.post") +def test_encrypt_credential_connection_error(mock_post): + """Test encrypt_credential with network error.""" + client = AmperityAPIClient() + + mock_post.side_effect = requests.exceptions.ConnectionError("Network error") + + with pytest.raises(ConnectionError, match="Connection error"): + client.encrypt_credential("secret", "test-token") + + +@patch("chuck_data.clients.amperity.requests.post") +def test_encrypt_credential_empty_response(mock_post): + """Test encrypt_credential when API returns empty encrypted value.""" + client = AmperityAPIClient() + + mock_response = Mock() + mock_response.status_code = 200 + mock_response.json.return_value = {"encrypted": ""} + mock_response.raise_for_status = Mock() + mock_post.return_value = mock_response + + with pytest.raises(ValueError, match="empty encrypted value"): + client.encrypt_credential("secret", "test-token")