1 files changed, 16 insertions, 2 deletions

diff --git a/db/models.py b/db/models.py
index 7e0afb0..8a01c8c 100644
--- a/db/models.py
+++ b/db/models.py
@@ -1,10 +1,13 @@
+import re
 from datetime import datetime, UTC
+from urllib.parse import urlsplit
 from uuid import uuid4, UUID, uuid7
 
 from sqlalchemy import DateTime
 from sqlmodel import SQLModel, Field, Relationship
-from authlib.integrations.sqla_oauth2 import OAuth2ClientMixin, OAuth2AuthorizationCodeMixin, OAuth2TokenMixin
+from starlette.authentication import BaseUser
 
+from authentication.mixins import OAuth2ClientMixin, OAuth2AuthorizationCodeMixin, OAuth2TokenMixin
 from authentication.utils import bcrypt_sha256_hash
 
 
@@ -24,7 +27,7 @@ class BaseModel(SQLModel):
     )
 
 
-class User(BaseModel, table=True):
+class User(BaseModel, BaseUser, table=True):
     username: str = Field(unique=True, index=True)
     password: str = Field(nullable=False)
 
@@ -33,6 +36,9 @@ class User(BaseModel, table=True):
     def set_password(self, value):
         self.password = bcrypt_sha256_hash(value)
 
+    def is_authenticated(self) -> bool:
+        return True
+
 class Session(BaseModel, table=True):
     user_id: UUID = Field(foreign_key='user.id', exclude=True)
     user: User = Relationship()
@@ -40,8 +46,16 @@ class Session(BaseModel, table=True):
 class ClientApplication(BaseModel, OAuth2ClientMixin, table=True):
     name: str
 
+    def check_redirect_uri(self, redirect_uri: str) -> bool:
+        return any(
+            re.fullmatch(re.escape(str(pattern)).replace(r"\*", r"\d+"), redirect_uri)
+            for pattern in self.redirect_uris
+        )
+
 class AuthCode(BaseModel, OAuth2AuthorizationCodeMixin, table=True):
     user_id: UUID = Field(foreign_key="user.id", ondelete="CASCADE")
+    user: User = Relationship()
 
 class AuthToken(BaseModel, OAuth2TokenMixin, table=True):
     user_id: UUID = Field(foreign_key="user.id", ondelete="CASCADE")
+    user: User = Relationship()