File: guards.py

package info (click to toggle)
litestar 2.19.0-2
  • links: PTS, VCS
  • area: main
  • in suites: sid
  • size: 12,500 kB
  • sloc: python: 70,169; makefile: 254; javascript: 105; sh: 60
file content (61 lines) | stat: -rw-r--r-- 1,582 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
from enum import Enum
from os import environ

from pydantic import UUID4, BaseModel

from litestar import Controller, Litestar, Router, get, post
from litestar.connection import ASGIConnection
from litestar.exceptions import NotAuthorizedException
from litestar.handlers.base import BaseRouteHandler


class UserRole(str, Enum):
    CONSUMER = "consumer"
    ADMIN = "admin"


class User(BaseModel):
    id: UUID4
    role: UserRole

    @property
    def is_admin(self) -> bool:
        """Determines whether the user is an admin user"""
        return self.role == UserRole.ADMIN


def admin_user_guard(connection: ASGIConnection, _: BaseRouteHandler) -> None:
    if not connection.user.is_admin:
        raise NotAuthorizedException()


@post(path="/user", guards=[admin_user_guard])
def create_user(data: User) -> User: ...


def my_guard(connection: ASGIConnection, handler: BaseRouteHandler) -> None: ...


# controller
class UserController(Controller):
    path = "/user"
    guards = [my_guard]


# router
admin_router = Router(path="admin", route_handlers=[UserController], guards=[my_guard])

# app
app = Litestar(route_handlers=[admin_router], guards=[my_guard])


def secret_token_guard(connection: ASGIConnection, route_handler: BaseRouteHandler) -> None:
    if (
        route_handler.opt.get("secret")
        and not connection.headers.get("Secret-Header", "") == route_handler.opt["secret"]
    ):
        raise NotAuthorizedException()


@get(path="/secret", guards=[secret_token_guard], opt={"secret": environ.get("SECRET")})
def secret_endpoint() -> None: ...