Merge pull request #695 from Dstack-TEE/feat/sdk-python-api-parity

feat(sdk/python): API parity with sdk/rust + guest-agent (0.5.4b1)

Author: Leechael

sdk/python/README.md

@@ -8,6 +8,16 @@ Access TEE features from your Python application running inside dstack. Derive d
 pip install dstack-sdk
 ```
 
+Blockchain helpers are optional extras:
+
+| Extra | Pulls in | Use when |
+|---|---|---|
+| `dstack-sdk[ethereum]` | `eth-account` | You want `to_account` / `to_account_secure` for Ethereum signing |
+| `dstack-sdk[solana]` | `solders` | You want `to_keypair` / `to_keypair_secure` for Solana signing |
+| `dstack-sdk[all]` | both | You need both |
+
+Aliases `[eth]` and `[sol]` are accepted for convenience.
+
 ## Quick Start
 
 ```python
@@ -24,10 +34,11 @@ quote = client.get_quote(b'my-app-state')
 print(quote.quote)
 ```
 
-The client automatically connects to `/var/run/dstack.sock`. For local development with the simulator, pass the endpoint explicitly:
+The client automatically connects to `/var/run/dstack.sock`. For local development with the simulator:
 
 ```python
 client = DstackClient('http://localhost:8090')
+# or export DSTACK_SIMULATOR_ENDPOINT=http://localhost:8090
 ```
 
 ## Core API
@@ -45,18 +56,19 @@ btc_key = client.get_key('wallet/bitcoin')
 mainnet_key = client.get_key('wallet/eth/mainnet')
 testnet_key = client.get_key('wallet/eth/testnet')
 
-# Different algorithm
+# Use a different signature algorithm (requires dstack OS >= 0.5.7)
 ed_key = client.get_key('signing/key', algorithm='ed25519')
 ```
 
 **Parameters:**
-- `path`: Key derivation path (determines the key)
-- `purpose` (optional): Included in signature chain message, does not affect the derived key
-- `algorithm` (optional): `'secp256k1'` (default) or `'ed25519'`
+- `path` (optional): Key derivation path. Defaults to `""` (root).
+- `purpose` (optional): Included in the signature chain message; does not affect the derived key.
+- `algorithm` (optional): `'secp256k1'` (default) or `'ed25519'`.
 
 **Returns:** `GetKeyResponse`
 - `key`: Hex-encoded private key
 - `signature_chain`: Signatures proving the key was derived in a genuine TEE
+- `decode_key()` / `decode_signature_chain()`: Helpers that return `bytes`
 
 ### Generate Attestation Quotes
 
@@ -71,12 +83,23 @@ print(rtmrs)
 ```
 
 **Parameters:**
-- `report_data`: Exactly 64 bytes recommended. If shorter, pad with zeros. If longer, hash it first (e.g., SHA-256).
+- `report_data`: Up to 64 bytes (`bytes` or `str`). Shorter inputs are padded with zeros; longer inputs should be hashed first (e.g., SHA-256).
 
 **Returns:** `GetQuoteResponse`
 - `quote`: Hex-encoded TDX quote
 - `event_log`: JSON string of measured events
-- `replay_rtmrs()`: Method to compute RTMR values from event log
+- `replay_rtmrs()`: Method to compute RTMR values from the event log
+- `decode_quote()` / `decode_event_log()`: Helpers
+
+### Versioned Attestation
+
+`attest()` returns a versioned attestation payload that newer verifier APIs can dispatch on without sniffing the quote format.
+
+```python
+result = client.attest(b'user:alice:nonce123')
+print(result.attestation)        # hex string
+print(result.decode_attestation())  # bytes
+```
 
 ### Get Instance Info
 
@@ -85,15 +108,16 @@ info = client.info()
 print(info.app_id)
 print(info.instance_id)
 print(info.tcb_info)
+print(info.cloud_vendor, info.cloud_product)  # 0.5.7+
 ```
 
 **Returns:** `InfoResponse`
