StarPunk/starpunk/tokens.py

"""
Token management for Micropub IndieAuth integration

Handles:
- Access token generation and verification
- Authorization code generation and exchange
- Token hashing for secure storage (SHA256)
- Scope validation
- Token expiry management

Security:
- Tokens stored as SHA256 hashes (never plain text)
- Authorization codes use single-use pattern with replay protection
- Optional PKCE support for enhanced security
"""

import hashlib
import secrets
from datetime import datetime, timedelta
from typing import Optional, Dict, Any
from flask import current_app


# V1 supported scopes
SUPPORTED_SCOPES = ["create"]
DEFAULT_SCOPE = "create"

# Token and code expiry defaults
TOKEN_EXPIRY_DAYS = 90
AUTH_CODE_EXPIRY_MINUTES = 10


class TokenError(Exception):
    """Base exception for token-related errors"""
    pass


class InvalidTokenError(TokenError):
    """Raised when token is invalid or expired"""
    pass


class InvalidAuthorizationCodeError(TokenError):
    """Raised when authorization code is invalid, expired, or already used"""
    pass


def generate_token() -> str:
    """
    Generate a cryptographically secure random token

    Returns:
        URL-safe base64-encoded random token (43 characters)
    """
    return secrets.token_urlsafe(32)


def hash_token(token: str) -> str:
    """
    Generate SHA256 hash of token for secure storage

    Args:
        token: Plain text token

    Returns:
        Hexadecimal SHA256 hash
    """
    return hashlib.sha256(token.encode()).hexdigest()


def create_access_token(me: str, client_id: str, scope: str) -> str:
    """
    Create and store an access token in the database

    Args:
        me: User's identity URL
        client_id: Client application URL
        scope: Space-separated list of scopes

    Returns:
        Plain text access token (return to client, never logged or stored)

    Raises:
        TokenError: If token creation fails
    """
    # Generate token
    token = generate_token()
    token_hash_value = hash_token(token)

    # Calculate expiry
    # Use UTC to match SQLite's datetime('now') which returns UTC
    expires_at = (datetime.utcnow() + timedelta(days=TOKEN_EXPIRY_DAYS)).strftime('%Y-%m-%d %H:%M:%S')

    # Store in database
    from starpunk.database import get_db

    try:
        db = get_db(current_app)
        db.execute("""
            INSERT INTO tokens (token_hash, me, client_id, scope, expires_at)
            VALUES (?, ?, ?, ?, ?)
        """, (token_hash_value, me, client_id, scope, expires_at))
        db.commit()

        current_app.logger.info(
            f"Created access token for client_id={client_id}, scope={scope}"
        )

        return token

    except Exception as e:
        current_app.logger.error(f"Failed to create access token: {e}")
        raise TokenError(f"Failed to create access token: {e}")


def verify_token(token: str) -> Optional[Dict[str, Any]]:
    """
    Verify an access token and return token information

    Args:
        token: Plain text token to verify

    Returns:
        Dictionary with token info: {me, client_id, scope}
        None if token is invalid, expired, or revoked
    """
    if not token:
        return None

    # Hash the token for lookup
    token_hash_value = hash_token(token)

    from starpunk.database import get_db

    try:
        db = get_db(current_app)
        row = db.execute("""
            SELECT me, client_id, scope, id
            FROM tokens
            WHERE token_hash = ?
            AND expires_at > datetime('now')
            AND revoked_at IS NULL
        """, (token_hash_value,)).fetchone()

        if row:
            # Update last_used_at
            db.execute("""
                UPDATE tokens
                SET last_used_at = datetime('now')
                WHERE id = ?
            """, (row['id'],))
            db.commit()

            return {
                'me': row['me'],
                'client_id': row['client_id'],
                'scope': row['scope']
            }

        return None

    except Exception as e:
        current_app.logger.error(f"Token verification failed: {e}")
        return None


def revoke_token(token: str) -> bool:
    """
    Revoke an access token (soft deletion)

    Args:
        token: Plain text token to revoke

    Returns:
        True if token was revoked, False if not found
    """
    token_hash_value = hash_token(token)

    from starpunk.database import get_db

    try:
        db = get_db(current_app)
        cursor = db.execute("""
            UPDATE tokens
            SET revoked_at = datetime('now')
            WHERE token_hash = ?
            AND revoked_at IS NULL
        """, (token_hash_value,))
        db.commit()

        return cursor.rowcount > 0

    except Exception as e:
        current_app.logger.error(f"Token revocation failed: {e}")
        return False


