uploader-bot/app/core/access/content_access_manager.py

from __future__ import annotations

import base64
import json
import logging
import os
import secrets
import time
from dataclasses import dataclass
from datetime import datetime, timedelta
from typing import Any, Dict, Optional, Tuple, Callable

from app.core._blockchain.ton.nft_license_manager import NFTLicenseManager
from app.core.crypto.content_cipher import ContentCipher

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class StreamingToken:
    token: str
    content_id: str
    owner_address: str
    issued_at: float
    expires_at: float

    def is_valid(self, now: Optional[float] = None) -> bool:
        now = now or time.time()
        return now < self.expires_at


class ContentAccessManager:
    """
    Управление доступом к зашифрованному контенту по NFT лицензиям в TON.

    Обязанности:
    - grant_access(): принять tonProof + content_id, проверить лицензию, выдать временный токен
    - verify_access(): валидация токена при запросе стрима/скачивания
    - create_streaming_token(): генерация подписанного/непредсказуемого токена с TTL
    - stream/decrypt: интеграция с ContentCipher — расшифровка возможна только при валидной лицензии/токене
    """

    DEFAULT_TOKEN_TTL_SEC = int(os.getenv("STREAM_TOKEN_TTL_SEC", "600"))  # 10 минут по умолчанию

    def __init__(
        self,
        nft_manager: Optional[NFTLicenseManager] = None,
        cipher: Optional[ContentCipher] = None,
    ):
        self.nft_manager = nft_manager or NFTLicenseManager()
        self.cipher = cipher or ContentCipher()
        # Простой in-memory storage токенов. Для продакшена стоит заменить на Redis или БД.
        self._tokens: Dict[str, StreamingToken] = {}
        logger.debug("ContentAccessManager initialized; token_ttl=%s", self.DEFAULT_TOKEN_TTL_SEC)

    def create_streaming_token(self, content_id: str, owner_address: str, ttl_sec: Optional[int] = None) -> StreamingToken:
        ttl = ttl_sec or self.DEFAULT_TOKEN_TTL_SEC
        token = base64.urlsafe_b64encode(secrets.token_bytes(32)).decode("ascii").rstrip("=")
        now = time.time()
        st = StreamingToken(
            token=token,
            content_id=content_id,
            owner_address=owner_address,
            issued_at=now,
            expires_at=now + ttl,
        )
        self._tokens[token] = st
        logger.info("Streaming token issued content_id=%s owner=%s ttl=%s", content_id, owner_address, ttl)
        return st

    def verify_access(self, token: str, content_id: str) -> Tuple[bool, Optional[str], Optional[StreamingToken]]:
        if not token:
            return False, "Missing token", None
        st = self._tokens.get(token)
        if not st:
            return False, "Token not found", None
        if not st.is_valid():
            # Удаляем просроченный
            self._tokens.pop(token, None)
            return False, "Token expired", None
        if st.content_id != content_id:
            return False, "Token/content mismatch", None
        logger.debug("Streaming token verified for content_id=%s owner=%s", st.content_id, st.owner_address)
        return True, None, st

    async def grant_access(
        self,
        ton_proof: Dict[str, Any],
        content_id: str,
        nft_address: Optional[str] = None,
        token_ttl_sec: Optional[int] = None,
    ) -> Tuple[bool, Optional[str], Optional[Dict[str, Any]]]:
        """
        Композитный сценарий: валидируем tonProof, проверяем владение NFT лицензией,
        создаем временный токен для стриминга.
        Возвращает: (ok, error, payload)
        payload: { token, expires_at, owner_address, nft_item }
        """
        try:
            ok, err, nft_item = await self.nft_manager.check_license_validity(
                ton_proof=ton_proof,
                content_id=content_id,
                nft_address=nft_address,
            )
            if not ok:
                return False, err, None

            owner_address = nft_proof_owner(ton_proof)
            token = self.create_streaming_token(content_id, owner_address, token_ttl_sec)
            payload = {
                "token": token.token,
                "expires_at": token.expires_at,
                "owner_address": token.owner_address,
                "nft_item": nft_item,
            }
            return True, None, payload
        except Exception as e:
            logger.exception("grant_access failed")
            return False, str(e), None

    def decrypt_for_stream(
        self,
        encrypted_obj: Dict[str, Any],
        content_key_provider: Callable[[str], bytes],
        token: str,
        content_id: str,
        associated_data: Optional[bytes] = None,
    ) -> Tuple[bool, Optional[str], Optional[bytes]]:
        """
        Расшифровка данных для стрима. Требует валидного стрим-токена.
        content_key_provider(content_id) -> bytes (32)
        """
        ok, err, st = self.verify_access(token, content_id)
        if not ok:
            return False, err, None

        try:
            # В идеале проверяем целостность до расшифровки
            # Здесь можем опционально вызвать verify_content_integrity, если есть сигнатуры
            # Но основной критерий — валидный токен.
            key = content_key_provider(content_id)
            pt = self.cipher.decrypt_content(
                ciphertext_b64=encrypted_obj["ciphertext_b64"],
                nonce_b64=encrypted_obj["nonce_b64"],
                tag_b64=encrypted_obj["tag_b64"],
                key=key,
                associated_data=associated_data,
            )
            logger.info("Decryption for stream succeeded content_id=%s owner=%s", content_id, st.owner_address)
            return True, None, pt
        except Exception as e:
            logger.exception("decrypt_for_stream failed")
            return False, str(e), None


def nft_proof_owner(ton_proof: Dict[str, Any]) -> str:
    """
    Извлечь адрес владельца из структуры tonProof запроса клиента.
    Совместимо с TonConnect unpack_wallet_info формой.
    """
    # Поддержка как плоской формы, так и вложенной ton_proof
    if "address" in ton_proof:
        return ton_proof["address"]
    if "account" in ton_proof and ton_proof["account"] and "address" in ton_proof["account"]:
        return ton_proof["account"]["address"]
    if "ton_proof" in ton_proof and ton_proof["ton_proof"] and "address" in ton_proof["ton_proof"]:
        return ton_proof["ton_proof"]["address"]
    # В противном случае бросаем: пусть вызывающий слой отловит
    raise ValueError("Cannot extract owner address from ton_proof")