-- `app_id`: Application identifier
-- `instance_id`: Instance identifier
-- `app_name`: Application name
-- `tcb_info`: TCB measurements (MRTD, RTMRs, event log)
+- `app_id`, `instance_id`, `app_name`, `device_id`
+- `tcb_info`: TCB measurements (MRTD, RTMRs, event log, compose hash, ...)
 - `compose_hash`: Hash of the app configuration
 - `app_cert`: Application certificate (PEM)
+- `key_provider_info`: Key management configuration
+- `cloud_vendor` / `cloud_product`: Cloud provider strings (empty on older OS)
 
 ### Generate TLS Certificates
 
@@ -103,27 +127,35 @@ print(info.tcb_info)
 tls = client.get_tls_key(
     subject='api.example.com',
     alt_names=['localhost'],
-    usage_ra_tls=True  # Embed attestation in certificate
+    usage_ra_tls=True,    # Embed attestation in certificate
+    # 0.5.7+ options below:
+    not_before=1700000000,   # seconds since UNIX epoch
+    not_after=1800000000,
+    with_app_info=True,
 )
-
-print(tls.key)                # PEM private key
-print(tls.certificate_chain)  # Certificate chain
+print(tls.key)                  # PEM private key
+print(tls.certificate_chain)    # Certificate chain
 ```
 
 **Parameters:**
-- `subject` (optional): Certificate common name (e.g., domain name)
-- `alt_names` (optional): List of subject alternative names
-- `usage_ra_tls` (optional): Embed TDX quote in certificate extension
-- `usage_server_auth` (optional): Enable for server authentication (default: `True`)
-- `usage_client_auth` (optional): Enable for client authentication (default: `False`)
+- `subject` (optional): Certificate Common Name (e.g., domain name)
+- `alt_names` (optional): Subject Alternative Names
+- `usage_ra_tls` (optional): Embed TDX quote in a certificate extension (default `False`)
+- `usage_server_auth` (optional): Enable for server authentication (default `True`)
+- `usage_client_auth` (optional): Enable for client authentication (default `False`)
+- `not_before` / `not_after` (optional, kw-only): Validity window in seconds since UNIX epoch. Requires dstack OS >= 0.5.7.
+- `with_app_info` (optional, kw-only): Embed app identity into the certificate. Requires dstack OS >= 0.5.7.
+
+When any of the 0.5.7-only options is set, the SDK probes `Version` first and raises `RuntimeError` on older guest agents that lack it.
 
 **Returns:** `GetTlsKeyResponse`
 - `key`: PEM-encoded private key
 - `certificate_chain`: List of PEM certificates
+- `as_uint8array(max_length=None)`: Returns the DER-encoded private key bytes (handy when feeding key material into low-level crypto libraries)
 
 ### Sign and Verify
 
-Sign data using TEE-derived keys (not yet released):
+Sign data using TEE-derived keys:
 
 ```python
 result = client.sign('ed25519', b'message to sign')
@@ -137,21 +169,15 @@ print(valid.valid)  # True
 
 **`sign()` Parameters:**
 - `algorithm`: `'ed25519'`, `'secp256k1'`, or `'secp256k1_prehashed'`
-- `data`: Data to sign (bytes or string)
+- `data`: Data to sign (`bytes` or `str`). For `secp256k1_prehashed`, must be a 32-byte digest.
 
 **`sign()` Returns:** `SignResponse`
 - `signature`: Hex-encoded signature
 - `public_key`: Hex-encoded public key
 - `signature_chain`: Signatures proving TEE origin
 
-**`verify()` Parameters:**
-- `algorithm`: Algorithm used for signing
-- `data`: Original data
-- `signature`: Signature to verify
-- `public_key`: Public key to verify against
-
 **`verify()` Returns:** `VerifyResponse`
-- `valid`: Boolean indicating if signature is valid
+- `valid`: Boolean indicating if the signature is valid
 
 ### Emit Events
 
