nexo-schemas 0.0.16__py3-none-any.whl

nexo/schemas/security/enums.py

@@ -0,0 +1,32 @@
+from enum import StrEnum
+from pydantic import BaseModel, Field
+from typing import Generic, TypeVar
+from nexo.types.string import ListOfStrs
+
+
+class Domain(StrEnum):
+    TENANT = "tenant"
+    SYSTEM = "system"
+
+    @classmethod
+    def choices(cls) -> ListOfStrs:
+        return [e.value for e in cls]
+
+
+DomainT = TypeVar("DomainT", bound=Domain)
+OptDomain = Domain | None
+OptDomainT = TypeVar("OptDomainT", bound=OptDomain)
+
+
+class DomainMixin(BaseModel, Generic[OptDomainT]):
+    domain: OptDomainT = Field(..., description="Domain")
+
+
+class RolePrefix(StrEnum):
+    MEDICAL = "medical"
+    TENANT = "tenant"
+    SYSTEM = "system"
+
+    @classmethod
+    def choices(cls) -> ListOfStrs:
+        return [e.value for e in cls]

nexo/schemas/security/impersonation.py

@@ -0,0 +1,179 @@
+from enum import StrEnum
+from fastapi import Header, Query
+from fastapi.requests import HTTPConnection, Request
+from fastapi.websockets import WebSocket
+from pydantic import BaseModel, Field
+from typing import (
+    Annotated,
+    Callable,
+    Generic,
+    Literal,
+    TypeVar,
+    overload,
+)
+from uuid import UUID
+from nexo.enums.connection import Header as HeaderEnum, Protocol, OptProtocol
+from nexo.types.string import ListOfStrs
+from nexo.types.uuid import OptUUID
+
+
+class Source(StrEnum):
+    HEADER = "header"
+    STATE = "state"
+
+    @classmethod
+    def choices(cls) -> ListOfStrs:
+        return [e.value for e in cls]
+
+
+class Impersonation(BaseModel):
+    user_id: Annotated[UUID, Field(..., description="User's ID")]
+    organization_id: Annotated[
+        OptUUID, Field(None, description="Organization's ID")
+    ] = None
+
+    @classmethod
+    def from_header(cls, conn: HTTPConnection) -> "Impersonation | None":
+        organization_id = conn.headers.get(HeaderEnum.X_ORGANIZATION_ID, None)
+        if organization_id is not None:
+            organization_id = UUID(organization_id)
+
+        user_id = conn.headers.get(HeaderEnum.X_USER_ID, None)
+        if user_id is not None:
+            user_id = UUID(user_id)
+
+            return cls(
+                user_id=user_id,
+                organization_id=organization_id,
+            )
+
+        return None
+
+    @classmethod
+    def from_query(cls, conn: HTTPConnection) -> "Impersonation | None":
+        organization_id = conn.query_params.get(
+            HeaderEnum.X_ORGANIZATION_ID.value, None
+        )
+        if organization_id is not None:
+            organization_id = UUID(organization_id)
+
+        user_id = conn.query_params.get(HeaderEnum.X_USER_ID.value, None)
+        if user_id is not None:
+            user_id = UUID(user_id)
+
+            return cls(
+                user_id=user_id,
+                organization_id=organization_id,
+            )
+
+        return None
+
+    @classmethod
+    def extract(cls, conn: HTTPConnection) -> "Impersonation | None":
+        impersonation = getattr(conn.state, "impersonation", None)
+        if isinstance(impersonation, Impersonation):
+            return impersonation
+
+        impersonation = cls.from_header(conn)
+        if impersonation is not None:
+            return impersonation
+
+        impersonation = cls.from_query(conn)
+        if impersonation is not None:
+            return impersonation
+
+        return None
+
+    @overload
+    @classmethod
+    def as_dependency(
+        cls,
+        protocol: None = None,
+        /,
+    ) -> Callable[[HTTPConnection, OptUUID, OptUUID], "Impersonation | None"]: ...
+    @overload
+    @classmethod
+    def as_dependency(
+        cls,
+        protocol: Literal[Protocol.HTTP],
+        /,
+    ) -> Callable[[Request, OptUUID, OptUUID], "Impersonation | None"]: ...
+    @overload
+    @classmethod
+    def as_dependency(
+        cls,
+        protocol: Literal[Protocol.WEBSOCKET],
+        /,
+    ) -> Callable[[WebSocket, OptUUID, OptUUID], "Impersonation | None"]: ...
+    @classmethod
+    def as_dependency(
+        cls,
+        protocol: OptProtocol = None,
+        /,
+    ) -> (
+        Callable[[HTTPConnection, OptUUID, OptUUID], "Impersonation | None"]
+        | Callable[[Request, OptUUID, OptUUID], "Impersonation | None"]
+        | Callable[[WebSocket, OptUUID, OptUUID], "Impersonation | None"]
+    ):
+        def _dependency(
+            conn: HTTPConnection,
+            # These are for documentation purpose only
+            _user_id: OptUUID = Header(
+                None,
+                alias=HeaderEnum.X_USER_ID.value,
+                description="User's ID",
+            ),
+            _organization_id: OptUUID = Header(
+                None,
+                alias=HeaderEnum.X_ORGANIZATION_ID.value,
+                description="Organization's ID",
+            ),
+        ) -> "Impersonation | None":
+            return cls.extract(conn)
+
+        def _request_dependency(
+            request: Request,
+            # These are for documentation purpose only
+            _user_id: OptUUID = Header(
+                None,
+                alias=HeaderEnum.X_USER_ID.value,
+                description="User's ID",
+            ),
+            _organization_id: OptUUID = Header(
+                None,
+                alias=HeaderEnum.X_ORGANIZATION_ID.value,
+                description="Organization's ID",
+            ),
+        ) -> "Impersonation | None":
+            return cls.extract(request)
+
+        def _websocket_dependency(
+            websocket: WebSocket,
+            # These are for documentation purpose only
+            _user_id: OptUUID = Query(
+                None,
+                alias=HeaderEnum.X_USER_ID.value,
+                description="User's ID",
+            ),
+            _organization_id: OptUUID = Query(
+                None,
+                alias=HeaderEnum.X_ORGANIZATION_ID.value,
+                description="Organization's ID",
+            ),
+        ) -> "Impersonation | None":
+            return cls.extract(websocket)
+
+        if protocol is None:
+            return _dependency
+        elif protocol is Protocol.HTTP:
+            return _request_dependency
+        elif protocol is Protocol.WEBSOCKET:
+            return _websocket_dependency
+
+
+OptImpersonation = Impersonation | None
+OptImpersonationT = TypeVar("OptImpersonationT", bound=OptImpersonation)
+
+
+class ImpersonationMixin(BaseModel, Generic[OptImpersonationT]):
+    impersonation: OptImpersonationT = Field(..., description="Impersonation")

