codekingpro/portable-devtools

1from __future__ import annotations
2
3from abc import abstractmethod
4from enum import Enum
5from typing import (
6    Any,
7    List,
8    Optional,
9    Dict,
10    Tuple,
11    TypeVar,
12)
13from dataclasses import dataclass
14
15from pydantic import SecretStr
16
17from chromadb.config import (
18    Component,
19    System,
20)
21
22T = TypeVar("T")
23S = TypeVar("S")
24
25
26class AuthError(Exception):
27    pass
28
29
30ClientAuthHeaders = Dict[str, SecretStr]
31
32
33class ClientAuthProvider(Component):
34    """
35    ClientAuthProvider is responsible for providing authentication headers for
36    client requests. Client implementations (in our case, just the FastAPI
37    client) must inject these headers into their requests.
38    """
39
40    def __init__(self, system: System) -> None:
41        super().__init__(system)
42
43    @abstractmethod
44    def authenticate(self) -> ClientAuthHeaders:
45        pass
46
47
48@dataclass
49class UserIdentity:
50    """
51    UserIdentity represents the identity of a user. In general, not all fields
52    will be populated, and the fields that are populated will depend on the
53    authentication provider.
54
55    The idea is that the AuthenticationProvider is responsible for populating
56    _all_ information known about the user, and the AuthorizationProvider is
57    responsible for making decisions based on that information.
58    """
59
60    user_id: str
61    tenant: Optional[str] = None
62    databases: Optional[List[str]] = None
63    # This can be used for any additional auth context which needs to be
64    # propagated from the authentication provider to the authorization
65    # provider.
66    attributes: Optional[Dict[str, Any]] = None
67
68
69class ServerAuthenticationProvider(Component):
70    """
71    ServerAuthenticationProvider is responsible for authenticating requests. If
72    a ServerAuthenticationProvider is configured, it will be called by the
73    server to authenticate requests. If no ServerAuthenticationProvider is
74    configured, all requests will be authenticated.
75
76    The ServerAuthenticationProvider should return a UserIdentity object if the
77    request is authenticated for use by the ServerAuthorizationProvider.
78    """
79
80    def __init__(self, system: System) -> None:
81        super().__init__(system)
82        self._ignore_auth_paths: Dict[
83            str, List[str]
84        ] = system.settings.chroma_server_auth_ignore_paths
85        self.overwrite_singleton_tenant_database_access_from_auth = (
86            system.settings.chroma_overwrite_singleton_tenant_database_access_from_auth
87        )
88
89    @abstractmethod
90    def authenticate_or_raise(self, headers: Dict[str, str]) -> UserIdentity:
91        pass
92
93    def ignore_operation(self, verb: str, path: str) -> bool:
94        if (
95            path in self._ignore_auth_paths.keys()
96            and verb.upper() in self._ignore_auth_paths[path]
97        ):
98            return True
99        return False
100
101    def read_creds_or_creds_file(self) -> List[str]:
102        _creds_file = None
103        _creds = None
104
105        if self._system.settings.chroma_server_authn_credentials_file:
106            _creds_file = str(
107                self._system.settings["chroma_server_authn_credentials_file"]
108            )
109        if self._system.settings.chroma_server_authn_credentials:
110            _creds = str(self._system.settings["chroma_server_authn_credentials"])
111        if not _creds_file and not _creds:
112            raise ValueError(
113                "No credentials file or credentials found in "
114                "[chroma_server_authn_credentials]."
115            )
116        if _creds_file and _creds:
117            raise ValueError(
118                "Both credentials file and credentials found."
119                "Please provide only one."
120            )
121        if _creds:
122            return [c for c in _creds.split("\n") if c]
123        elif _creds_file:
124            with open(_creds_file, "r") as f:
125                return f.readlines()
126        raise ValueError("Should never happen")
127
128    def singleton_tenant_database_if_applicable(
129        self, user: Optional[UserIdentity]
130    ) -> Tuple[Optional[str], Optional[str]]:
131        """
132        If settings.chroma_overwrite_singleton_tenant_database_access_from_auth
133        is False, this function always returns (None, None).
134
135        If settings.chroma_overwrite_singleton_tenant_database_access_from_auth
136        is True, follows the following logic:
137        - If the user only has access to a single tenant, this function will
138          return that tenant as its first return value.
139        - If the user only has access to a single database, this function will
140          return that database as its second return value. If the user has
141          access to multiple tenants and/or databases, including "*", this
142          function will return None for the corresponding value(s).
143        - If the user has access to multiple tenants and/or databases this
144          function will return None for the corresponding value(s).
145        """
146        if not self.overwrite_singleton_tenant_database_access_from_auth or not user:
147            return None, None
148        tenant = None
149        database = None
150        if user.tenant and user.tenant != "*":
151            tenant = user.tenant
152        if user.databases and len(user.databases) == 1 and user.databases[0] != "*":
153            database = user.databases[0]
154        return tenant, database
155
156
157class AuthzAction(str, Enum):
158    """
159    The set of actions that can be authorized by the authorization provider.
160    """
161
162    RESET = "system:reset"
163    CREATE_TENANT = "tenant:create_tenant"
164    GET_TENANT = "tenant:get_tenant"
165    CREATE_DATABASE = "db:create_database"
166    GET_DATABASE = "db:get_database"
167    DELETE_DATABASE = "db:delete_database"
168    LIST_DATABASES = "db:list_databases"
169    LIST_COLLECTIONS = "db:list_collections"
170    COUNT_COLLECTIONS = "db:count_collections"
171    CREATE_COLLECTION = "db:create_collection"
172    GET_OR_CREATE_COLLECTION = "db:get_or_create_collection"
173    GET_COLLECTION = "collection:get_collection"
174    DELETE_COLLECTION = "collection:delete_collection"
175    UPDATE_COLLECTION = "collection:update_collection"
176    ADD = "collection:add"
177    DELETE = "collection:delete"
178    GET = "collection:get"
179    QUERY = "collection:query"
180    COUNT = "collection:count"
181    UPDATE = "collection:update"
182    UPSERT = "collection:upsert"
183
184
185@dataclass
186class AuthzResource:
187    """
188    The resource being accessed in an authorization request.
189    """
190
191    tenant: Optional[str]
192    database: Optional[str]
193    collection: Optional[str]
194
195
196class ServerAuthorizationProvider(Component):
197    """
198    ServerAuthorizationProvider is responsible for authorizing requests. If a
199    ServerAuthorizationProvider is configured, it will be called by the server
200    to authorize requests. If no ServerAuthorizationProvider is configured, all
201    requests will be authorized.
202
203    ServerAuthorizationProvider should raise an exception if the request is not
204    authorized.
205    """
206
207    def __init__(self, system: System) -> None:
208        super().__init__(system)
209
210    @abstractmethod
211    def authorize_or_raise(
212        self, user: UserIdentity, action: AuthzAction, resource: AuthzResource
213    ) -> None:
214        pass
215
216    def read_config_or_config_file(self) -> List[str]:
217        _config_file = None
218        _config = None
219        if self._system.settings.chroma_server_authz_config_file:
220            _config_file = self._system.settings["chroma_server_authz_config_file"]
221        if self._system.settings.chroma_server_authz_config:
222            _config = str(self._system.settings["chroma_server_authz_config"])
223        if not _config_file and not _config:
224            raise ValueError(
225                "No authz configuration file or authz configuration found."
226            )
227        if _config_file and _config:
228            raise ValueError(
229                "Both authz configuration file and authz configuration found."
230                "Please provide only one."
231            )
232        if _config:
233            return [c for c in _config.split("\n") if c]
234        elif _config_file:
235            with open(_config_file, "r") as f:
236                return f.readlines()
237        raise ValueError("Should never happen")
238