@@ -162,25 +188,28 @@ client.emit_event('config_loaded', 'production')
 client.emit_event('plugin_initialized', 'auth-v2')
 ```
 
-**Parameters:**
-- `event`: Event name (string identifier)
-- `payload`: Event value (bytes or string)
+### Diagnostics
+
+```python
+client.version()        # VersionResponse(version, rev) — raises on OS < 0.5.7
+client.is_reachable()   # Quick connectivity probe; never raises
+```
 
 ## Async Client
 
-For async applications, use `AsyncDstackClient`:
+For async applications, use `AsyncDstackClient`. The API surface is identical, but every method is a coroutine:
 
 ```python
-from dstack_sdk import AsyncDstackClient
 import asyncio
+from dstack_sdk import AsyncDstackClient
 
 async def main():
     client = AsyncDstackClient()
 
     info = await client.info()
     key = await client.get_key('wallet/eth')
 
-    # Concurrent operations
+    # Run requests concurrently
     keys = await asyncio.gather(
         client.get_key('user/alice'),
         client.get_key('user/bob'),
@@ -189,79 +218,91 @@ async def main():
 asyncio.run(main())
 ```
 
+`AsyncDstackClient` accepts the same constructor as `DstackClient` plus `use_sync_http: bool = False` for callers that need to issue sync HTTP from within an async context.
+
 ## Blockchain Integration
 
 ### Ethereum
 
 ```python
-from dstack_sdk.ethereum import to_account
+from dstack_sdk.ethereum import to_account_secure
 
 key = client.get_key('wallet/ethereum')
-account = to_account(key)
+account = to_account_secure(key)
 print(account.address)
 ```
 
+`to_account_secure(key)` hashes the full key material with SHA-256 before deriving the Ethereum private key. The legacy `to_account()` is kept for backward compatibility but uses raw key bytes—prefer the secure variant for new code.
+
 ### Solana
 
 ```python
-from dstack_sdk.solana import to_keypair
+from dstack_sdk.solana import to_keypair_secure
 
 key = client.get_key('wallet/solana')
-keypair = to_keypair(key)
-print(keypair.public_key)
-```
-
-## Development
-
-For local development without TDX hardware, use the simulator:
-
-```bash
-git clone https://github.com/Dstack-TEE/dstack.git
-cd dstack/sdk/simulator
-./build.sh
-./dstack-simulator
-```
-
-Then set the endpoint:
-
-```bash
-export DSTACK_SIMULATOR_ENDPOINT=http://localhost:8090
+keypair = to_keypair_secure(key)
+print(keypair.pubkey())
 ```
 
-Run tests with PDM:
-
-```bash
-pdm install -d
-pdm run pytest -s
-```
+Same pattern: `to_keypair_secure(key)` SHA-256-hashes the key material; `to_keypair()` is the legacy raw-bytes variant.
 
 ---
 
 ## Deployment Utilities
 
 These utilities are for deployment scripts, not runtime SDK operations.
 
-### Encrypt Environment Variables
+### Encrypted Environment Variables
 
-Encrypt secrets before deploying to dstack:
+The KMS returns a fresh X25519 public key (with a secp256k1 signature) that you encrypt secrets against before submitting them with your deployment. Always verify the signer before trusting the key:
 
 ```python
-from dstack_sdk import encrypt_env_vars, verify_env_encrypt_public_key, EnvVar
+from dstack_sdk import (
+    encrypt_env_vars,
+    verify_env_encrypt_public_key,
+    verify_env_encrypt_public_key_legacy,
+    EnvVar,
+)
 
-# Get and verify the KMS public key
-# (obtain public_key and signature from KMS API)
-kms_identity = verify_env_encrypt_public_key(public_key_bytes, signature_bytes, app_id)
-if not kms_identity:
-    raise RuntimeError('Invalid KMS key')
+# `public_key`, `signature_v1`, `timestamp` come from KMS /GetAppEnvEncryptPubKey.
+signer = verify_env_encrypt_public_key(
+    public_key=public_key_bytes,
+    signature=signature_v1_bytes,
+    app_id=app_id_hex,
+    timestamp=timestamp,
+)
+if signer is None:
+    # Fallback for older KMS builds that only emit the unprotected legacy
+    # signature. Vulnerable to replay; warn loudly if you must use it.
+    signer = verify_env_encrypt_public_key_legacy(
+        public_key=public_key_bytes,
+        signature=legacy_signature_bytes,
+        app_id=app_id_hex,
+    )
+    if signer is None:
+        raise RuntimeError('invalid KMS env-encrypt public key')
+
+# Always compare the recovered signer against a known-good KMS signer
+# address, obtained out-of-band from the DstackKms contract or your
+# deployment configuration. Without this check, an attacker could sign
+# their own env-encrypt key and the verification above would still pass.
+EXPECTED_KMS_SIGNER = '0x...'  # replace with your known KMS signer address
+if signer != EXPECTED_KMS_SIGNER:
+    raise RuntimeError(
+        f'unexpected KMS signer: got {signer}, '
+        f'expected {EXPECTED_KMS_SIGNER}'
+    )
 
-# Encrypt variables
 env_vars = [
     EnvVar(key='DATABASE_URL', value='postgresql://...'),
     EnvVar(key='API_KEY', value='secret'),
 ]
-encrypted = encrypt_env_vars(env_vars, public_key)
+encrypted = await encrypt_env_vars(env_vars, public_key_hex)
+# encrypt_env_vars_sync(...) is also available for non-async callers.
 ```
 
+`verify_env_encrypt_public_key` returns the recovered compressed secp256k1 signer (`0x`-prefixed hex) on success, or `None` for any failure (bad length, expired/future timestamp, malformed `app_id`, invalid signature). The default `max_age_seconds` is 300; pass a larger value if your deployment workflow legitimately holds the response longer.
+
 ### Calculate Compose Hash
 
 ```python
@@ -272,6 +313,43 @@ hash_value = get_compose_hash(app_compose_dict)
 
 ---
 
+## Compatibility
+
+| Feature | Required dstack OS |
+|---|---|
+| `get_key`, `get_quote`, `get_tls_key` (legacy fields), `info` (legacy fields) | 0.3+ |
+| `emit_event` | 0.5.0+ |
+| `attest`, `sign` / `verify`, `is_reachable` | 0.5.0+ (sign/verify require server build with the feature) |
+| `version`, `algorithm='ed25519'` on `get_key`, `info.cloud_vendor` / `cloud_product`, `not_before` / `not_after` / `with_app_info` on `get_tls_key` | 0.5.7+ |
+| `verify_env_encrypt_public_key` (signature_v1 with timestamp) | Requires KMS build that emits `signature_v1`; legacy variant remains available |
+
+Calls that require 0.5.7-only fields probe the `Version` RPC first and raise a clear `RuntimeError` on older guest agents.
+
+## Development
+
+For local development without TDX hardware, use the simulator:
+
+```bash
+git clone https://github.com/Dstack-TEE/dstack.git
+cd dstack/sdk/simulator
+./build.sh
+./dstack-simulator
+```
+
+Then set the endpoint:
+
+```bash
+export DSTACK_SIMULATOR_ENDPOINT=http://localhost:8090
+```
+
+Install dev dependencies and run tests with PDM:
+
+```bash
+cd sdk/python
+make install
+make test
+```
+
 ## Migration from TappdClient
 
 Replace `TappdClient` with `DstackClient`:
@@ -288,7 +366,7 @@ client = DstackClient()
 
 Method changes:
 - `derive_key()` → `get_tls_key()` for TLS certificates
-- `tdx_quote()` → `get_quote()`
+- `tdx_quote()` → `get_quote()` (raw data only, no hash algorithms)
 - Socket path: `/var/run/tappd.sock` → `/var/run/dstack.sock`
 
 ## License