scale-agentex/agentex/src/adapters/authorization/adapter_agentex_authz_proxy.py at d7448dfa3157c5e6ba1e3d456b3667f53caedcda · scaleapi/scale-agentex · GitHub

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
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
from collections.abc import Iterable
from functools import lru_cache
from typing import Annotated

from fastapi import Depends
from src.adapters.authorization.port import (
    AuthorizationGateway,
)
from src.api.schemas.authorization_types import (
    AgentexResource,
    AgentexResourceType,
    AuthorizedOperationType,
)
from src.api.schemas.principal_context import AgentexAuthPrincipalContext
from src.config.dependencies import DEnvironmentVariable
from src.config.environment_variables import EnvVarKeys
from src.utils.http_request_handler import HttpRequestHandler


class AgentexAuthorizationProxy(AuthorizationGateway[AgentexAuthPrincipalContext]):
    def __init__(
        self,
        agentex_auth_url: DEnvironmentVariable(EnvVarKeys.AGENTEX_AUTH_URL),
    ):
        self.agentex_auth_url = agentex_auth_url

    async def grant(
        self,
        principal: AgentexAuthPrincipalContext,
        resource: AgentexResource,
        operation: AuthorizedOperationType,
    ) -> None:
        payload = {
            "principal": principal,
            "resource": resource.model_dump(),
            "operation": operation,
        }
        await HttpRequestHandler.post_with_error_handling(
            self.agentex_auth_url, "/v1/authz/grant", json=payload
        )

    async def revoke(
        self,
        principal: AgentexAuthPrincipalContext,
        resource: AgentexResource,
        operation: AuthorizedOperationType,
    ) -> None:
        payload = {
            "principal": principal,
            "resource": resource.model_dump(),
            "operation": operation,
        }
        await HttpRequestHandler.post_with_error_handling(
            self.agentex_auth_url, "/v1/authz/revoke", json=payload
        )

    async def check(
        self,
        principal: AgentexAuthPrincipalContext,
        resource: AgentexResource,
        operation: AuthorizedOperationType,
    ) -> bool:
        payload = {
            "principal": principal,
            "resource": resource.model_dump(),
            "operation": operation,
        }
        await HttpRequestHandler.post_with_error_handling(
            self.agentex_auth_url, "/v1/authz/check", json=payload
        )
        return True  # request was successful

    async def list_resources(
        self,
        principal: AgentexAuthPrincipalContext,
        filter_resource: AgentexResourceType,
        filter_operation: AuthorizedOperationType = AuthorizedOperationType.read,
    ) -> Iterable[str]:
        payload = {
            "principal": principal,
            "filter_resource": filter_resource,
            "filter_operation": filter_operation,
        }
        response = await HttpRequestHandler.post_with_error_handling(
            self.agentex_auth_url, "/v1/authz/search", json=payload
        )
        return response["items"]


@lru_cache(maxsize=1)
def _get_cached_agentex_authorization() -> AgentexAuthorizationProxy:
    """Cached AgentexAuthorizationProxy instance."""
    from src.config.dependencies import resolve_environment_variable_dependency

    url = resolve_environment_variable_dependency(EnvVarKeys.AGENTEX_AUTH_URL)
    return AgentexAuthorizationProxy(agentex_auth_url=url)


def _get_agentex_authorization() -> AgentexAuthorizationProxy:
    return _get_cached_agentex_authorization()


DAgentexAuthorization = Annotated[
    AgentexAuthorizationProxy, Depends(_get_agentex_authorization)
]