"""License bootstrap for the polygen demo. The polygen SDK validates a license at import time. Two modes: - **Deploy (HuggingFace Space):** set the ``POLYGEN_LICENSE`` secret to a real JWT issued for the deployed wheel. Leave ``POLYGEN_DEMO_LOCAL_MOCK`` unset; ``ensure_license`` is then a no-op and the SDK validates normally. - **Local dev:** set ``POLYGEN_DEMO_LOCAL_MOCK=1``. This generates a self-consistent RSA keypair, signs a short-lived JWT with it, and patches the SDK's public-key lookup to the matching key -- the same pattern ``tests/conftest.py`` uses for ad-hoc scripts. No real license needed. ``ensure_license`` MUST run before ``import polygen``. """ import os def apply_local_mock() -> None: """Generate a self-consistent license and patch the SDK key lookup.""" import datetime from pathlib import Path from unittest import mock import jwt from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import rsa private_key = rsa.generate_private_key(public_exponent=65537, key_size=4096) private_pem = private_key.private_bytes( encoding=serialization.Encoding.PEM, format=serialization.PrivateFormat.PKCS8, encryption_algorithm=serialization.NoEncryption(), ) public_pem = private_key.public_key().public_bytes( encoding=serialization.Encoding.PEM, format=serialization.PublicFormat.SubjectPublicKeyInfo, ) now = datetime.datetime.now(datetime.timezone.utc) os.environ["POLYGEN_LICENSE"] = jwt.encode( { "iss": "https://license.datasent.com", "sub": "polygen demo (local mock)", "aud": "polygen-sdk", "email": "demo@datasent.com", "iat": now, "exp": now + datetime.timedelta(hours=12), }, key=private_pem, algorithm="PS512", ) mock_path = mock.Mock(spec=Path) mock_path.joinpath.return_value = mock_path mock_path.__truediv__ = mock.Mock(return_value=mock_path) mock_path.read_text.return_value = public_pem # Scope the patch to polygen's license lookup ONLY. A global patch of # importlib.resources.files (as in tests/conftest.py) breaks any other # library that locates package resources -- e.g. gradio/starlette loading # Jinja templates. Delegate every other anchor to the real ``files``. import importlib.resources as _ir real_files = _ir.files def fake_files(*args, **kwargs): if args and args[0] == "polygen._license": return mock_path return real_files(*args, **kwargs) mock.patch("importlib.resources.files", side_effect=fake_files).start() def ensure_license() -> None: """Apply the local mock when requested; otherwise rely on POLYGEN_LICENSE.""" if os.environ.get("POLYGEN_DEMO_LOCAL_MOCK") == "1": apply_local_mock()