auth.py

"""
auth.py — Authentication & RBAC System
Provides user authentication and role-based access control
"""
from __future__ import annotations

import hashlib
import logging
import os
import secrets
import time
import uuid
from datetime import datetime, timedelta
from functools import wraps
from typing import Any, Dict, List, Optional, Set

from fastapi import HTTPException, Request, Depends
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from pydantic import BaseModel

logger = logging.getLogger("auth")

# ─── Configuration ─────────────────────────────────────────────────────────────
SECRET_KEY = os.environ.get("AUTH_SECRET_KEY", secrets.token_urlsafe(32))
TOKEN_EXPIRY_HOURS = int(os.environ.get("AUTH_TOKEN_EXPIRY_HOURS", "24"))
REFRESH_TOKEN_EXPIRY_DAYS = int(os.environ.get("AUTH_REFRESH_TOKEN_EXPIRY_DAYS", "7"))

security = HTTPBearer(auto_error=False)

# ─── User Models ───────────────────────────────────────────────────────────────
class UserRole(BaseModel):
    name: str
    permissions: List[str]


class User(BaseModel):
    id: str
    username: str
    email: str
    password_hash: str
    roles: List[str]
    created_at: float
    last_login: Optional[float] = None
    is_active: bool = True
    metadata: Dict[str, Any] = {}


class UserCreate(BaseModel):
    username: str
    email: str
    password: str
    roles: List[str] = ["user"]


class TokenData(BaseModel):
    user_id: str
    username: str
    roles: List[str]
    exp: float
    token_type: str = "access"


class RefreshTokenData(BaseModel):
    user_id: str
    token_id: str
    exp: float


# ─── Permission Definitions ────────────────────────────────────────────────────
PERMISSIONS = {
    # Chat & Conversation
    "chat:send": "Send messages",
    "chat:read": "Read conversations",
    "chat:delete": "Delete conversations",
    
    # Agent Operations
    "agent:run": "Run autonomous agent",
    "agent:plan": "Generate plans",
    "agent:cancel": "Cancel agent tasks",
    
    # Code Execution
    "code:execute": "Execute code",
    "code:read": "Read execution history",
    
    # Memory
    "memory:read": "Read memory",
    "memory:write": "Write to memory",
    "memory:delete": "Delete memories",
    
    # Files
    "files:read": "Read files",
    "files:write": "Write files",
    "files:delete": "Delete files",
    
    # Developer Tools
    "dev:generate": "Generate code",
    "dev:deploy": "Deploy applications",
    "dev:github": "GitHub operations",
    
    # Admin
    "admin:users": "Manage users",
    "admin:settings": "Modify settings",
    "admin:logs": "View logs",
    "admin:metrics": "View metrics",
}

# ─── Default Roles ──────────────────────────────────────────────────────────────
DEFAULT_ROLES = {
    "anonymous": ["chat:send", "chat:read", "memory:read"],
    "user": [
        "chat:send", "chat:read", "chat:delete",
        "agent:run", "agent:plan", "agent:cancel",
        "code:execute", "code:read",
        "memory:read", "memory:write",
        "files:read", "files:write",
        "dev:generate",
    ],
    "developer": [
        "chat:send", "chat:read", "chat:delete",
        "agent:run", "agent:plan", "agent:cancel",
        "code:execute", "code:read",
        "memory:read", "memory:write", "memory:delete",
        "files:read", "files:write", "files:delete",
        "dev:generate", "dev:deploy", "dev:github",
    ],
    "admin": list(PERMISSIONS.keys()),
}


