replypier/backend/app/core/security.py

"""
Copyright 2024-2026 ReplyPier

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import json
from datetime import datetime, timedelta, timezone
from typing import Optional
import re
import uuid
from jose import JWTError, jwt
from passlib.context import CryptContext
from app.core.config import  settings
from app.core.encryption import CURRENT_KEY_ID, get_keys
from app.core.logger import get_logger
import base64
import hashlib

logger = get_logger(__name__)

# Constants
SECRET_KEY = settings.SECRET_KEY
CONVERSATION_SECRET_KEY = settings.CONVERSATION_SECRET_KEY
ALGORITHM = "HS256"
ACCESS_TOKEN_EXPIRE_MINUTES = 30
REFRESH_TOKEN_EXPIRE_DAYS = 7

pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")

# Password policy for passwords one person sets on behalf of another (today:
# an admin resetting an agent's password). Mirrors the strength meter the UI
# draws — minimum length plus three of the four character classes — so the
# form and the API agree on what "strong enough" means.
MIN_PASSWORD_LENGTH = 8
MIN_PASSWORD_CHARACTER_CLASSES = 3

_PASSWORD_CHARACTER_CLASSES = (r"[A-Z]", r"[a-z]", r"[0-9]", r"[^A-Za-z0-9]")


def validate_password_strength(password: str) -> str:
    """Return the password, or raise ValueError describing what it is missing."""
    if len(password or "") < MIN_PASSWORD_LENGTH:
        raise ValueError(
            f"Password must be at least {MIN_PASSWORD_LENGTH} characters long"
        )

    classes = sum(1 for pattern in _PASSWORD_CHARACTER_CLASSES if re.search(pattern, password))
    if classes < MIN_PASSWORD_CHARACTER_CLASSES:
        raise ValueError(
            f"Password must include at least {MIN_PASSWORD_CHARACTER_CLASSES} of: "
            "uppercase letter, lowercase letter, number, special character"
        )

    return password


def _fernet():
    """The shared encryption-at-rest key, loaded by app.core.encryption.

    API keys predate the ``enc:v1:`` wire format and stay bare Fernet tokens, so
    they are encrypted here rather than through encrypt_value/decrypt_value — but
    the key itself has exactly one owner.
    """
    return get_keys()[CURRENT_KEY_ID]


def create_access_token(data: dict) -> str:
    to_encode = data.copy()
    expire = datetime.now(timezone.utc) + timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
    to_encode.update({"exp": expire, "type": "access"})
    return jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)


def create_refresh_token(data: dict) -> str:
    to_encode = data.copy()
    expire = datetime.now(timezone.utc) + timedelta(days=REFRESH_TOKEN_EXPIRE_DAYS)
    to_encode.update({"exp": expire, "type": "refresh"})
    return jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)


def verify_token(token: str) -> Optional[dict]:
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        return payload
    except JWTError:
        return None
    
def encrypt_api_key(api_key: str) -> str:
    """Encrypt API key before storing in database"""
    try:
        return _fernet().encrypt(api_key.encode()).decode()
    except Exception as e:
        logger.error(f"Encryption error: {str(e)}")
        raise ValueError("Failed to encrypt API key")


def decrypt_credentials(encrypted: str) -> dict:
    """Decrypt a JSON credential blob stored with encrypt_api_key."""
    return json.loads(decrypt_api_key(encrypted))


def decrypt_api_key(encrypted_key: str) -> str:
    """Decrypt API key from database"""
    try:
        return _fernet().decrypt(encrypted_key.encode()).decode()
    except Exception as e:
        logger.error(f"Decryption error: {str(e)}")
        raise ValueError("Failed to decrypt API key")


def create_conversation_token(widget_id: str, customer_id: Optional[str] = None, customer_email: Optional[str] = None, **extra_data) -> str:
    """
    Create a JWT token for widget conversations with revocation support via Redis
    
    Includes a JTI (JWT ID) claim for revocation. Token is stored in Redis with TTL.
    Optionally stores customer email for session tracking and bulk revocation.
    
    Args:
        widget_id: The widget ID
        customer_id: Optional customer ID
        customer_email: Optional customer email for session management
        **extra_data: Additional data to include in token
    """
    jti = str(uuid.uuid4())
    data = {
        "widget_id": widget_id,
        "type": "conversation",
        "jti": jti
    }

    # Only include `sub` when a customer_id is present. RFC 7519 (and python-jose)
    # require `sub` to be a string — encoding it as JSON null causes decode to fail
    # with JWTError("Subject must be a string") for anonymous widget visitors
    # (e.g. Explore flow), which silently drops all downstream claims (like `source`).
    if customer_id is not None:
        data["sub"] = customer_id

    # Add email to token if provided
    if customer_email:
        data["email"] = customer_email
    
    # Add any extra data to the token
    data.update(extra_data)
    
    # Set expiration to 30 days
    now = datetime.now(timezone.utc)
    expire = now + timedelta(days=30)
    data.update({"exp": expire})

    token = jwt.encode(data, CONVERSATION_SECRET_KEY, algorithm=ALGORITHM)

    # Store token JTI in Redis with TTL and optional email mapping (with widget_id)
    _store_token_in_redis(jti, int((expire - now).total_seconds()), customer_email, widget_id)
    
    return token


def verify_conversation_token(token: str) -> Optional[dict]:
    """
    Verify conversation token signature and check if revoked

    Returns:
        Token payload dict if valid and not revoked, None otherwise
    """
    try:
        # verify_sub=False tolerates existing tokens (already in browsers' localStorage)
        # that were issued with `sub: null`. New tokens from create_conversation_token
        # omit `sub` entirely when anonymous, so this option only matters for in-flight
        # tokens until the old 30-day TTLs expire.
        payload = jwt.decode(token, CONVERSATION_SECRET_KEY,
                             algorithms=[ALGORITHM],
                             options={"verify_sub": False})

        if payload.get("type") != "conversation":
            return None

        # Check if token JTI exists in Redis (if revoked, it won't exist)
        jti = payload.get("jti")
        if jti and not _is_token_in_redis(jti):
            logger.warning(f"Token revoked or expired: jti={jti[:8]}...")
            return None

        return payload
    except JWTError:
        return None


def verify_password(plain_password: str, hashed_password: str) -> bool:
    """Verify a password against its hash"""
    return pwd_context.verify(plain_password, hashed_password)


def get_password_hash(password: str) -> str:
    """Hash a password"""
    return pwd_context.hash(password)


# ============================================================================
# TOKEN STORAGE IN REDIS - Minimal approach
# ============================================================================

def _store_token_in_redis(jti: str, ttl_seconds: int, email: Optional[str] = None, widget_id: Optional[str] = None) -> bool:
    """
    Store token JTI in Redis with TTL and optionally store email/widget_id mapping

    Args:
        jti: JWT ID for the token
        ttl_seconds: Time to live in seconds
        email: Optional email to store for session tracking
        widget_id: Optional widget ID for widget-specific session tracking
    """
    try:
        from app.core.redis import redis_client
        if redis_client:
            # Store the token JTI
            redis_client.setex(f"token:{jti}", ttl_seconds, "1")

            # Store email mapping if provided (key: user_sessions:email, value: set of JTIs)
            if email:
                redis_client.sadd(f"user_sessions:{email}", jti)
                # Set same TTL for email->JTI mapping
                redis_client.expire(f"user_sessions:{email}", ttl_seconds)
                
                # If widget_id provided, also store widget-specific sessions
                if widget_id:
                    widget_key = f"user_sessions:{email}:widget:{widget_id}"
                    redis_client.sadd(widget_key, jti)
                    redis_client.expire(widget_key, ttl_seconds)
                    
                    # Maintain reverse index: track all widget keys for this email
                    # for efficient bulk deletion without scanning
                    redis_client.sadd(f"user_sessions:{email}:widget_keys", widget_key)
                    redis_client.expire(f"user_sessions:{email}:widget_keys", ttl_seconds)
            
            return True
    except Exception as e:
        logger.warning(f"Failed to store token in Redis: {str(e)}")
    return False


def _is_token_in_redis(jti: str) -> bool:
    """Check if token JTI exists in Redis"""
    try:
        from app.core.redis import redis_client
        if redis_client:
            return redis_client.exists(f"token:{jti}") > 0
    except Exception as e:
        logger.debug(f"Failed to check token in Redis: {str(e)}")
    # If Redis unavailable, assume valid (don't block requests)
    return True


def revoke_token(jti: str) -> bool:
    """Revoke token by deleting JTI from Redis"""
    try:
        from app.core.redis import redis_client
        if redis_client:
            redis_client.delete(f"token:{jti}")
            logger.info(f"Token revoked: {jti[:8]}...")
            return True
    except Exception as e:
        logger.error(f"Failed to revoke token: {str(e)}")
    return False


def revoke_user_sessions(email: str) -> bool:
    """
    Revoke all sessions for a user by email.
    
    Args:
        email: The user's email address
    
    Returns:
        True if successful, False otherwise
    """
    try:
        from app.core.redis import redis_client
        if redis_client:
            # Get all JTIs for this email across all widgets
            jti_set = redis_client.smembers(f"user_sessions:{email}")
            
            if jti_set:
                # Delete each JTI
                for jti in jti_set:
                    redis_client.delete(f"token:{jti.decode() if isinstance(jti, bytes) else jti}")
                
                # Delete the email mapping
                redis_client.delete(f"user_sessions:{email}")
                
                # Also delete all widget-specific mappings using the reverse index set
                widget_keys_set = f"user_sessions:{email}:widget_keys"
                widget_keys = redis_client.smembers(widget_keys_set)
                if widget_keys:
                    for key in widget_keys:
                        redis_client.delete(key.decode() if isinstance(key, bytes) else key)
                    redis_client.delete(widget_keys_set)
                # NOTE: When creating or deleting widget-specific session keys elsewhere,
                # be sure to add/remove the key from the widget_keys_set for this user.
                
                logger.info(f"Revoked {len(jti_set)} sessions for user: {email}")
                return True
            else:
                logger.info(f"No active sessions found for user: {email}")
                return True
    except Exception as e:
        logger.error(f"Failed to revoke user sessions: {str(e)}")
    return False


def revoke_user_sessions_by_widget(email: str, widget_id: str) -> bool:
    """
    Revoke all sessions for a user in a specific widget.
    
    Args:
        email: The user's email address
        widget_id: The widget ID to revoke sessions for
    
    Returns:
        True if successful, False otherwise
    """
    try:
        from app.core.redis import redis_client
        if redis_client:
            # Get all JTIs for this email in this specific widget
            jti_set = redis_client.smembers(f"user_sessions:{email}:widget:{widget_id}")
            
            if jti_set:
                # Delete each JTI
                for jti in jti_set:
                    jti_str = jti.decode() if isinstance(jti, bytes) else jti
                    redis_client.delete(f"token:{jti_str}")
                    # Also remove from general user_sessions set
                    redis_client.srem(f"user_sessions:{email}", jti_str)
                
                # Delete the widget-specific mapping
                redis_client.delete(f"user_sessions:{email}:widget:{widget_id}")
                
                logger.info(f"Revoked {len(jti_set)} sessions for user: {email} in widget: {widget_id}")
                return True
            else:
                logger.info(f"No active sessions found for user: {email} in widget: {widget_id}")
                return True
    except Exception as e:
        logger.error(f"Failed to revoke user sessions by widget: {str(e)}")
    return False


def get_user_active_sessions(email: str, widget_id: Optional[str] = None) -> list:
    """
    Get all active session JTIs for a user by email.
    
    Args:
        email: The user's email address
        widget_id: Optional widget ID to get sessions only for that widget
    
    Returns:
        List of active JTI strings for the user (filtered by widget if provided)
    """
    try:
        from app.core.redis import redis_client
        if redis_client:
            if widget_id:
                # Get sessions only for this specific widget
                jti_set = redis_client.smembers(f"user_sessions:{email}:widget:{widget_id}")
            else:
                # Get all sessions for this user across all widgets
                jti_set = redis_client.smembers(f"user_sessions:{email}")
            
            return [jti.decode() if isinstance(jti, bytes) else jti for jti in jti_set]
    except Exception as e:
        logger.error(f"Failed to get user sessions: {str(e)}")
    return []


def get_existing_valid_token_jti(email: str, widget_id: str, customer_id: str) -> Optional[str]:
    """
    Check if there's an existing valid token JTI for this customer/widget combination.
    
    Returns the JTI (JWT ID) of an existing valid token, not the token itself.
    This prevents token multiplication when the same user refreshes the page multiple times.
    
    Args:
        email: Customer email
        widget_id: Widget ID
        customer_id: Customer ID
    
    Returns:
        The JTI (JWT ID) string of an existing valid token if found, None otherwise
    """
    try:
        from app.core.redis import redis_client
        if not redis_client:
            return None
        
        # Get all active JTIs for this email in this widget
        jti_set = redis_client.smembers(f"user_sessions:{email}:widget:{widget_id}")
        
        if not jti_set:
            logger.debug(f"No existing sessions found for {email} in widget {widget_id}")
            return None
        
        # If we have active tokens, return the first valid one
        # In a typical flow, there should only be one per widget per user
        for jti in jti_set:
            jti_str = jti.decode() if isinstance(jti, bytes) else jti
            # Check if the token still exists in Redis (not revoked)
            if redis_client.exists(f"token:{jti_str}") > 0:
                logger.info(f"Found existing valid token for {email} in widget {widget_id}, reusing it")
                return jti_str
        
        logger.debug(f"Active sessions found for {email} in widget {widget_id}, but none are valid in Redis")
        return None
        
    except Exception as e:
        logger.error(f"Error checking existing tokens: {str(e)}")
        return None


# ============================================================================
# WIDGET API KEY MANAGEMENT
# ============================================================================

def generate_widget_api_key() -> str:
    """
    Generate a secure random API key with 'wak_' prefix.

    Format: wak_<base64url> (e.g., wak_Xu8f3kLp9mN2...)
    Entropy: 256 bits (32 random bytes)
    Length: ~47 characters total (4 + 43)

    Returns:
        str: API key in format 'wak_<base64url>'
    """
    import secrets

    # Generate 32 random bytes (256 bits of entropy)
    random_bytes = secrets.token_bytes(32)
    # Encode as URL-safe base64 without padding
    key_suffix = base64.urlsafe_b64encode(random_bytes).decode('utf-8').rstrip('=')
    return f"wak_{key_suffix}"


def hash_widget_api_key(api_key: str) -> str:
    """
    Hash API key using bcrypt (same as password hashing).

    Args:
        api_key: Plain text API key

    Returns:
        str: bcrypt hash of the API key
    """
    return get_password_hash(api_key)


def verify_widget_api_key(plain_key: str, hashed_key: str) -> bool:
    """
    Verify API key against its hash.

    Args:
        plain_key: Plain text API key from request
        hashed_key: Stored bcrypt hash

    Returns:
        bool: True if key matches hash
    """
    return verify_password(plain_key, hashed_key)


def widget_api_key_lookup(api_key: str) -> str:
    """
    Indexed lookup value for a widget API key: its SHA-256.

    A fast hash is safe here because the key is 256 random bits, nothing to
    brute-force; it lets a request find its app in one query instead of
    bcrypt-checking every app.
    """
    return hashlib.sha256(api_key.encode()).hexdigest()