Add sandbox token derivation helper

 * Add a sandbox auth helper that derives per-run HMAC tokens from the internal platform token.

Author: md-mq

cli/polyaxon/_sandbox/__init__.py

@@ -0,0 +1,12 @@
+from polyaxon._sandbox.auth import (
+    DERIVATION_PREFIX,
+    derive_sandbox_token,
+    derive_sandbox_token_from_env,
+)
+
+
+__all__ = [
+    "DERIVATION_PREFIX",
+    "derive_sandbox_token",
+    "derive_sandbox_token_from_env",
+]

cli/polyaxon/_sandbox/auth.py

@@ -0,0 +1,32 @@
+import hashlib
+import hmac
+import os
+
+from polyaxon._env_vars.keys import ENV_KEYS_SECRET_INTERNAL_TOKEN
+from polyaxon.exceptions import PolyaxonConverterError
+
+
+DERIVATION_PREFIX = b"sandbox:v1:"
+
+
+def derive_sandbox_token(signing_key: str, run_uuid: str) -> str:
+    if not signing_key or not signing_key.strip():
+        raise ValueError("A non-empty signing key is required.")
+    if not run_uuid or not run_uuid.strip():
+        raise ValueError("A non-empty run uuid is required.")
+
+    return hmac.new(
+        signing_key.encode(),
+        DERIVATION_PREFIX + run_uuid.encode(),
+        hashlib.sha256,
+    ).hexdigest()
+
+
+def derive_sandbox_token_from_env(run_uuid: str) -> str:
+    signing_key = os.environ.get(ENV_KEYS_SECRET_INTERNAL_TOKEN)
+    if not signing_key or not signing_key.strip():
+        raise PolyaxonConverterError(
+            "plugins.sandbox is enabled but "
+            f"{ENV_KEYS_SECRET_INTERNAL_TOKEN} is not set."
+        )
+    return derive_sandbox_token(signing_key, run_uuid)

cli/tests/test_sandbox/__init__.py

@@ -0,0 +1 @@
+

cli/tests/test_sandbox/test_auth.py

@@ -0,0 +1,70 @@
+import hashlib
+import hmac
+import os
+import uuid
+
+from unittest import TestCase
+
+import pytest
+
+from polyaxon._env_vars.keys import ENV_KEYS_SECRET_INTERNAL_TOKEN
+from polyaxon._sandbox.auth import (
+    DERIVATION_PREFIX,
+    derive_sandbox_token,
+    derive_sandbox_token_from_env,
+)
+from polyaxon.exceptions import PolyaxonConverterError
+
+
+@pytest.mark.utils_mark
+class TestSandboxAuth(TestCase):
+    def setUp(self):
+        super().setUp()
+        self.current_internal_token = os.environ.get(ENV_KEYS_SECRET_INTERNAL_TOKEN)
+
+    def tearDown(self):
+        if self.current_internal_token is None:
+            os.environ.pop(ENV_KEYS_SECRET_INTERNAL_TOKEN, None)
+        else:
+            os.environ[ENV_KEYS_SECRET_INTERNAL_TOKEN] = self.current_internal_token
+        super().tearDown()
+
+    def test_derive_sandbox_token(self):
+        signing_key = "internal-token"
+        run_uuid = uuid.uuid4().hex
+        expected = hmac.new(
+            signing_key.encode(),
+            DERIVATION_PREFIX + run_uuid.encode(),
+            hashlib.sha256,
+        ).hexdigest()
+
+        assert derive_sandbox_token(signing_key, run_uuid) == expected
+
+    def test_derive_sandbox_token_scopes_by_key_and_run_uuid(self):
+        run_uuid = uuid.uuid4().hex
+        token = derive_sandbox_token("internal-token", run_uuid)
+
+        assert token != derive_sandbox_token("other-token", run_uuid)
+        assert token != derive_sandbox_token("internal-token", uuid.uuid4().hex)
+
+    def test_derive_sandbox_token_rejects_empty_inputs(self):
+        with self.assertRaises(ValueError):
+            derive_sandbox_token("", uuid.uuid4().hex)
+
+        with self.assertRaises(ValueError):
+            derive_sandbox_token("internal-token", "")
+
+    def test_derive_sandbox_token_from_env(self):
+        os.environ[ENV_KEYS_SECRET_INTERNAL_TOKEN] = "internal-token"
+        run_uuid = uuid.uuid4().hex
+
+        assert derive_sandbox_token_from_env(run_uuid) == derive_sandbox_token(
+            "internal-token", run_uuid
+        )
+
+    def test_derive_sandbox_token_from_env_requires_internal_token(self):
+        os.environ.pop(ENV_KEYS_SECRET_INTERNAL_TOKEN, None)
+
+        with self.assertRaises(PolyaxonConverterError) as ctx:
+            derive_sandbox_token_from_env(uuid.uuid4().hex)
+        assert ENV_KEYS_SECRET_INTERNAL_TOKEN in str(ctx.exception)