| from contextlib import asynccontextmanager |
| from dataclasses import asdict, dataclass |
| from enum import Enum |
| import re |
| from typing import ( |
| TYPE_CHECKING, |
| Any, |
| AsyncGenerator, |
| Dict, |
| MutableMapping, |
| Optional, |
| cast, |
| ) |
| import uuid |
|
|
| from asgiref.typing import ( |
| ASGI3Application, |
| ASGIReceiveCallable, |
| ASGIReceiveEvent, |
| ASGISendCallable, |
| ASGISendEvent, |
| Scope as ASGIScope, |
| ) |
| from loguru import logger |
| from starlette.requests import Request |
|
|
| from open_webui.env import AUDIT_LOG_LEVEL, MAX_BODY_LOG_SIZE |
| from open_webui.utils.auth import get_current_user, get_http_authorization_cred |
| from open_webui.models.users import UserModel |
|
|
|
|
| if TYPE_CHECKING: |
| from loguru import Logger |
|
|
|
|
| @dataclass(frozen=True) |
| class AuditLogEntry: |
| |
| id: str |
| user: dict[str, Any] |
| audit_level: str |
| verb: str |
| request_uri: str |
| user_agent: Optional[str] = None |
| source_ip: Optional[str] = None |
| |
| request_object: Any = None |
| |
| response_object: Any = None |
| response_status_code: Optional[int] = None |
|
|
|
|
| class AuditLevel(str, Enum): |
| NONE = "NONE" |
| METADATA = "METADATA" |
| REQUEST = "REQUEST" |
| REQUEST_RESPONSE = "REQUEST_RESPONSE" |
|
|
|
|
| class AuditLogger: |
| """ |
| A helper class that encapsulates audit logging functionality. It uses Loguru’s logger with an auditable binding to ensure that audit log entries are filtered correctly. |
| |
| Parameters: |
| logger (Logger): An instance of Loguru’s logger. |
| """ |
|
|
| def __init__(self, logger: "Logger"): |
| self.logger = logger.bind(auditable=True) |
|
|
| def write( |
| self, |
| audit_entry: AuditLogEntry, |
| *, |
| log_level: str = "INFO", |
| extra: Optional[dict] = None, |
| ): |
|
|
| entry = asdict(audit_entry) |
|
|
| if extra: |
| entry["extra"] = extra |
|
|
| self.logger.log( |
| log_level, |
| "", |
| **entry, |
| ) |
|
|
|
|
| class AuditContext: |
| """ |
| Captures and aggregates the HTTP request and response bodies during the processing of a request. It ensures that only a configurable maximum amount of data is stored to prevent excessive memory usage. |
| |
| Attributes: |
| request_body (bytearray): Accumulated request payload. |
| response_body (bytearray): Accumulated response payload. |
| max_body_size (int): Maximum number of bytes to capture. |
| metadata (Dict[str, Any]): A dictionary to store additional audit metadata (user, http verb, user agent, etc.). |
| """ |
|
|
| def __init__(self, max_body_size: int = MAX_BODY_LOG_SIZE): |
| self.request_body = bytearray() |
| self.response_body = bytearray() |
| self.max_body_size = max_body_size |
| self.metadata: Dict[str, Any] = {} |
|
|
| def add_request_chunk(self, chunk: bytes): |
| if len(self.request_body) < self.max_body_size: |
| self.request_body.extend( |
| chunk[: self.max_body_size - len(self.request_body)] |
| ) |
|
|
| def add_response_chunk(self, chunk: bytes): |
| if len(self.response_body) < self.max_body_size: |
| self.response_body.extend( |
| chunk[: self.max_body_size - len(self.response_body)] |
| ) |
|
|
|
|
| class AuditLoggingMiddleware: |
| """ |
| ASGI middleware that intercepts HTTP requests and responses to perform audit logging. It captures request/response bodies (depending on audit level), headers, HTTP methods, and user information, then logs a structured audit entry at the end of the request cycle. |
| """ |
|
|
| AUDITED_METHODS = {"PUT", "PATCH", "DELETE", "POST"} |
|
|
| def __init__( |
| self, |
| app: ASGI3Application, |
| *, |
| excluded_paths: Optional[list[str]] = None, |
| max_body_size: int = MAX_BODY_LOG_SIZE, |
| audit_level: AuditLevel = AuditLevel.NONE, |
| ) -> None: |
| self.app = app |
| self.audit_logger = AuditLogger(logger) |
| self.excluded_paths = excluded_paths or [] |
| self.max_body_size = max_body_size |
| self.audit_level = audit_level |
|
|
| async def __call__( |
| self, |
| scope: ASGIScope, |
| receive: ASGIReceiveCallable, |
| send: ASGISendCallable, |
| ) -> None: |
| if scope["type"] != "http": |
| return await self.app(scope, receive, send) |
|
|
| request = Request(scope=cast(MutableMapping, scope)) |
|
|
| if self._should_skip_auditing(request): |
| return await self.app(scope, receive, send) |
|
|
| async with self._audit_context(request) as context: |
|
|
| async def send_wrapper(message: ASGISendEvent) -> None: |
| if self.audit_level == AuditLevel.REQUEST_RESPONSE: |
| await self._capture_response(message, context) |
|
|
| await send(message) |
|
|
| original_receive = receive |
|
|
| async def receive_wrapper() -> ASGIReceiveEvent: |
| nonlocal original_receive |
| message = await original_receive() |
|
|
| if self.audit_level in ( |
| AuditLevel.REQUEST, |
| AuditLevel.REQUEST_RESPONSE, |
| ): |
| await self._capture_request(message, context) |
|
|
| return message |
|
|
| await self.app(scope, receive_wrapper, send_wrapper) |
|
|
| @asynccontextmanager |
| async def _audit_context( |
| self, request: Request |
| ) -> AsyncGenerator[AuditContext, None]: |
| """ |
| async context manager that ensures that an audit log entry is recorded after the request is processed. |
| """ |
| context = AuditContext() |
| try: |
| yield context |
| finally: |
| await self._log_audit_entry(request, context) |
|
|
| async def _get_authenticated_user(self, request: Request) -> UserModel: |
|
|
| auth_header = request.headers.get("Authorization") |
| assert auth_header |
| user = get_current_user(request, None, get_http_authorization_cred(auth_header)) |
|
|
| return user |
|
|
| def _should_skip_auditing(self, request: Request) -> bool: |
| if ( |
| request.method not in {"POST", "PUT", "PATCH", "DELETE"} |
| or AUDIT_LOG_LEVEL == "NONE" |
| or not request.headers.get("authorization") |
| ): |
| return True |
| |
| pattern = re.compile( |
| r"^/api(?:/v1)?/(" + "|".join(self.excluded_paths) + r")\b" |
| ) |
| if pattern.match(request.url.path): |
| return True |
|
|
| return False |
|
|
| async def _capture_request(self, message: ASGIReceiveEvent, context: AuditContext): |
| if message["type"] == "http.request": |
| body = message.get("body", b"") |
| context.add_request_chunk(body) |
|
|
| async def _capture_response(self, message: ASGISendEvent, context: AuditContext): |
| if message["type"] == "http.response.start": |
| context.metadata["response_status_code"] = message["status"] |
|
|
| elif message["type"] == "http.response.body": |
| body = message.get("body", b"") |
| context.add_response_chunk(body) |
|
|
| async def _log_audit_entry(self, request: Request, context: AuditContext): |
| try: |
| user = await self._get_authenticated_user(request) |
|
|
| entry = AuditLogEntry( |
| id=str(uuid.uuid4()), |
| user=user.model_dump(include={"id", "name", "email", "role"}), |
| audit_level=self.audit_level.value, |
| verb=request.method, |
| request_uri=str(request.url), |
| response_status_code=context.metadata.get("response_status_code", None), |
| source_ip=request.client.host if request.client else None, |
| user_agent=request.headers.get("user-agent"), |
| request_object=context.request_body.decode("utf-8", errors="replace"), |
| response_object=context.response_body.decode("utf-8", errors="replace"), |
| ) |
|
|
| self.audit_logger.write(entry) |
| except Exception as e: |
| logger.error(f"Failed to log audit entry: {str(e)}") |
|
|
