style(sealed-env): satisfy Python checks across CLI and SDK

Signed-off-by: Kevin Wang <wy721@qq.com>

Author: kvinwang

dstack/vmm/src/tests/test_vmm_cli.py

@@ -64,7 +64,10 @@ def test_kms_urls_are_sent_to_upgrade_app(self) -> None:
 
 
 class SealedEnvTests(unittest.TestCase):
+    """Exercise signed recipient selection and SealedEnv compatibility."""
+
     def setUp(self):
+        """Create synthetic keys, a signed response and an isolated client."""
         self.app_id = "11" * 20
         self.epoch = 3
         self.timestamp = 1700000000
@@ -102,9 +105,12 @@ def setUp(self):
         self.envs = {"SECRET": "synthetic test value", "UNICODE": "测试"}
 
     def decrypt(self, payload):
+        """Check plaintext recovery and rejection under a different epoch context."""
         self.assertNotIn("encrypted_env", payload)
         envelope = json.loads(bytes.fromhex(payload["sealed_env"]))
-        self.assertEqual(set(envelope), {"version", "root_epoch", "app_epoch", "enc", "ciphertext"})
+        self.assertEqual(
+            set(envelope), {"version", "root_epoch", "app_epoch", "enc", "ciphertext"}
+        )
         self.assertEqual(envelope["version"], 1)
         self.assertEqual(envelope["root_epoch"], self.epoch)
         enc = base64.b64decode(envelope["enc"], validate=True)
@@ -125,19 +131,25 @@ def decrypt(self, payload):
             wrong_epoch.open(ciphertext, aad=b"")
 
     def test_sealed_payload_round_trip_and_normalized_pins(self):
+        """Accept equivalent pin encodings without using legacy encryption."""
         for pin in [self.pin, "0x" + self.pin.upper()]:
             with (
                 self.subTest(pin=pin),
                 patch.object(VMM_CLI, "load_whitelist", return_value=[pin]),
             ):
-                self.decrypt(self.cli.encrypt_env_payload(self.envs, self.app_id, None, True))
+                self.decrypt(
+                    self.cli.encrypt_env_payload(self.envs, self.app_id, None, True)
+                )
         self.cli.get_app_env_encrypt_pub_key.assert_not_called()
 
     def test_app_epoch_is_signed_and_used_for_sealing(self):
+        """Bind the app epoch in both the signature and HPKE context."""
         self.response["app_epoch"] = 2
         digest = keccak(
-            b"dstack-sealed-env-pubkey-v1:" + bytes.fromhex(self.app_id)
-            + self.epoch.to_bytes(4, "big") + (2).to_bytes(4, "big")
+            b"dstack-sealed-env-pubkey-v1:"
+            + bytes.fromhex(self.app_id)
+            + self.epoch.to_bytes(4, "big")
+            + (2).to_bytes(4, "big")
             + bytes.fromhex(self.response["public_key"])
             + self.timestamp.to_bytes(8, "big")
         )
@@ -147,16 +159,20 @@ def test_app_epoch_is_signed_and_used_for_sealing(self):
         self.assertEqual(envelope["version"], 1)
         self.assertEqual(envelope["app_epoch"], 2)
         receiver = self.suite.create_recipient_context(
-            base64.b64decode(envelope["enc"]), self.recipient.private_key,
+            base64.b64decode(envelope["enc"]),
+            self.recipient.private_key,
             info=f"dstack-sealed-env-v1:{self.epoch}:2".encode(),
         )
         plaintext = json.loads(receiver.open(base64.b64decode(envelope["ciphertext"])))
-        self.assertEqual(plaintext, {"env": [{"key": k, "value": v} for k, v in self.envs.items()]})
+        self.assertEqual(
+            plaintext, {"env": [{"key": k, "value": v} for k, v in self.envs.items()]}
+        )
         self.response["app_epoch"] = 1
         with self.assertRaises(ValueError):
             self.cli.encrypt_env_payload(self.envs, self.app_id, None, True)
 
     def test_unpinned_signer_never_becomes_trusted(self):
