audit.py

"""
Audit Logging Service

Provides audit logging functionality for security and compliance.
Immutable entries - append only, no updates allowed.
"""

import uuid
from datetime import datetime
from typing import Any, Dict, Optional
from enum import Enum

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from backend.db.models import AuditLog, User, Tenant
from backend.db.session import get_db_session


class AuditAction(str, Enum):
    """Standard audit action types."""
    # Job actions
    JOB_CREATED = "JOB_CREATED"
    JOB_STARTED = "JOB_STARTED"
    JOB_COMPLETED = "JOB_COMPLETED"
    JOB_FAILED = "JOB_FAILED"
    JOB_CANCELLED = "JOB_CANCELLED"
    JOB_DELETED = "JOB_DELETED"
    
    # User actions
    USER_CREATED = "USER_CREATED"
    USER_UPDATED = "USER_UPDATED"
    USER_DELETED = "USER_DELETED"
    USER_LOGIN = "USER_LOGIN"
    USER_LOGOUT = "USER_LOGOUT"
    ROLE_CHANGED = "ROLE_CHANGED"
    
    # API Key actions
    API_KEY_CREATED = "API_KEY_CREATED"
    API_KEY_DELETED = "API_KEY_DELETED"
    API_KEY_ROTATED = "API_KEY_ROTATED"
    
    # Release actions
    RELEASE_APPROVED = "RELEASE_APPROVED"
    RELEASE_REJECTED = "RELEASE_REJECTED"
    RELEASE_BLOCKED = "RELEASE_BLOCKED"
    
    # Report actions
    REPORT_CREATED = "REPORT_CREATED"
    REPORT_EXPORTED = "REPORT_EXPORTED"
    REPORT_DELETED = "REPORT_DELETED"
    
    # Monitoring actions
    ALERT_CREATED = "ALERT_CREATED"
    ALERT_RESOLVED = "ALERT_RESOLVED"
    BASELINE_UPDATED = "BASELINE_UPDATED"
    
    # Tenant actions
    TENANT_CREATED = "TENANT_CREATED"
    TENANT_UPDATED = "TENANT_UPDATED"
    TENANT_DEACTIVATED = "TENANT_DEACTIVATED"
    
    # Model/Dataset actions
    MODEL_REGISTERED = "MODEL_REGISTERED"
    DATASET_REGISTERED = "DATASET_REGISTERED"


class ResourceType(str, Enum):
    """Resource types for audit logging."""
    JOB = "job"
    USER = "user"
    API_KEY = "api_key"
    RELEASE = "release"
    REPORT = "report"
    ALERT = "alert"
    TENANT = "tenant"
    MODEL = "model"
    DATASET = "dataset"


