chore/backend/utils/digest_token.py

import hashlib
import hmac
import json
import os
import time
import uuid
from datetime import datetime, timedelta, timezone

from db.digest_action_tokens import insert_token, get_token_by_id, mark_token_used
from models.digest_action_token import DigestActionToken


def _get_secret() -> bytes:
    secret = os.environ.get('DIGEST_TOKEN_SECRET')
    if not secret:
        raise RuntimeError('DIGEST_TOKEN_SECRET environment variable is not set.')
    return secret.encode('utf-8')


def _sign(payload: dict) -> str:
    """Return HMAC-SHA256 hex digest of the canonical JSON payload."""
    canonical = json.dumps(payload, sort_keys=True, separators=(',', ':'))
    return hmac.new(_get_secret(), canonical.encode('utf-8'), hashlib.sha256).hexdigest()


def create_action_token(
    user_id: str,
    child_id: str,
    entity_id: str,
    entity_type: str,
    action: str,
    expiry_hours: int = 24,
) -> DigestActionToken:
    """Create and persist a signed action token. Returns the saved token."""
    token_id = str(uuid.uuid4())
    expires_at = (datetime.now(timezone.utc) + timedelta(hours=expiry_hours)).isoformat()

    payload = {
        'id': token_id,
        'user_id': user_id,
        'child_id': child_id,
        'entity_id': entity_id,
        'entity_type': entity_type,
        'action': action,
        'expires_at': expires_at,
    }
    signature = _sign(payload)

    token = DigestActionToken(
        id=token_id,
        user_id=user_id,
        child_id=child_id,
        entity_id=entity_id,
        entity_type=entity_type,
        action=action,
        expires_at=expires_at,
        used=False,
        signature=signature,
    )
    insert_token(token)
    return token


def validate_and_consume_token(token_id: str) -> DigestActionToken | None:
    """
    Validate a token. Returns the token if valid and marks it as used.
    Returns None if token is missing, expired, used, or has an invalid signature.
    """
    token = get_token_by_id(token_id)
    if not token:
        return None

    if token.used:
        return None

    # Check expiry
    try:
        expires_at = datetime.fromisoformat(token.expires_at)
        if expires_at.tzinfo is None:
            expires_at = expires_at.replace(tzinfo=timezone.utc)
        if datetime.now(timezone.utc) > expires_at:
            return None
    except (ValueError, TypeError):
        return None

    # Verify HMAC signature
    payload = {
        'id': token.id,
        'user_id': token.user_id,
        'child_id': token.child_id,
        'entity_id': token.entity_id,
        'entity_type': token.entity_type,
        'action': token.action,
        'expires_at': token.expires_at,
    }
    expected_sig = _sign(payload)
    if not hmac.compare_digest(token.signature, expected_sig):
        return None

    mark_token_used(token_id)
    return token


def peek_token(token_id: str) -> DigestActionToken | None:
    """
    Validate a token (signature, expiry, not-used) WITHOUT consuming it.
    Returns the token if valid, None otherwise.
    """
    token = get_token_by_id(token_id)
    if not token or token.used:
        return None

    try:
        expires_at = datetime.fromisoformat(token.expires_at)
        if expires_at.tzinfo is None:
            expires_at = expires_at.replace(tzinfo=timezone.utc)
        if datetime.now(timezone.utc) > expires_at:
            return None
    except (ValueError, TypeError):
        return None

    payload = {
        'id': token.id,
        'user_id': token.user_id,
        'child_id': token.child_id,
        'entity_id': token.entity_id,
        'entity_type': token.entity_type,
        'action': token.action,
        'expires_at': token.expires_at,
    }
    expected_sig = _sign(payload)
    if not hmac.compare_digest(token.signature, expected_sig):
        return None

    return token


def create_unsubscribe_token(user_id: str) -> str:
    """
    Create a simple signed unsubscribe token string (not persisted in DB).
    Format: '<user_id>.<expiry_ts>.<signature>'
    """
    expiry_ts = int(time.time()) + 30 * 24 * 3600  # 30 days
    payload_str = f"{user_id}:{expiry_ts}"
    sig = hmac.new(_get_secret(), payload_str.encode('utf-8'), hashlib.sha256).hexdigest()
    import base64
    token = base64.urlsafe_b64encode(f"{payload_str}:{sig}".encode()).decode()
    return token


def validate_unsubscribe_token(token: str) -> str | None:
    """
    Validate an unsubscribe token.
    Returns user_id if valid, None otherwise.
    """
    import base64
    try:
        decoded = base64.urlsafe_b64decode(token.encode()).decode()
        parts = decoded.rsplit(':', 1)
        if len(parts) != 2:
            return None
        payload_str, sig = parts
        expected_sig = hmac.new(_get_secret(), payload_str.encode('utf-8'), hashlib.sha256).hexdigest()
        if not hmac.compare_digest(sig, expected_sig):
            return None
        sub_parts = payload_str.split(':', 1)
        if len(sub_parts) != 2:
            return None
        user_id, expiry_str = sub_parts
        expiry_ts = int(expiry_str)
        if time.time() > expiry_ts:
            return None
        return user_id
    except Exception:
        return None