nexo/schemas/security/token.py

@@ -0,0 +1,402 @@
+from collections.abc import Iterable, Sequence
+from Crypto.PublicKey.RSA import RsaKey
+from datetime import datetime, timedelta, timezone
+from pydantic import BaseModel, Field, ValidationError, model_validator
+from typing import (
+    Annotated,
+    Generic,
+    Literal,
+    Self,
+    TypeGuard,
+    TypeVar,
+    overload,
+)
+from uuid import UUID
+from nexo.crypto.token import decode, encode
+from nexo.enums.expiration import Expiration
+from nexo.types.datetime import OptDatetime
+from nexo.types.integer import OptInt, DoubleInts
+from nexo.types.misc import BytesOrStr
+from nexo.types.string import OptListOfStrs, OptStr
+from nexo.types.uuid import OptUUID
+from .enums import Domain, OptDomain, DomainT
+
+
+class TokenV1(BaseModel):
+    iss: Annotated[OptStr, Field(None, description="Issuer")] = None
+    sub: Annotated[str, Field(..., description="Subject")]
+    sr: Annotated[str, Field(..., description="System role")]
+    u_i: Annotated[int, Field(..., description="User's ID")]
+    u_uu: Annotated[UUID, Field(..., description="User's UUID")]
+    u_u: Annotated[str, Field(..., description="User's Username")]
+    u_e: Annotated[str, Field(..., description="User's Email")]
+    u_ut: Annotated[str, Field(..., description="User's type")]
+    o_i: Annotated[OptInt, Field(None, description="Organization's ID")] = None
+    o_uu: Annotated[OptUUID, Field(None, description="Organization's UUID")] = None
+    o_k: Annotated[OptStr, Field(None, description="Organization's Key")] = None
+    o_ot: Annotated[OptStr, Field(None, description="Organization's type")] = None
+    uor: Annotated[
+        OptListOfStrs,
+        Field(None, description="User's organization role", min_length=1),
+    ] = None
+    iat_dt: Annotated[datetime, Field(..., description="Issued At Timestamp")]
+    iat: Annotated[int, Field(..., description="Issued at")]
+    exp_dt: Annotated[datetime, Field(..., description="Expired At Timestamp")]
+    exp: Annotated[int, Field(..., description="Expired at")]
+
+    @classmethod
+    def from_string(
+        cls,
+        token: str,
+        *,
+        key: BytesOrStr | RsaKey,
+        audience: str | Iterable[str] | None = None,
+        subject: OptStr = None,
+        issuer: str | Sequence[str] | None = None,
+        leeway: float | timedelta = 0,
+    ) -> "TokenV1":
+        obj = decode(
+            token,
+            key=key,
+            audience=audience,
+            subject=subject,
+            issuer=issuer,
+            leeway=leeway,
+        )
+        return cls.model_validate(obj)
+
+    @overload
+    def to_string(
+        self,
+        key: RsaKey,
+    ) -> str: ...
+    @overload
+    def to_string(
+        self,
+        key: BytesOrStr,
+        *,
+        password: OptStr = None,
+    ) -> str: ...
+    @overload
+    def to_string(
+        self,
+        key: BytesOrStr | RsaKey,
+        *,
+        password: OptStr = None,
+    ) -> str: ...
+    def to_string(
+        self,
+        key: BytesOrStr | RsaKey,
+        *,
+        password: OptStr = None,
+    ) -> str:
+        if isinstance(key, RsaKey):
+            return encode(
+                payload=self.model_dump(mode="json", exclude_none=True),
+                key=key,
+            )
+        else:
+            return encode(
+                payload=self.model_dump(mode="json", exclude_none=True),
+                key=key,
+                password=password,
+            )
+
+
+class Claim(BaseModel):
+    iss: Annotated[OptStr, Field(None, description="Issuer")] = None
+    sub: Annotated[UUID, Field(..., description="Subject")]
+    aud: Annotated[OptStr, Field(None, description="Audience")] = None
+    exp: Annotated[int, Field(..., description="Expired at")]
+    iat: Annotated[int, Field(..., description="Issued at")]
+
+    @classmethod
+    def new_timestamp(
+        cls, iat_dt: OptDatetime = None, exp_in: Expiration = Expiration.EXP_15MN
+    ) -> DoubleInts:
+        if iat_dt is None:
+            iat_dt = datetime.now(tz=timezone.utc)
+        exp_dt = iat_dt + timedelta(seconds=exp_in.value)
+        return int(iat_dt.timestamp()), int(exp_dt.timestamp())
+
+
+OrganizationT = TypeVar("OrganizationT", bound=OptUUID)
+
+
+class Credential(BaseModel, Generic[DomainT, OrganizationT]):
+    d: DomainT = Field(..., description="Domain")
+    o: OrganizationT = Field(..., description="Organization")
+
+
+class GenericToken(
+    Credential[DomainT, OrganizationT],
+    Claim,
+    Generic[DomainT, OrganizationT],
+):
+    @classmethod
+    def from_string(
+        cls,
+        token: str,
+        *,
+        key: BytesOrStr | RsaKey,
+        audience: str | Iterable[str] | None = None,
+        subject: OptStr = None,
+        issuer: str | Sequence[str] | None = None,
+        leeway: float | timedelta = 0,
+    ) -> Self:
+        obj = decode(
+            token,
+            key=key,
+            audience=audience,
+            subject=subject,
+            issuer=issuer,
+            leeway=leeway,
+        )
+        return cls.model_validate(obj)
+
+    @model_validator(mode="after")
+    def validate_credential(self) -> Self:
+        return self
+
+    @overload
+    def to_string(
+        self,
+        key: RsaKey,
+    ) -> str: ...
+    @overload
+    def to_string(
+        self,
+        key: BytesOrStr,
+        *,
+        password: OptStr = None,
+    ) -> str: ...
+    @overload
+    def to_string(
+        self,
+        key: BytesOrStr | RsaKey,
+        *,
+        password: OptStr = None,
+    ) -> str: ...
+    def to_string(
+        self,
+        key: BytesOrStr | RsaKey,
+        *,
+        password: OptStr = None,
+    ) -> str:
+        if isinstance(key, RsaKey):
+            return encode(
+                payload=self.model_dump(mode="json", exclude_none=True),
+                key=key,
+            )
+        else:
+            return encode(
+                payload=self.model_dump(mode="json", exclude_none=True),
+                key=key,
+                password=password,
+            )
+
+
+class TenantToken(GenericToken[Literal[Domain.TENANT], UUID]):
+    d: Annotated[Literal[Domain.TENANT], Field(Domain.TENANT, description="Domain")] = (
+        Domain.TENANT
+    )
+
+    @model_validator(mode="after")
+    def validate_identity(self) -> Self:
+        if self.d is not Domain.TENANT:
+            raise ValueError(f"Value of 'd' claim must be {Domain.TENANT}")
+        if not isinstance(self.o, UUID):
+            raise ValueError(f"Value of 'o' claim must be an UUID. Value: {self.o}")
+        return self
+
+    @classmethod
+    def new(
+        cls,
+        *,
+        sub: UUID,
+        o: UUID,
+        iss: OptStr = None,
+        aud: OptStr = None,
+        iat_dt: OptDatetime = None,
+        exp_in: Expiration = Expiration.EXP_15MN,
+    ) -> "TenantToken":
+        iat, exp = cls.new_timestamp(iat_dt, exp_in)
+        return cls(iss=iss, sub=sub, aud=aud, exp=exp, iat=iat, o=o)
+
+
+class SystemToken(GenericToken[Literal[Domain.SYSTEM], None]):
+    d: Annotated[Literal[Domain.SYSTEM], Field(Domain.SYSTEM, description="Domain")] = (
+        Domain.SYSTEM
+    )
+    o: None = None
+
+    @model_validator(mode="after")
+    def validate_identity(self) -> Self:
+        if self.d is not Domain.SYSTEM:
+            raise ValueError(f"Value of 'd' claim must be {Domain.SYSTEM}")
+        if self.o is not None:
+            raise ValueError(f"Value of 'o' claim must be None. Value: {self.o}")
+        return self
+
+    @classmethod
+    def new(
+        cls,
+        *,
+        sub: UUID,
+        iss: OptStr = None,
+        aud: OptStr = None,
+        iat_dt: OptDatetime = None,
+        exp_in: Expiration = Expiration.EXP_15MN,
+    ) -> "SystemToken":
+        iat, exp = cls.new_timestamp(iat_dt, exp_in)
+        return cls(iss=iss, sub=sub, aud=aud, exp=exp, iat=iat)
+
+
+AnyToken = TenantToken | SystemToken
+AnyTokenT = TypeVar("AnyTokenT", bound=AnyToken)
+OptAnyToken = AnyToken | None
+OptAnyTokenT = TypeVar("OptAnyTokenT", bound=OptAnyToken)
+
+
+class TokenMixin(BaseModel, Generic[OptAnyTokenT]):
+    token: OptAnyTokenT = Field(..., description="Token")
+
+
+def is_tenant_token(token: AnyToken) -> TypeGuard[TenantToken]:
+    return (
+        isinstance(token, TenantToken)
+        and token.d is Domain.TENANT
+        and token.o is not None
+    )
+
+
+def is_system_token(token: AnyToken) -> TypeGuard[SystemToken]:
+    return (
+        isinstance(token, SystemToken) and token.d is Domain.SYSTEM and token.o is None
+    )
+
+
+class TokenFactory:
+    @overload
+    @staticmethod
+    def from_string(
+        token: str,
+        domain: Literal[Domain.TENANT],
+        *,
+        key: BytesOrStr | RsaKey,
+        audience: str | Iterable[str] | None = None,
+        subject: OptStr = None,
+        issuer: str | Sequence[str] | None = None,
+        leeway: float | timedelta = 0,
+    ) -> TenantToken: ...
+    @overload
+    @staticmethod
+    def from_string(
+        token: str,
+        domain: Literal[Domain.SYSTEM],
+        *,
+        key: BytesOrStr | RsaKey,
+        audience: str | Iterable[str] | None = None,
+        subject: OptStr = None,
+        issuer: str | Sequence[str] | None = None,
+        leeway: float | timedelta = 0,
+    ) -> SystemToken: ...
+    @overload
+    @staticmethod
+    def from_string(
+        token: str,
+        domain: None = None,
+        *,
+        key: BytesOrStr | RsaKey,
+        audience: str | Iterable[str] | None = None,
+        subject: OptStr = None,
+        issuer: str | Sequence[str] | None = None,
+        leeway: float | timedelta = 0,
+    ) -> AnyToken: ...
+    @staticmethod
+    def from_string(
+        token: str,
+        domain: OptDomain = None,
+        *,
+        key: BytesOrStr | RsaKey,
+        audience: str | Iterable[str] | None = None,
+        subject: OptStr = None,
+        issuer: str | Sequence[str] | None = None,
+        leeway: float | timedelta = 0,
+    ) -> AnyToken:
+        validated_token = None
+        models = (TokenV1, TenantToken, SystemToken)
+        for model in models:
+            try:
+                validated_token = model.from_string(
+                    token,
+                    key=key,
+                    audience=audience,
+                    subject=subject,
+                    issuer=issuer,
+                    leeway=leeway,
+                )
+            except ValidationError:
+                continue
+        if validated_token is None:
+            raise ValueError("Unable to validate raw token into known token model")
+
+        if isinstance(validated_token, (TenantToken, SystemToken)):
+            result_token = validated_token
+        else:
+            if validated_token.sr == "administrator":
+                if (
+                    validated_token.o_i is not None
+                    or validated_token.o_uu is not None
+                    or validated_token.o_k is not None
+                    or validated_token.o_ot is not None
+                    or validated_token.uor is not None
+                ):
+                    raise ValueError(
+                        "All organization-related claims must be None for System token"
+                    )
+                result_token = SystemToken(
+                    iss=validated_token.iss,
+                    sub=validated_token.u_uu,
+                    aud=None,
+                    exp=validated_token.exp,
+                    iat=validated_token.iat,
+                )
+            elif validated_token.sr == "user":
+                if (
+                    validated_token.o_i is None
+                    or validated_token.o_uu is None
+                    or validated_token.o_k is None
+                    or validated_token.o_ot is None
+                    or validated_token.uor is None
+                ):
+                    raise ValueError(
+                        "All organization-related claims can not be None for Tenant Token"
+                    )
+                result_token = TenantToken(
+                    iss=validated_token.iss,
+                    sub=validated_token.u_uu,
+                    aud=None,
+                    exp=validated_token.exp,
+                    iat=validated_token.iat,
+                    o=validated_token.o_uu,
+                )
+            else:
+                raise ValueError(
+                    f"Claim 'sr' can only be either 'administrator' or 'user' but received {validated_token.sr}"
+                )
+
+        if domain is None:
+            return result_token
+        elif domain is Domain.TENANT:
+            if not is_tenant_token(result_token):
+                raise ValueError(
+                    "Failed parsing Tenant Token from string, raw token did not qualify as Tenant Token"
+                )
+            return result_token
+        elif domain is Domain.SYSTEM:
+            if not is_system_token(result_token):
+                raise ValueError(
+                    "Failed parsing System token from string, raw token did not qualify as System Token"
+                )
+            return result_token

nexo/schemas/security/types.py

@@ -0,0 +1,17 @@
+from typing import TypeVar
+from nexo.enums.organization import ListOfOrganizationRoles, SeqOfOrganizationRoles
+from nexo.enums.system import ListOfSystemRoles, SeqOfSystemRoles
+
+
+ListOfDomainRoles = ListOfOrganizationRoles | ListOfSystemRoles
+ListOfDomainRolesT = TypeVar("ListOfDomainRolesT", bound=ListOfDomainRoles)
+
+OptListOfDomainRoles = ListOfDomainRoles | None
+OptListOfDomainRolesT = TypeVar("OptListOfDomainRolesT", bound=OptListOfDomainRoles)
+
+
+SeqOfDomainRoles = SeqOfOrganizationRoles | SeqOfSystemRoles
+SeqOfDomainRolesT = TypeVar("SeqOfDomainRolesT", bound=SeqOfDomainRoles)
+
+OptSeqOfDomainRoles = SeqOfDomainRoles | None
+OptSeqOfDomainRolesT = TypeVar("OptSeqOfDomainRolesT", bound=OptSeqOfDomainRoles)