class AuditLogger:
    """
    Audit logging service for tracking user actions.
    
    All entries are immutable - append only, no updates allowed.
    """
    
    def __init__(self, db_session: AsyncSession):
        self.db = db_session
    
    async def log(
        self,
        tenant_id: uuid.UUID,
        user_id: uuid.UUID,
        action: AuditAction,
        resource_type: ResourceType,
        resource_id: Optional[uuid.UUID] = None,
        metadata: Optional[Dict[str, Any]] = None,
        request_id: Optional[str] = None,
        ip_address: Optional[str] = None,
        status: str = "success",
        error_message: Optional[str] = None,
    ) -> AuditLog:
        """
        Create an audit log entry.
        
        Args:
            tenant_id: Tenant ID
            user_id: User who performed the action
            action: Action performed
            resource_type: Type of resource affected
            resource_id: ID of the affected resource
            metadata: Additional action metadata
            request_id: Request ID for tracing
            ip_address: Client IP address
            status: Action status (success, failure, pending)
            error_message: Error message if action failed
            
        Returns:
            Created AuditLog entry
        """
        audit_log = AuditLog(
            tenant_id=tenant_id,
            user_id=user_id,
            action=action.value,
            resource_type=resource_type.value,
            resource_id=resource_id,
            metadata=metadata,
            request_id=request_id,
            ip_address=ip_address,
            status=status,
            error_message=error_message,
            timestamp=datetime.utcnow(),
        )
        
        self.db.add(audit_log)
        await self.db.commit()
        await self.db.refresh(audit_log)
        
        return audit_log
    
    async def log_job_created(
        self,
        tenant_id: uuid.UUID,
        user_id: uuid.UUID,
        job_id: uuid.UUID,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs,
    ) -> AuditLog:
        """Log job creation."""
        return await self.log(
            tenant_id=tenant_id,
            user_id=user_id,
            action=AuditAction.JOB_CREATED,
            resource_type=ResourceType.JOB,
            resource_id=job_id,
            metadata=metadata,
            **kwargs,
        )
    
    async def log_job_completed(
        self,
        tenant_id: uuid.UUID,
        user_id: uuid.UUID,
        job_id: uuid.UUID,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs,
    ) -> AuditLog:
        """Log job completion."""
        return await self.log(
            tenant_id=tenant_id,
            user_id=user_id,
            action=AuditAction.JOB_COMPLETED,
            resource_type=ResourceType.JOB,
            resource_id=job_id,
            metadata=metadata,
            **kwargs,
        )
    
    async def log_job_failed(
        self,
        tenant_id: uuid.UUID,
        user_id: uuid.UUID,
        job_id: uuid.UUID,
        error_message: str,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs,
    ) -> AuditLog:
        """Log job failure."""
        return await self.log(
            tenant_id=tenant_id,
            user_id=user_id,
            action=AuditAction.JOB_FAILED,
            resource_type=ResourceType.JOB,
            resource_id=job_id,
            metadata=metadata,
            status="failure",
            error_message=error_message,
            **kwargs,
        )
    
    async def log_api_key_created(
        self,
        tenant_id: uuid.UUID,
        user_id: uuid.UUID,
        api_key_id: uuid.UUID,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs,
    ) -> AuditLog:
        """Log API key creation."""
        return await self.log(
            tenant_id=tenant_id,
            user_id=user_id,
            action=AuditAction.API_KEY_CREATED,
            resource_type=ResourceType.API_KEY,
            resource_id=api_key_id,
            metadata=metadata,
            **kwargs,
        )
    
    async def log_release_approved(
        self,
        tenant_id: uuid.UUID,
        user_id: uuid.UUID,
        release_id: uuid.UUID,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs,
    ) -> AuditLog:
        """Log release approval."""
        return await self.log(
            tenant_id=tenant_id,
            user_id=user_id,
            action=AuditAction.RELEASE_APPROVED,
            resource_type=ResourceType.RELEASE,
            resource_id=release_id,
            metadata=metadata,
            **kwargs,
        )
    
    async def log_user_login(
        self,
        tenant_id: uuid.UUID,
        user_id: uuid.UUID,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs,
    ) -> AuditLog:
        """Log user login."""
        return await self.log(
            tenant_id=tenant_id,
            user_id=user_id,
            action=AuditAction.USER_LOGIN,
            resource_type=ResourceType.USER,
            resource_id=user_id,
            metadata=metadata,
            **kwargs,
        )
    
    async def log_role_changed(
        self,
        tenant_id: uuid.UUID,
        user_id: uuid.UUID,
        target_user_id: uuid.UUID,
        old_role: str,
        new_role: str,
        **kwargs,
    ) -> AuditLog:
        """Log role change."""
        return await self.log(
            tenant_id=tenant_id,
            user_id=user_id,
            action=AuditAction.ROLE_CHANGED,
            resource_type=ResourceType.USER,
            resource_id=target_user_id,
            metadata={"old_role": old_role, "new_role": new_role},
            **kwargs,
        )
    
    async def log_report_exported(
        self,
        tenant_id: uuid.UUID,
        user_id: uuid.UUID,
        report_id: uuid.UUID,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs,
    ) -> AuditLog:
        """Log report export."""
        return await self.log(
            tenant_id=tenant_id,
            user_id=user_id,
            action=AuditAction.REPORT_EXPORTED,
            resource_type=ResourceType.REPORT,
            resource_id=report_id,
            metadata=metadata,
            **kwargs,
        )
    
    async def get_tenant_audit_logs(
        self,
        tenant_id: uuid.UUID,
        limit: int = 100,
        offset: int = 0,
        action: Optional[AuditAction] = None,
        user_id: Optional[uuid.UUID] = None,
        resource_type: Optional[ResourceType] = None,
    ) -> list[AuditLog]:
        """
        Get audit logs for a tenant.
        
        Args:
            tenant_id: Tenant ID
            limit: Maximum number of entries to return
            offset: Number of entries to skip
            action: Filter by action type
            user_id: Filter by user
            resource_type: Filter by resource type
            
        Returns:
            List of audit log entries
        """
        query = select(AuditLog).where(
            AuditLog.tenant_id == tenant_id
        )
        
        if action:
            query = query.where(AuditLog.action == action.value)
        
        if user_id:
            query = query.where(AuditLog.user_id == user_id)
        
        if resource_type:
            query = query.where(AuditLog.resource_type == resource_type.value)
        
        query = (
            query
            .order_by(AuditLog.timestamp.desc())
            .limit(limit)
            .offset(offset)
        )
        
        result = await self.db.execute(query)
        return list(result.scalars().all())


