bearer_token.py

import ast
import base64
import hashlib
import hmac
import os
from datetime import datetime, timezone

class BearerTokenGenerator:
    @staticmethod
    def get_bearer(body: str, path: str = "/chats/stream") -> tuple:
        # 从环境变量获取 SECRET_AUTH_PREFIX 和 SECRET_KEY
        secret_auth_prefix_str = os.environ.get('SECRET_AUTH_PREFIX')
        secret_key_str = os.environ.get('SECRET_KEY')

        if not secret_auth_prefix_str or not secret_key_str:
            raise ValueError("SECRET_AUTH_PREFIX or SECRET_KEY environment variables are not set")

        # 将字符串转换为字节列表
        secret_auth_prefix = bytes(ast.literal_eval(secret_auth_prefix_str))
        secret_key = bytes(ast.literal_eval(secret_key_str))

        timestamp = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
        prefix = f"POST:{path}:{timestamp}\n".encode()
        to_sign = prefix + body.encode()
        
        signature = BearerTokenGenerator.generate_signature(to_sign, secret_key)
        auth_prefix_base64 = base64.b64encode(secret_auth_prefix).decode()
        
        bearer_token = f"Bearer {auth_prefix_base64}.{signature}"
        return bearer_token, timestamp

    @staticmethod
    def generate_signature(to_sign: bytes, secret_key: bytes) -> str:
        h = hmac.new(secret_key, to_sign, hashlib.sha256)
        return base64.b64encode(h.digest()).decode()