Arag / app /core /subscription_crypto.py
AuthorBot
Restructure project for HF Spaces deployment
772f852
Raw
History Blame Contribute Delete
2.94 kB
"""Author RAG Chatbot SaaS — Subscription Token Cryptography.
Generates and verifies HMAC-SHA256 signed subscription tokens.
Token format: base64url(payload_json).base64url(hmac_signature)
"""
import base64
import hashlib
import hmac
import json
from datetime import datetime, timezone
def generate_subscription_token(
author_id: str,
plan: str,
expires_at: str,
secret: str,
) -> str:
"""Generate a signed subscription token.
Args:
author_id: UUID of the author.
plan: Plan name ('monthly', 'quarterly', etc.).
expires_at: ISO 8601 expiry datetime string.
secret: HMAC secret key.
Returns:
Signed token string (payload.signature, both base64url-encoded).
"""
payload = json.dumps({
"author_id": author_id,
"plan": plan,
"expires_at": expires_at,
}, separators=(",", ":"))
encoded_payload = base64.urlsafe_b64encode(payload.encode()).decode().rstrip("=")
signature = _sign(encoded_payload, secret)
return f"{encoded_payload}.{signature}"
def verify_subscription_token(token: str, secret: str) -> dict:
"""Verify and decode a subscription token.
Args:
token: Token string from generate_subscription_token.
secret: HMAC secret key.
Returns:
Decoded payload dict with author_id, plan, expires_at.
Raises:
ValueError: If signature is invalid or token is malformed.
ValueError: If token has expired (message contains 'expired').
"""
try:
parts = token.split(".")
if len(parts) != 2:
raise ValueError("Malformed token: expected 2 parts")
encoded_payload, provided_sig = parts
# Verify signature
expected_sig = _sign(encoded_payload, secret)
if not hmac.compare_digest(expected_sig, provided_sig):
raise ValueError("Token signature is invalid")
# Decode payload
padded = encoded_payload + "=" * (4 - len(encoded_payload) % 4)
payload = json.loads(base64.urlsafe_b64decode(padded).decode())
# Check expiry
expires_at = datetime.fromisoformat(payload["expires_at"])
if expires_at.tzinfo is None:
expires_at = expires_at.replace(tzinfo=timezone.utc)
if expires_at < datetime.now(timezone.utc):
raise ValueError(f"Token has expired at {payload['expires_at']}")
return payload
except ValueError:
raise
except Exception as e:
raise ValueError(f"Token verification failed: {e}") from e
def _sign(payload: str, secret: str) -> str:
"""Compute HMAC-SHA256 signature of payload.
Args:
payload: String to sign.
secret: HMAC secret key.
Returns:
Base64url-encoded signature (no padding).
"""
sig = hmac.new(secret.encode(), payload.encode(), hashlib.sha256).digest()
return base64.urlsafe_b64encode(sig).decode().rstrip("=")