async def create_audit_log(
    tenant_id: uuid.UUID,
    user_id: uuid.UUID,
    action: AuditAction,
    resource_type: ResourceType,
    resource_id: Optional[uuid.UUID] = None,
    metadata: Optional[Dict[str, Any]] = None,
    request_id: Optional[str] = None,
    ip_address: Optional[str] = None,
    status: str = "success",
    error_message: Optional[str] = None,
) -> AuditLog:
    """
    Convenience function to create an audit log entry.
    
    Creates a new database session for the operation.
    """
    async with get_db_session() as session:
        logger = AuditLogger(session)
        return await logger.log(
            tenant_id=tenant_id,
            user_id=user_id,
            action=action,
            resource_type=resource_type,
            resource_id=resource_id,
            metadata=metadata,
            request_id=request_id,
            ip_address=ip_address,
            status=status,
            error_message=error_message,
        )


# =============================================================================
# Audit context manager for automatic logging
# =============================================================================

class AuditContext:
    """
    Context manager for automatic audit logging.
    
    Usage:
        async with AuditContext(
            tenant_id=tenant_id,
            user_id=user_id,
            action=AuditAction.JOB_CREATED,
            resource_type=ResourceType.JOB,
            resource_id=job_id,
        ) as audit:
            # Do work
            await process_job(job_id)
            # Success automatically logged on exit
    """
    
    def __init__(
        self,
        tenant_id: uuid.UUID,
        user_id: uuid.UUID,
        action: AuditAction,
        resource_type: ResourceType,
        resource_id: Optional[uuid.UUID] = None,
        metadata: Optional[Dict[str, Any]] = None,
        request_id: Optional[str] = None,
        ip_address: Optional[str] = None,
    ):
        self.tenant_id = tenant_id
        self.user_id = user_id
        self.action = action
        self.resource_type = resource_type
        self.resource_id = resource_id
        self.metadata = metadata or {}
        self.request_id = request_id
        self.ip_address = ip_address
        self.status = "success"
        self.error_message = None
        self._audit_log: Optional[AuditLog] = None
    
    async def __aenter__(self) -> "AuditContext":
        return self
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        if exc_type is not None:
            self.status = "failure"
            self.error_message = str(exc_val) if exc_val else "Unknown error"
        
        # Log the action
        async with get_db_session() as session:
            logger = AuditLogger(session)
            self._audit_log = await logger.log(
                tenant_id=self.tenant_id,
                user_id=self.user_id,
                action=self.action,
                resource_type=self.resource_type,
                resource_id=self.resource_id,
                metadata=self.metadata,
                request_id=self.request_id,
                ip_address=self.ip_address,
                status=self.status,
                error_message=self.error_message,
            )
        
        return False  # Don't suppress exceptions
    
    @property
    def audit_log(self) -> Optional[AuditLog]:
        """Get the created audit log entry."""
        return self._audit_log