def create_authorization_code(
    me: str,
    client_id: str,
    redirect_uri: str,
    scope: str = "",
    state: Optional[str] = None,
    code_challenge: Optional[str] = None,
    code_challenge_method: Optional[str] = None
) -> str:
    """
    Create and store an authorization code for token exchange

    Args:
        me: User's identity URL
        client_id: Client application URL
        redirect_uri: Client's redirect URI (must match during exchange)
        scope: Space-separated list of requested scopes (can be empty)
        state: Client's state parameter (optional)
        code_challenge: PKCE code challenge (optional)
        code_challenge_method: PKCE method, typically 'S256' (optional)

    Returns:
        Plain text authorization code (return to client)

    Raises:
        TokenError: If code creation fails
    """
    # Generate authorization code
    code = generate_token()
    code_hash_value = hash_token(code)

    # Calculate expiry (short-lived)
    # Use UTC to match SQLite's datetime('now') which returns UTC
    expires_at = (datetime.utcnow() + timedelta(minutes=AUTH_CODE_EXPIRY_MINUTES)).strftime('%Y-%m-%d %H:%M:%S')

    # Store in database
    from starpunk.database import get_db

    try:
        db = get_db(current_app)
        db.execute("""
            INSERT INTO authorization_codes (
                code_hash, me, client_id, redirect_uri, scope, state,
                code_challenge, code_challenge_method, expires_at
            )
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
        """, (
            code_hash_value, me, client_id, redirect_uri, scope, state,
            code_challenge, code_challenge_method, expires_at
        ))
        db.commit()

        current_app.logger.info(
            f"Created authorization code for client_id={client_id}, scope={scope}"
        )

        return code

    except Exception as e:
        current_app.logger.error(f"Failed to create authorization code: {e}")
        raise TokenError(f"Failed to create authorization code: {e}")


def exchange_authorization_code(
    code: str,
    client_id: str,
    redirect_uri: str,
    me: str,
    code_verifier: Optional[str] = None
) -> Dict[str, Any]:
    """
    Exchange authorization code for access token

    Args:
        code: Authorization code to exchange
        client_id: Client application URL (must match original request)
        redirect_uri: Redirect URI (must match original request)
        me: User's identity URL (must match original request)
        code_verifier: PKCE verifier (required if code_challenge was provided)

    Returns:
        Dictionary with: {me, client_id, scope}

    Raises:
        InvalidAuthorizationCodeError: If code is invalid, expired, used, or validation fails
    """
    if not code:
        raise InvalidAuthorizationCodeError("No authorization code provided")

    code_hash_value = hash_token(code)

    from starpunk.database import get_db

    try:
        db = get_db(current_app)

        # Look up authorization code
        row = db.execute("""
            SELECT me, client_id, redirect_uri, scope, code_challenge,
                   code_challenge_method, used_at
            FROM authorization_codes
            WHERE code_hash = ?
            AND expires_at > datetime('now')
        """, (code_hash_value,)).fetchone()

        if not row:
            raise InvalidAuthorizationCodeError(
                "Authorization code is invalid or expired"
            )

        # Check if already used (prevent replay attacks)
        if row['used_at']:
            raise InvalidAuthorizationCodeError(
                "Authorization code has already been used"
            )

        # Validate parameters match original authorization request
        if row['client_id'] != client_id:
            raise InvalidAuthorizationCodeError(
                "client_id does not match authorization request"
            )

        if row['redirect_uri'] != redirect_uri:
            raise InvalidAuthorizationCodeError(
                "redirect_uri does not match authorization request"
            )

        if row['me'] != me:
            raise InvalidAuthorizationCodeError(
                "me parameter does not match authorization request"
            )

        # Validate PKCE if code_challenge was provided
        if row['code_challenge']:
            if not code_verifier:
                raise InvalidAuthorizationCodeError(
                    "code_verifier required (PKCE was used during authorization)"
                )

            # Verify PKCE challenge
            if row['code_challenge_method'] == 'S256':
                # SHA256 hash of verifier
                computed_challenge = hashlib.sha256(
                    code_verifier.encode()
                ).hexdigest()
            else:
                # Plain (not recommended, but spec allows it)
                computed_challenge = code_verifier

            if computed_challenge != row['code_challenge']:
                raise InvalidAuthorizationCodeError(
                    "code_verifier does not match code_challenge"
                )

        # Mark code as used
        db.execute("""
            UPDATE authorization_codes
            SET used_at = datetime('now')
            WHERE code_hash = ?
        """, (code_hash_value,))
        db.commit()

        # Return authorization info for token creation
        return {
            'me': row['me'],
            'client_id': row['client_id'],
            'scope': row['scope']
        }

    except InvalidAuthorizationCodeError:
        # Re-raise validation errors
        raise

    except Exception as e:
        current_app.logger.error(f"Authorization code exchange failed: {e}")
        raise InvalidAuthorizationCodeError(f"Code exchange failed: {e}")


def validate_scope(requested_scope: str) -> str:
    """
    Validate and filter requested scopes to supported ones

    Args:
        requested_scope: Space-separated list of requested scopes

    Returns:
        Space-separated list of valid scopes (may be empty)
    """
    if not requested_scope:
        return ""

    requested = set(requested_scope.split())
    supported = set(SUPPORTED_SCOPES)
    valid_scopes = requested & supported

    return " ".join(sorted(valid_scopes)) if valid_scopes else ""


def check_scope(required: str, granted: str) -> bool:
    """
    Check if granted scopes include required scope

    Args:
        required: Required scope (single scope string)
        granted: Granted scopes (space-separated string)

    Returns:
        True if required scope is in granted scopes
    """
    if not granted:
        # IndieAuth spec: no scope means no access
        return False

    granted_scopes = set(granted.split())
    return required in granted_scopes