py-app-runner 0.5.49.dev0__py3-none-any.whl

py_app_runner/request_handler/decorators.py

@@ -0,0 +1,304 @@
+import logging
+from collections.abc import Awaitable, Callable
+from functools import wraps
+from typing import Any, Concatenate, TypeAlias, TypeVar, cast
+
+from database_wrapper_pgsql import DBWrapperPgsqlAsync
+from typing_extensions import ParamSpec
+
+from py_app_runner.http_exception import HTTPException
+from py_app_runner.request_handler.handlers import RequestHandlerHelper
+from py_app_runner.throttle import Throttle
+
+rate_limit_logger = logging.getLogger(__name__ + ".rate_limit")
+audit_logger = logging.getLogger(__name__ + ".audit")
+
+P = ParamSpec("P")
+R = TypeVar("R")
+T = TypeVar("T", bound=type)
+SelfT = TypeVar("SelfT", bound=RequestHandlerHelper)
+
+AsyncMethod: TypeAlias = Callable[Concatenate[SelfT, P], Awaitable[R]]
+SyncMethod: TypeAlias = Callable[Concatenate[SelfT, P], R]
+
+# Receives (db_wrapper, *, event_type, user_id, details, ip, user_agent).
+AuditSink: TypeAlias = Callable[..., Awaitable[None]]
+
+
+def action(name: str) -> Callable[[SyncMethod], SyncMethod]:
+    def decorator(func: SyncMethod) -> SyncMethod:
+        # function objects don't have this attribute in typeshed -> tell the type checker
+        cast(object, func)._action_name = name  # type: ignore[attr-defined]
+
+        @wraps(func)
+        def wrapper(self: SelfT, *args: P.args, **kwargs: P.kwargs) -> R:  # pyrefly: ignore[not-a-type]
+            return func(self, *args, **kwargs)
+
+        cast(object, wrapper)._action_name = name  # type: ignore[attr-defined]
+        return wrapper
+
+    return decorator
+
+
+def authenticated(func: AsyncMethod) -> AsyncMethod:
+    @wraps(func)
+    async def wrapper(self: SelfT, *args: P.args, **kwargs: P.kwargs) -> R:  # pyrefly: ignore[not-a-type]
+        bh = getattr(self, "bridge_handler", None)
+        if bh is None or getattr(bh, "current_user", None) is None:
+            raise HTTPException("Not authorized", 4010, http_status=401)
+        return await func(self, *args, **kwargs)
+
+    # Marker so @require_auth_for_actions does not wrap an already-guarded method twice.
+    # functools.wraps copies __dict__ outwards, so decorators applied above this one
+    # (e.g. @action) carry the marker too.
+    cast(object, wrapper)._auth_wrapped = True  # type: ignore[attr-defined]
+    return wrapper
+
+
+def require_auth_for_actions(cls: T) -> T:
+    """
+    Class decorator: wraps every method that has `_action_name` with @authenticated.
+
+    Walks the full MRO, not just `vars(cls)`: dispatch in RequestHandlerHelper resolves
+    actions via `dir()`, so inherited @action methods are routable and must be guarded
+    as well.
+    """
+
+    for name in dir(cls):
+        attr = getattr(cls, name, None)
+        action_name = getattr(attr, "_action_name", None)
+        if action_name is None or not callable(attr):
+            continue
+
+        if getattr(attr, "_auth_wrapped", False):
+            continue
+
+        # Important: keep the marker so dispatch still works
+        wrapped = authenticated(attr)  # type: ignore[reportUnknownLambdaType]
+        cast(object, wrapped)._action_name = action_name  # type: ignore[attr-defined]
+        setattr(cls, name, wrapped)
+
+    return cls
+
+
+def with_tx(func: AsyncMethod) -> AsyncMethod:
+    @wraps(func)
+    async def wrapper(self: SelfT, *args: P.args, **kwargs: P.kwargs) -> R:  # pyrefly: ignore[not-a-type]
+        pg_conn = getattr(self, "pg_conn", None)
+        if pg_conn is None:
+            raise HTTPException("Database connection not linked (pg_conn missing)", 5000)
+
+        async with pg_conn.transaction():
+            return await func(self, *args, **kwargs)
+
+    return wrapper
+
+
+def with_db(func: AsyncMethod) -> AsyncMethod:
+    @wraps(func)
+    async def wrapper(self: SelfT, *args: P.args, **kwargs: P.kwargs) -> R:  # pyrefly: ignore[not-a-type]
+        bh = getattr(self, "bridge_handler", None)
+        if bh is None:
+            raise HTTPException("Bridge handler not linked", 5000)
+
+        async with bh.db_pools.main_db_pool as (pg_conn, pg_cur):
+            if not pg_conn or not pg_cur:
+                raise HTTPException("Failed to connect to database server", 5000)
+
+            self.pg_conn = pg_conn
+            self.pg_cur = pg_cur
+            self.db_wrapper = DBWrapperPgsqlAsync(pg_cur)
+
+            try:
+                return await func(self, *args, **kwargs)
+            finally:
+                for attr in ("pg_conn", "pg_cur", "db_wrapper"):
+                    if hasattr(self, attr):
+                        delattr(self, attr)
+
+    return wrapper
+
+
+def with_cache(func: AsyncMethod) -> AsyncMethod:
+    @wraps(func)
+    async def wrapper(self: SelfT, *args: P.args, **kwargs: P.kwargs) -> R:  # pyrefly: ignore[not-a-type]
+        bh = getattr(self, "bridge_handler", None)
+        if bh is None:
+            raise HTTPException("Bridge handler not linked", 5000)
+
+        async with bh.db_pools.cache_db_pool as redis_con:
+            if not redis_con:
+                raise HTTPException("Failed to connect to cache server", 5000)
+
+            self.redis_con = redis_con
+            try:
+                return await func(self, *args, **kwargs)
+            finally:
+                if hasattr(self, "redis_con"):
+                    delattr(self, "redis_con")
+
+    return wrapper
+
+
+def with_cache_and_db(func: AsyncMethod) -> AsyncMethod:
+    @wraps(func)
+    async def wrapper(self: SelfT, *args: P.args, **kwargs: P.kwargs) -> R:  # pyrefly: ignore[not-a-type]
+        bh = getattr(self, "bridge_handler", None)
+        if bh is None:
+            raise HTTPException("Bridge handler not linked", 5000)
+
+        async with bh.db_pools.cache_db_pool as redis_con:
+            if not redis_con:
+                raise HTTPException("Failed to connect to cache server", 5000)
+
+            self.redis_con = redis_con
+            try:
+                async with bh.db_pools.main_db_pool as (pg_conn, pg_cur):
+                    if not pg_conn or not pg_cur:
+                        raise HTTPException("Failed to connect to database server", 5000)
+
+                    self.pg_conn = pg_conn
+                    self.pg_cur = pg_cur
+                    self.db_wrapper = DBWrapperPgsqlAsync(pg_cur)
+
+                    try:
+                        return await func(self, *args, **kwargs)
+                    finally:
+                        for attr in ("pg_conn", "pg_cur", "db_wrapper"):
+                            if hasattr(self, attr):
+                                delattr(self, attr)
+            finally:
+                if hasattr(self, "redis_con"):
+                    delattr(self, "redis_con")
+
+    return wrapper
+
+
+def rate_limit(max_requests: int, window_seconds: int) -> Callable[[AsyncMethod], AsyncMethod]:
+    """
+    Redis-backed fixed-window rate limiter.
+    Key is derived from user ID (if authenticated) or client IP.
+    Must be applied AFTER @with_cache or @with_cache_and_db (so self.redis_con exists).
+    Returns HTTP 429 with Retry-After header when limit is exceeded.
+
+    Counting is done by py_app_runner.throttle, which is also callable directly for
+    anything that is not a bridge action - a login form keyed by email, say. Sharing the
+    implementation is what keeps the window semantics in one place; before, this decorator
+    was the only thing that could count, so every other caller invented its own.
+    """
+
+    def decorator(func: AsyncMethod) -> AsyncMethod:
+        @wraps(func)
+        async def wrapper(self: SelfT, *args: P.args, **kwargs: P.kwargs) -> R:  # pyrefly: ignore[not-a-type]
+            redis_con = getattr(self, "redis_con", None)
+            if redis_con is None:
+                rate_limit_logger.warning("rate_limit requires Redis; skipping enforcement")
+                return await func(self, *args, **kwargs)
+
+            bh = getattr(self, "bridge_handler", None)
+            user = getattr(bh, "current_user", None) if bh else None
+
+            if user and getattr(user, "id", None):
+                identity = f"user:{user.id}"
+            elif bh:
+                identity = f"ip:{bh.request.remote_ip}"
+            else:
+                identity = "unknown"
+
+            action_name = getattr(func, "_action_name", func.__name__)
+            # Scope by handler class as well: the same action name in two services
+            # would otherwise share a single bucket.
+            scope = f"{type(self).__module__}.{type(self).__qualname__}"
+            key = f"rl:{scope}:{action_name}:{identity}"
+
+            attempt = await Throttle(redis_con).hit(key, max_requests, window_seconds)
+
+            if not attempt.allowed:
+                raise HTTPException(
+                    f"Rate limit exceeded. Try again in {attempt.retry_after} seconds.",
+                    code=4029,
+                    http_status=429,
+                )
+
+            return await func(self, *args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+def audited(
+    event_type: str | None = None,
+    *,
+    sensitive_fields: tuple[str, ...] = ("password",),
+    sink: AuditSink | None = None,
+) -> Callable[[AsyncMethod], AsyncMethod]:
+    """
+    Audit-log decorator. Records after the wrapped action succeeds.
+
+    Place it as the innermost decorator, closest to the method body; with a `sink`
+    it must come AFTER @with_db in the stack (needs self.db_wrapper). The record
+    goes to `sink(db_wrapper, event_type=..., user_id=..., details=..., ip=...,
+    user_agent=...)`; without a sink it is written to the audit logger only.
+    If *event_type* is None the action name from @action is used.
+    """
+
+    def decorator(func: AsyncMethod) -> AsyncMethod:
+        @wraps(func)
+        async def wrapper(self: SelfT, *args: P.args, **kwargs: P.kwargs) -> R:  # pyrefly: ignore[not-a-type]
+            result = await func(self, *args, **kwargs)
+
+            # Best-effort: an audit failure must never fail the action itself.
+            try:
+                db_wrapper = getattr(self, "db_wrapper", None)
+                if sink is not None and db_wrapper is None:
+                    audit_logger.warning("@audited with a sink requires @with_db; skipping audit")
+                    return result
+
+                resolved_event = event_type or getattr(func, "_action_name", func.__name__)
+
+                bh = getattr(self, "bridge_handler", None)
+                user = getattr(bh, "current_user", None) if bh else None
+                user_id: int | None = getattr(user, "id", None)
+
+                ip: str | None = None
+                user_agent: str | None = None
+                if bh:
+                    ip = bh.request.remote_ip
+                    user_agent = bh.request.headers.get("User-Agent")
+
+                # Build details from input_data (first positional arg after self)
+                details = _sanitize_input(args[0] if args else {}, sensitive_fields)
+
+                if sink is not None:
+                    await sink(
+                        db_wrapper,
+                        event_type=resolved_event,
+                        user_id=user_id,
+                        details=details,
+                        ip=ip,
+                        user_agent=user_agent,
+                    )
+                else:
+                    audit_logger.info(
+                        "audit event=%s user_id=%s ip=%s details=%s",
+                        resolved_event,
+                        user_id,
+                        ip,
+                        details,
+                    )
+            except Exception:
+                audit_logger.exception("@audited failed for event_type=%s", event_type)
+
+            return result
+
+        return wrapper
+
+    return decorator
+
+
+def _sanitize_input(input_data: Any, sensitive_fields: tuple[str, ...]) -> dict[str, Any]:
+    """Strip sensitive fields from input_data before storing in audit details."""
+    if not isinstance(input_data, dict):
+        return {}
+    return {k: "***" if k in sensitive_fields else v for k, v in input_data.items()}