license_bootstrap.py

"""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()