+        """Reject recipients without an independently trusted signer."""
         for pins in [
             [],
             [keys.PrivateKey(b"\x33" * 32).public_key.to_compressed_bytes().hex()],
@@ -172,6 +188,7 @@ def test_unpinned_signer_never_becomes_trusted(self):
         self.cli.get_app_env_encrypt_pub_key.assert_not_called()
 
     def test_signature_binds_every_field_and_app(self):
+        """Reject mutations of the signed claim or application identity."""
         cases = [
             ("public_key", "44" * 32),
             ("root_epoch", 4),
@@ -201,20 +218,26 @@ def test_signature_binds_every_field_and_app(self):
         self.cli.get_app_env_encrypt_pub_key.assert_not_called()
 
     def test_pinned_signed_key_timestamp_window(self):
+        """Enforce the freshness window for a trusted signer."""
         for age in [-60, 0, 300]:
-            with self.subTest(age=age), patch.object(
-                VMM_CLI.time, "time", return_value=self.timestamp + age
+            with (
+                self.subTest(age=age),
+                patch.object(VMM_CLI.time, "time", return_value=self.timestamp + age),
             ):
-                self.decrypt(self.cli.encrypt_env_payload(self.envs, self.app_id, None, True))
+                self.decrypt(
+                    self.cli.encrypt_env_payload(self.envs, self.app_id, None, True)
+                )
         for age, reason in [(-61, "future"), (301, "stale")]:
-            with self.subTest(age=age), patch.object(
-                VMM_CLI.time, "time", return_value=self.timestamp + age
+            with (
+                self.subTest(age=age),
+                patch.object(VMM_CLI.time, "time", return_value=self.timestamp + age),
             ):
                 with self.assertRaisesRegex(ValueError, reason):
                     self.cli.encrypt_env_payload(self.envs, self.app_id, None, True)
         self.cli.get_app_env_encrypt_pub_key.assert_not_called()
 
     def test_public_key_request_contains_only_app_id(self):
+        """Keep epoch selection on the KMS rather than in the lookup request."""
         response = self.cli.get_sealed_env_pub_key(self.app_id)
         self.assertEqual(
             self.cli.client.request.call_args.kwargs["body"],
@@ -223,8 +246,10 @@ def test_public_key_request_contains_only_app_id(self):
         self.assertEqual(response["root_epoch"], self.epoch)
 
     def test_missing_rpc_cannot_replay_a_legacy_recipient(self):
+        """Reject capability downgrade attempts using a captured legacy key."""
         legacy_digest = keccak(
-            b"dstack-env-encrypt-pubkey:" + bytes.fromhex(self.app_id)
+            b"dstack-env-encrypt-pubkey:"
+            + bytes.fromhex(self.app_id)
             + bytes.fromhex(self.response["public_key"])
         )
         legacy = {
@@ -234,27 +259,48 @@ def test_missing_rpc_cannot_replay_a_legacy_recipient(self):
         missing = {"error": "Service not found: GetSealedEnvPubKey"}
         errors = [(status, missing) for status in [400, 404]]
         errors += [
-            (400, {"error": f"Request failed with status={status}, error={json.dumps(missing)}"})
+            (
+                400,
+                {
+                    "error": f"Request failed with status={status}, error={json.dumps(missing)}"
+                },
+            )
             for status in ["400 Bad Request", "404 Not Found"]
         ]
         for error in errors:
-            with self.subTest(error=error), patch.object(VMM_CLI, "VmmClient") as client:
+            with (
+                self.subTest(error=error),
+                patch.object(VMM_CLI, "VmmClient") as client,
+            ):
                 client.return_value.request.side_effect = [error, (200, legacy)]
                 with self.assertRaises(RuntimeError):
-                    self.cli.encrypt_env_payload(self.envs, self.app_id, "https://kms.example", True)
+                    self.cli.encrypt_env_payload(
+                        self.envs, self.app_id, "https://kms.example", True
+                    )
                 client.return_value.request.assert_called_once()
         self.cli.get_app_env_encrypt_pub_key.assert_not_called()
 
     def test_errors_and_bad_responses_do_not_downgrade(self):
+        """Fail closed on lookup errors and malformed key responses."""
         for status, body in [
             (401, {"error": "Service not found: GetSealedEnvPubKey"}),
             (403, {"error": "forbidden"}),
             (404, "Not Found"),
             (400, {"error": "invalid app"}),
             (500, {"error": "unavailable"}),
-            (400, {"error": 'Request failed with status=500 Internal Server Error, error={"error":"Service not found: GetSealedEnvPubKey"}'}),
-            (400, {"error": 'Request failed with status=404 Not Found, error={"error":"other service"}'}),
-            (400, {"error": 'Request failed with status=404 Not Found, error=invalid'}),
+            (
+                400,
+                {
+                    "error": 'Request failed with status=500 Internal Server Error, error={"error":"Service not found: GetSealedEnvPubKey"}'
+                },
+            ),
+            (
+                400,
+                {
+                    "error": 'Request failed with status=404 Not Found, error={"error":"other service"}'
+                },
+            ),
+            (400, {"error": "Request failed with status=404 Not Found, error=invalid"}),
             (200, {}),
             (200, "invalid JSON"),
         ]:
@@ -268,6 +314,7 @@ def test_errors_and_bad_responses_do_not_downgrade(self):
         self.cli.get_app_env_encrypt_pub_key.assert_not_called()
 
     def test_direct_kms_verifies_pin_and_does_not_forward_vmm_credentials(self):
+        """Authenticate direct KMS keys without forwarding VMM credentials."""
         with patch.object(VMM_CLI, "VmmClient") as client:
             client.return_value.request.return_value = (200, self.response)
             self.decrypt(
@@ -279,17 +326,34 @@ def test_direct_kms_verifies_pin_and_does_not_forward_vmm_credentials(self):
         self.cli.get_app_env_encrypt_pub_key.assert_not_called()
 
     def test_image_downgrade_and_env_replacement_share_one_request(self):
+        """Replace the image and environment format in the same update."""
+
         def rpc(method, params=None):
             if method == "GetInfo":
-                return {"found": True, "info": {"app_id": self.app_id, "configuration": {
-                    "image": "new", "compose_file": '{"kms_enabled":true}',
-                }}}
+                return {
+                    "found": True,
+                    "info": {
+                        "app_id": self.app_id,
+                        "configuration": {
+                            "image": "new",
+                            "compose_file": '{"kms_enabled":true}',
+                        },
+                    },
+                }
             if method == "ListImages":
-                return {"images": [{"name": "new", "sealed_env": True}, {"name": "old"}]}
+                return {
+                    "images": [{"name": "new", "sealed_env": True}, {"name": "old"}]
+                }
             return {}
+
         self.cli.rpc_call = Mock(side_effect=rpc)
-        with patch.object(VMM_CLI, "parse_env_file", return_value=self.envs), contextlib.redirect_stdout(io.StringIO()):
-            self.cli.update_vm("vm", image="old", env_file="env", kms_urls=["https://kms.example"])
+        with (
+            patch.object(VMM_CLI, "parse_env_file", return_value=self.envs),
+            contextlib.redirect_stdout(io.StringIO()),
+        ):
+            self.cli.update_vm(
+                "vm", image="old", env_file="env", kms_urls=["https://kms.example"]
+            )
         calls = self.cli.rpc_call.call_args_list
         self.assertNotIn("ResizeVm", [call.args[0] for call in calls])
         method, payload = calls[-1].args
@@ -300,6 +364,7 @@ def rpc(method, params=None):
         self.cli.client.request.assert_not_called()
 
     def test_create_update_and_update_env_send_only_the_negotiated_field(self):
+        """Use the selected ciphertext field consistently across CLI commands."""
         compose = json.dumps({"kms_enabled": True})
         args = argparse.Namespace(
             compose="compose.json",

dstack/vmm/src/vmm-cli.py

@@ -324,9 +324,7 @@ def seal_env(envs, hex_public_key: str, root_epoch: int, app_epoch: int = 0) ->
     public_key = suite.kem.deserialize_public_key(
         bytes.fromhex(hex_public_key.removeprefix("0x"))
     )
-    enc, context = suite.create_sender_context(
-        public_key, info=info.encode("utf-8")
-    )
+    enc, context = suite.create_sender_context(public_key, info=info.encode("utf-8"))
     plaintext = json.dumps({"env": envs}).encode("utf-8")
     envelope = {
         "version": 1,
@@ -341,6 +339,7 @@ def seal_env(envs, hex_public_key: str, root_epoch: int, app_epoch: int = 0) ->
 
 
 def validate_uint(value: int, bits: int, name: str) -> None:
+    """Require an integer in the specified unsigned range."""
     if type(value) is not int or not 0 <= value < 1 << bits:
         raise ValueError(f"{name} must be a uint{bits}")
 
@@ -849,10 +848,14 @@ def list_images(self, json_output: bool = False) -> None:
         print(format_table(rows, headers))
 
     def image_supports_sealed_env(self, image: Optional[str]) -> bool:
+        """Read the selected image capability, treating missing metadata as legacy."""
         if not image:
             return False
         images = self.rpc_call("ListImages").get("images", [])
-        return any(item.get("name") == image and item.get("sealed_env") is True for item in images)
+        return any(
+            item.get("name") == image and item.get("sealed_env") is True
+            for item in images
+        )
 
     def get_sealed_env_pub_key(
         self,
@@ -877,15 +880,23 @@ def get_sealed_env_pub_key(
         return response
 
     def encrypt_env_payload(
-        self, envs: Dict[str, str], app_id: str, kms_url: Optional[str],
+        self,
+        envs: Dict[str, str],
+        app_id: str,
+        kms_url: Optional[str],
         supports_sealed_env: bool,
     ) -> Dict:
         """Select encryption using the target image capability and KMS RPC support."""
         envs_list = [{"key": k, "value": v} for k, v in envs.items()]
         if supports_sealed_env:
             key = self.get_sealed_env_pub_key(app_id, kms_url)
             return {
-                "sealed_env": seal_env(envs_list, key["public_key"], key["root_epoch"], key.get("app_epoch", 0))
+                "sealed_env": seal_env(
+                    envs_list,
+                    key["public_key"],
+                    key["root_epoch"],
+                    key.get("app_epoch", 0),
+                )
             }
         public_key = self.get_app_env_encrypt_pub_key(app_id, kms_url)
         return {"encrypted_env": encrypt_env(envs_list, public_key)}
@@ -1153,7 +1164,14 @@ def create_vm(self, args) -> None:
             encrypt_url = args.kms_encrypt_url or (
                 args.kms_url[0] if args.kms_url else None
             )
-            params.update(self.encrypt_env_payload(envs, app_id, encrypt_url, self.image_supports_sealed_env(args.image)))
+            params.update(
+                self.encrypt_env_payload(
+                    envs,
+                    app_id,
+                    encrypt_url,
+                    self.image_supports_sealed_env(args.image),
+                )
+            )
         response = self.rpc_call("CreateVm", params)
         print(f"Created VM with ID: {response.get('id')}")
         return response.get("id")
@@ -1179,9 +1197,14 @@ def update_vm_env(
         compose_file = vm_configuration.get("compose_file")
 
         payload = {"id": vm_id}
-        payload.update(self.encrypt_env_payload(
-            envs, app_id, kms_urls[0], self.image_supports_sealed_env(vm_configuration.get("image"))
-        ))
+        payload.update(
+            self.encrypt_env_payload(
+                envs,
+                app_id,
+                kms_urls[0],
+                self.image_supports_sealed_env(vm_configuration.get("image")),
+            )
+        )
 
         if compose_file:
             try:
@@ -1345,10 +1368,18 @@ def update_vm(
                 vm_configuration = vm_info_response["info"].get("configuration") or {}
                 compose_file_content = vm_configuration.get("compose_file")
 
-                upgrade_params.update(self.encrypt_env_payload(
-                    envs, app_id, kms_urls[0],
-                    self.image_supports_sealed_env(image if image is not None else vm_configuration.get("image")),
-                ))
+                upgrade_params.update(
+                    self.encrypt_env_payload(
+                        envs,
+                        app_id,
+                        kms_urls[0],
+                        self.image_supports_sealed_env(
+                            image
+                            if image is not None
+                            else vm_configuration.get("image")
+                        ),
+                    )
+                )
                 updates.append("environment variables")
 
                 # update allowed_envs in compose file if needed

sdk/python/src/dstack_sdk/__init__.py

@@ -27,17 +27,17 @@
 from .dstack_client_v1 import InfoResponseV1
 from .dstack_client_v1 import IssueCertResponseV1
 from .dstack_client_v1 import VersionResponseV1
-from .seal_env import seal_env
 from .encrypt_env_vars import EnvVar
 from .encrypt_env_vars import encrypt_env_vars
 from .encrypt_env_vars import encrypt_env_vars_sync
 from .get_compose_hash import AppCompose
 from .get_compose_hash import DockerConfig
 from .get_compose_hash import Requirements
 from .get_compose_hash import get_compose_hash
+from .seal_env import seal_env
 from .verify_env_encrypt_public_key import verify_env_encrypt_public_key
-from .verify_env_encrypt_public_key import verify_sealed_env_public_key
 from .verify_env_encrypt_public_key import verify_env_encrypt_public_key_legacy
+from .verify_env_encrypt_public_key import verify_sealed_env_public_key
 
 __all__ = [
     "seal_env",