# ─── Auth Manager ─────────────────────────────────────────────────────────────
class AuthManager:
    """
    Manages authentication and authorization.
    
    Features:
    - User registration/login
    - JWT token management
    - Role-based access control
    - Permission checking
    """
    
    def __init__(self):
        self._users: Dict[str, User] = {}
        self._tokens: Dict[str, TokenData] = {}
        self._refresh_tokens: Dict[str, RefreshTokenData] = {}
        self._load_default_users()
    
    def _load_default_users(self):
        """Load default users from environment or create admin."""
        admin_email = os.environ.get("ADMIN_EMAIL", "admin@example.com")
        admin_password = os.environ.get("ADMIN_PASSWORD")
        
        if admin_password:
            self.create_user(
                username="admin",
                email=admin_email,
                password=admin_password,
                roles=["admin"],
            )
    
    def _hash_password(self, password: str) -> str:
        """Hash password with salt."""
        salt = secrets.token_hex(16)
        hash_obj = hashlib.pbkdf2_hmac(
            'sha256',
            password.encode(),
            salt.encode(),
            100000,
        )
        return f"{salt}${hash_obj.hex()}"
    
    def _verify_password(self, password: str, password_hash: str) -> bool:
        """Verify password against hash."""
        try:
            salt, stored_hash = password_hash.split("$")
            hash_obj = hashlib.pbkdf2_hmac(
                'sha256',
                password.encode(),
                salt.encode(),
                100000,
            )
            return hash_obj.hex() == stored_hash
        except Exception:
            return False
    
    def _generate_token(self, user_id: str, username: str, roles: List[str]) -> str:
        """Generate JWT-style token."""
        token_id = secrets.token_urlsafe(32)
        exp = time.time() + (TOKEN_EXPIRY_HOURS * 3600)
        
        token_data = TokenData(
            user_id=user_id,
            username=username,
            roles=roles,
            exp=exp,
        )
        
        # Simple token format (not JWT, for simplicity)
        token = f"{token_id}.{user_id}.{exp}"
        self._tokens[token_id] = token_data
        return token
    
    def _generate_refresh_token(self, user_id: str) -> str:
        """Generate refresh token."""
        token_id = secrets.token_urlsafe(32)
        exp = time.time() + (REFRESH_TOKEN_EXPIRY_DAYS * 86400)
        
        refresh_data = RefreshTokenData(
            user_id=user_id,
            token_id=token_id,
            exp=exp,
        )
        
        self._refresh_tokens[token_id] = refresh_data
        return token_id
    
    def create_user(
        self,
        username: str,
        email: str,
        password: str,
        roles: List[str] = None,
    ) -> Optional[User]:
        """Create a new user."""
        # Check for existing user
        if any(u.username == username for u in self._users.values()):
            logger.warning(f"User already exists: {username}")
            return None
        
        if any(u.email == email for u in self._users.values()):
            logger.warning(f"Email already registered: {email}")
            return None
        
        user = User(
            id=str(uuid.uuid4()),
            username=username,
            email=email,
            password_hash=self._hash_password(password),
            roles=roles or ["user"],
            created_at=time.time(),
        )
        
        self._users[user.id] = user
        logger.info(f"User created: {username}")
        return user
    
    def authenticate(self, username: str, password: str) -> Optional[Dict]:
        """Authenticate user and return tokens."""
        user = next((u for u in self._users.values() if u.username == username), None)
        
        if not user or not self._verify_password(password, user.password_hash):
            logger.warning(f"Authentication failed for: {username}")
            return None
        
        if not user.is_active:
            logger.warning(f"Inactive user attempted login: {username}")
            return None
        
        # Update last login
        user.last_login = time.time()
        
        # Generate tokens
        access_token = self._generate_token(user.id, user.username, user.roles)
        refresh_token = self._generate_refresh_token(user.id)
        
        return {
            "access_token": access_token,
            "refresh_token": refresh_token,
            "token_type": "Bearer",
            "expires_in": TOKEN_EXPIRY_HOURS * 3600,
            "user": {
                "id": user.id,
                "username": user.username,
                "email": user.email,
                "roles": user.roles,
            },
        }
    
    def refresh_access_token(self, refresh_token: str) -> Optional[Dict]:
        """Refresh access token using refresh token."""
        token_data = self._refresh_tokens.get(refresh_token)
        
        if not token_data or token_data.exp < time.time():
            return None
        
        user = self._users.get(token_data.user_id)
        if not user or not user.is_active:
            return None
        
        # Generate new access token
        access_token = self._generate_token(user.id, user.username, user.roles)
        
        return {
            "access_token": access_token,
            "token_type": "Bearer",
            "expires_in": TOKEN_EXPIRY_HOURS * 3600,
        }
    
    def verify_token(self, token: str) -> Optional[TokenData]:
        """Verify and decode access token."""
        try:
            parts = token.split(".")
            if len(parts) != 3:
                return None
            
            token_id, user_id, exp = parts
            token_data = self._tokens.get(token_id)
            
            if not token_data or token_data.exp < time.time():
                return None
            
            return token_data
        except Exception:
            return None
    
    def get_user(self, user_id: str) -> Optional[User]:
        """Get user by ID."""
        return self._users.get(user_id)
    
    def get_user_by_username(self, username: str) -> Optional[User]:
        """Get user by username."""
        return next((u for u in self._users.values() if u.username == username), None)
    
    def has_permission(self, user_id: str, permission: str) -> bool:
        """Check if user has a specific permission."""
        user = self._users.get(user_id)
        if not user or not user.is_active:
            return False
        
        for role in user.roles:
            role_perms = DEFAULT_ROLES.get(role, [])
            if permission in role_perms:
                return True
        
        return False
    
    def has_any_permission(self, user_id: str, permissions: List[str]) -> bool:
        """Check if user has any of the specified permissions."""
        return any(self.has_permission(user_id, p) for p in permissions)
    
    def has_all_permissions(self, user_id: str, permissions: List[str]) -> bool:
        """Check if user has all specified permissions."""
        return all(self.has_permission(user_id, p) for p in permissions)
    
    def list_users(self) -> List[Dict]:
        """List all users (admin only)."""
        return [
            {
                "id": u.id,
                "username": u.username,
                "email": u.email,
                "roles": u.roles,
                "is_active": u.is_active,
                "created_at": u.created_at,
                "last_login": u.last_login,
            }
            for u in self._users.values()
        ]
    
    def update_user(self, user_id: str, **kwargs) -> Optional[User]:
        """Update user properties."""
        user = self._users.get(user_id)
        if not user:
            return None
        
        for key, value in kwargs.items():
            if hasattr(user, key) and key != "id":
                setattr(user, key, value)
        
        return user
    
    def delete_user(self, user_id: str) -> bool:
        """Delete user (admin only)."""
        if user_id in self._users:
            del self._users[user_id]
            logger.info(f"User deleted: {user_id}")
            return True
        return False


# ─── Global auth manager ───────────────────────────────────────────────────────
auth_manager = AuthManager()


# ─── Dependency Functions ──────────────────────────────────────────────────────
async def get_current_user(
    request: Request,
    credentials: HTTPAuthorizationCredentials = Depends(security),
) -> Optional[TokenData]:
    """Get current authenticated user from token."""
    if not credentials:
        # Try to get from session/cookie
        token = request.cookies.get("access_token") or request.headers.get("X-User-ID", "anonymous")
        
        # For anonymous access, create guest token data
        if token == "anonymous" or not auth_manager.verify_token(token):
            return TokenData(
                user_id="anonymous",
                username="anonymous",
                roles=["anonymous"],
                exp=time.time() + 86400,
            )
    
    token_data = auth_manager.verify_token(credentials.credentials)
    if not token_data:
        raise HTTPException(401, "Invalid or expired token")
    
    return token_data


async def require_permission(permission: str):
    """Dependency to require a specific permission."""
    async def check_permission(user: TokenData = Depends(get_current_user)):
        if not auth_manager.has_permission(user.user_id, permission):
            raise HTTPException(403, f"Permission denied: {permission}")
        return user
    return check_permission


async def require_any_permission(permissions: List[str]):
    """Dependency to require any of the specified permissions."""
    async def check_permissions(user: TokenData = Depends(get_current_user)):
        if not auth_manager.has_any_permission(user.user_id, permissions):
            raise HTTPException(403, f"Permission denied: requires one of {permissions}")
        return user
    return check_permissions


def require_auth(func):
    """Decorator to require authentication."""
    @wraps(func)
    async def wrapper(*args, user: TokenData = Depends(get_current_user), **kwargs):
        if user.user_id == "anonymous":
            raise HTTPException(401, "Authentication required")
        return await func(*args, user=user, **kwargs)
    return wrapper


def require_permission_check(permission: str):
    """Decorator to require specific permission."""
    def decorator(func):
        @wraps(func)
        async def wrapper(*args, user: TokenData = Depends(get_current_user), **kwargs):
            if not auth_manager.has_permission(user.user_id, permission):
                raise HTTPException(403, f"Permission denied: {permission}")
            return await func(*args, user=user, **kwargs)
        return wrapper
    return decorator


# ─── Auth API Routes ───────────────────────────────────────────────────────────
def register_auth_routes(app):
    """Register authentication routes with FastAPI app."""
    from fastapi import APIRouter, HTTPException
    
    router = APIRouter(prefix="/auth", tags=["auth"])
    
    @router.post("/register")
    async def register(user_data: UserCreate):
        """Register a new user."""
        user = auth_manager.create_user(
            username=user_data.username,
            email=user_data.email,
            password=user_data.password,
            roles=user_data.roles,
        )
        if not user:
            raise HTTPException(400, "User already exists")
        return {
            "message": "User created successfully",
            "user_id": user.id,
        }
    
    @router.post("/login")
    async def login(username: str, password: str):
        """Login and get access token."""
        result = auth_manager.authenticate(username, password)
        if not result:
            raise HTTPException(401, "Invalid credentials")
        return result
    
    @router.post("/refresh")
    async def refresh(refresh_token: str):
        """Refresh access token."""
        result = auth_manager.refresh_access_token(refresh_token)
        if not result:
            raise HTTPException(401, "Invalid refresh token")
        return result
    
    @router.get("/me")
    async def get_me(user: TokenData = Depends(get_current_user)):
        """Get current user info."""
        full_user = auth_manager.get_user(user.user_id)
        if full_user:
            return {
                "id": full_user.id,
                "username": full_user.username,
                "email": full_user.email,
                "roles": full_user.roles,
                "created_at": full_user.created_at,
            }
        return {
            "id": user.user_id,
            "username": user.username,
            "roles": user.roles,
        }
    
    @router.get("/permissions")
    async def list_permissions():
        """List all available permissions."""
        return {"permissions": PERMISSIONS}
    
    @router.get("/users")
    async def list_users(user: TokenData = Depends(get_current_user)):
        """List all users (admin only)."""
        if not auth_manager.has_permission(user.user_id, "admin:users"):
            raise HTTPException(403, "Admin access required")
        return {"users": auth_manager.list_users()}
    
    app.include_router(router)