squirrels 0.5.0b4__py3-none-any.whl → 0.5.1__py3-none-any.whl

squirrels/_auth.py

@@ -1,57 +1,52 @@
-from typing import Callable
+from typing import Callable, Any
 from datetime import datetime, timedelta, timezone
-from enum import Enum
 from functools import cached_property
 from jwt.exceptions import InvalidTokenError
 from passlib.context import CryptContext
 from pydantic import ValidationError
-from pydantic_core import PydanticUndefined
 from sqlalchemy import create_engine, Engine, func, inspect, text, ForeignKey
-from sqlalchemy import Column, String, Integer, Float, Boolean
 from sqlalchemy.orm import declarative_base, sessionmaker, Mapped, mapped_column
-import jwt, types, typing as _t, uuid, secrets, json
+import jwt, uuid, secrets, json
 
 from ._manifest import PermissionScope
 from ._py_module import PyModule
 from ._exceptions import InvalidInputError, ConfigurationError
 from ._arguments.init_time_args import AuthProviderArgs
 from ._schemas.auth_models import (
-    BaseUser, ApiKey, UserField, AuthProvider, ProviderConfigs, ClientRegistrationRequest, ClientUpdateRequest,
-    ClientDetailsResponse, ClientRegistrationResponse, ClientUpdateResponse, TokenResponse
+    CustomUserFields, AbstractUser, GuestUser, RegisteredUser, ApiKey, UserField, AuthProvider, ProviderConfigs, 
+    ClientRegistrationRequest, ClientUpdateRequest, ClientDetailsResponse, ClientRegistrationResponse, 
+    ClientUpdateResponse, TokenResponse
 )
 from . import _utils as u, _constants as c
 
 pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
 
-reserved_fields = ["username", "is_admin"]
-disallowed_fields = ["password", "password_hash", "created_at", "token_id", "exp"]
-
-User = _t.TypeVar('User', bound=BaseUser)
-
 ProviderFunctionType = Callable[[AuthProviderArgs], AuthProvider]
 
 
-class Authenticator(_t.Generic[User]):
+class Authenticator:
     providers: list[ProviderFunctionType] = []  # static variable to stage providers
 
     def __init__(
         self, logger: u.Logger, base_path: str, auth_args: AuthProviderArgs, provider_functions: list[ProviderFunctionType], 
-        user_cls: type[User], *, sa_engine: Engine | None = None
+        custom_user_fields_cls: type[CustomUserFields], *, sa_engine: Engine | None = None, external_only: bool = False
     ):
         self.logger = logger
         self.env_vars = auth_args.env_vars
         self.secret_key = self.env_vars.get(c.SQRL_SECRET_KEY)
+        self.external_only = external_only
 
         # Create a new declarative base for this instance
         self.Base = declarative_base()
         
-        # Define DbBaseUser class for this instance
-        class DbBaseUser(self.Base):
+        # Define DbUser class for this instance
+        class DbUser(self.Base):
             __tablename__ = 'users'
             __table_args__ = {'extend_existing': True}
             username: Mapped[str] = mapped_column(primary_key=True)
-            is_admin: Mapped[bool] = mapped_column(nullable=False, default=False)
+            access_level: Mapped[str] = mapped_column(nullable=False, default="member")
             password_hash: Mapped[str] = mapped_column(nullable=False)
+            custom_fields: Mapped[str] = mapped_column(nullable=False, default="{}")
             created_at: Mapped[datetime] = mapped_column(nullable=False, server_default=func.now())
         
         # Define DbApiKey class for this instance
@@ -124,20 +119,19 @@ class Authenticator(_t.Generic[User]):
             def __repr__(self):
                 return f"<DbOAuthToken(token_id='{self.token_id}', client_id='{self.client_id}', username='{self.username}')>"
         
-        self.DbBaseUser = DbBaseUser
         self.DbApiKey = DbApiKey
         self.DbOAuthClient = DbOAuthClient
         self.DbAuthorizationCode = DbAuthorizationCode
         self.DbOAuthToken = DbOAuthToken
         
-        self.User = user_cls
-        self.DbUser: type[DbBaseUser] = self._initialize_db_user_model(self.User)
+        self.CustomUserFields = custom_user_fields_cls
+        self.DbUser = DbUser
 
         self.auth_providers = [provider_function(auth_args) for provider_function in provider_functions]
         
         if sa_engine is None:
-            sqlite_relative_path = self.env_vars.get(c.SQRL_AUTH_DB_FILE_PATH, f"{c.TARGET_FOLDER}/{c.DB_FILE}")
-            sqlite_path = u.Path(base_path, sqlite_relative_path)
+            raw_sqlite_path = self.env_vars.get(c.SQRL_AUTH_DB_FILE_PATH, f"{{project_path}}/{c.TARGET_FOLDER}/{c.DB_FILE}")
+            sqlite_path = u.Path(raw_sqlite_path.format(project_path=base_path))
             sqlite_path.parent.mkdir(parents=True, exist_ok=True)
             self.engine = create_engine(f"sqlite:///{str(sqlite_path)}")
         else:
@@ -153,67 +147,21 @@ class Authenticator(_t.Generic[User]):
 
         self.Session = sessionmaker(bind=self.engine)
 
-        self._initialize_db(self.User, self.DbUser, self.engine, self.Session)
+        self._initialize_db()
     
-    def _get_user_model(self, base_path: str) -> type[BaseUser]:
-        user_module_path = u.Path(base_path, c.PYCONFIGS_FOLDER, c.USER_FILE)
-        user_module = PyModule(user_module_path)
-        User = user_module.get_func_or_class("User", default_attr=BaseUser)
-        if not issubclass(User, BaseUser):
-            raise ConfigurationError(f"User class in '{c.USER_FILE}' must inherit from BaseUser")
-        return User
-
-    def _initialize_db_user_model(self, *args) -> type:
-        """Get the user model with any custom attributes defined in user.py"""
-        attrs = {}
-
-        # Iterate over all fields in the User model
-        for field_name, field in self.User.model_fields.items():
-            if field_name in reserved_fields:
-                continue
-            if field_name in disallowed_fields:
-                raise ConfigurationError(f"Field name '{field_name}' is disallowed in the User model and cannot be used")
-            
-            field_type = field.annotation
-            if _t.get_origin(field_type) in (_t.Union, types.UnionType):
-                field_type = _t.get_args(field_type)[0]
-                nullable = True
-            else:
-                nullable = False
+    def _convert_db_user_to_user(self, db_user) -> RegisteredUser:
+        """Convert a database user to an AbstractUser object"""
+        # Deserialize custom_fields JSON and merge with defaults
+        custom_fields_json = json.loads(db_user.custom_fields) if db_user.custom_fields else {}
+        custom_fields = self.CustomUserFields(**custom_fields_json)
+        
+        return RegisteredUser(
+            username=db_user.username, 
+            access_level=db_user.access_level,
+            custom_fields=custom_fields
+        )
     
-            if _t.get_origin(field_type) == _t.Literal:
-                field_type = str
-
-            # Map Python types and default values to SQLAlchemy columns
-            default_value = field.default
-            if default_value is PydanticUndefined:
-                raise ConfigurationError(f"No default value found for field '{field_name}' in User model")
-            elif not nullable and default_value is None:
-                raise ConfigurationError(f"Default value for non-nullable field '{field_name}' was set as None in User model")
-            elif default_value is not None and type(default_value) is not field_type:
-                raise ConfigurationError(f"Default value for field '{field_name}' does not match field type in User model")
-            
-            if field_type == str:
-                col_type = String
-            elif field_type == int:
-                col_type = Integer
-            elif field_type == float:
-                col_type = Float
-            elif field_type == bool:
-                col_type = Boolean
-            elif isinstance(field_type, type) and issubclass(field_type, Enum):
-                col_type = String
-                default_value = default_value.value
-            else:
-                continue
-
-            attrs[field_name] = Column(col_type, nullable=nullable, default=default_value) # type: ignore
-
-        # Create the sqlalchemy model class
-        DbUser = type('DbUser', (self.DbBaseUser,), attrs)
-        return DbUser
-
-    def _initialize_db(self, *args): # TODO: Use logger instead of print
+    def _initialize_db(self): # TODO: Use logger instead of print
         session = self.Session()
         try:
             # Get existing columns in the database
@@ -221,14 +169,12 @@ class Authenticator(_t.Generic[User]):
             existing_columns = {col['name'] for col in inspector.get_columns('users')}
 
             # Get all columns defined in the model
-            dropped_columns = set(self.User.dropped_columns())
-            model_columns = set(self.DbUser.__table__.columns.keys()) - dropped_columns
+            model_columns = set(self.DbUser.__table__.columns.keys())
 
             # Find columns that are in the model but not in the database
             new_columns = model_columns - existing_columns
             if new_columns:
                 add_columns_msg = f"Adding columns to database: {new_columns}"
-                print("NOTE -", add_columns_msg)
                 self.logger.info(add_columns_msg)
                 
                 for col_name in new_columns:
@@ -245,24 +191,6 @@ class Authenticator(_t.Generic[User]):
                     session.execute(text(alter_stmt))
                 
                 session.commit()
-            
-            # Determine columns to drop
-            columns_to_drop = dropped_columns.intersection(existing_columns)
-            if columns_to_drop:
-                drop_columns_msg = f"Dropping columns from database: {columns_to_drop}"
-                print("NOTE -", drop_columns_msg)
-                self.logger.info(drop_columns_msg)
-                
-                for col_name in columns_to_drop:
-                    session.execute(text(f"ALTER TABLE users DROP COLUMN {col_name}"))
-                
-                session.commit()
-
-            # Find columns that are in the database but not in the model
-            extra_db_columns = existing_columns - columns_to_drop - model_columns
-            if extra_db_columns:
-                self.logger.warning(f"The following database columns are not in the User model: {extra_db_columns}\n"
-                    "If you want to drop these columns, please use the `dropped_columns` class method of the User model.")
 
             # Get admin password from environment variable if exists
             admin_password = self.env_vars.get(c.SQRL_SECRET_ADMIN_PASSWORD)
@@ -272,7 +200,7 @@ class Authenticator(_t.Generic[User]):
                 password_hash = pwd_context.hash(admin_password)
                 admin_user = session.get(self.DbUser, c.ADMIN_USERNAME)
                 if admin_user is None:
-                    admin_user = self.DbUser(username=c.ADMIN_USERNAME, password_hash=password_hash, is_admin=True)
+                    admin_user = self.DbUser(username=c.ADMIN_USERNAME, password_hash=password_hash, access_level="admin")
                     session.add(admin_user)
                 else:
                     admin_user.password_hash = password_hash
@@ -285,7 +213,7 @@ class Authenticator(_t.Generic[User]):
     @cached_property
     def user_fields(self) -> list[UserField]:
         """
-        Get the fields of the User model as a list of dictionaries
+        Get the fields of the CustomUserFields model as a list of dictionaries
         
         Each dictionary contains the following keys:
         - name: The name of the field
@@ -294,11 +222,15 @@ class Authenticator(_t.Generic[User]):
         - enum: The possible values of the field (or None if not applicable)
         - default: The default value of the field (or None if field is required)
         """
-        schema = self.User.model_json_schema()
-
-        fields = []
-        
-        properties: dict[str, dict[str, _t.Any]] = schema.get("properties", {})
+        # Add the base fields
+        fields = [
+            UserField(name="username", type="string", nullable=False, enum=None, default=None),
+            UserField(name="access_level", type="string", nullable=False, enum=["admin", "member"], default="member")
+        ]
+        
+        # Add custom fields
+        schema = self.CustomUserFields.model_json_schema()
+        properties: dict[str, dict[str, Any]] = schema.get("properties", {})
         for field_name, field_schema in properties.items():
             if choices := field_schema.get("anyOf"):
                 field_type = choices[0]["type"]
@@ -315,37 +247,53 @@ class Authenticator(_t.Generic[User]):
     def add_user(self, username: str, user_fields: dict, *, update_user: bool = False) -> None:
         session = self.Session()
 
-        # Validate the user data
+        # Separate custom fields from base fields
+        access_level = user_fields.get('access_level', 'member')
+        password = user_fields.get('password')
+        
+        # Validate access level - cannot add/update users with guest access level
+        if access_level == "guest":
+            raise InvalidInputError(400, "invalid_access_level", "Cannot create or update users with 'guest' access level")
+        
+        # Extract custom fields (everything except username, access_level, password)
+        custom_fields_dict = {k: v for k, v in user_fields.items() if k not in ['username', 'access_level', 'password']}
+        
+        # Validate the custom fields
         try:
-            user_data = self.User(**user_fields, username=username).model_dump(mode='json')
+            custom_fields = self.CustomUserFields(**custom_fields_dict)
+            custom_fields_json = json.dumps(custom_fields.model_dump(mode='json'))
         except ValidationError as e:
-            raise InvalidInputError(400, "Invalid user data", f"Invalid user field '{e.errors()[0]['loc'][0]}': {e.errors()[0]['msg']}")
+            raise InvalidInputError(400, "invalid_user_data", f"Invalid user field '{e.errors()[0]['loc'][0]}': {e.errors()[0]['msg']}")
 
-        # Add a new user
+        # Add or update user
         try:
             # Check if the user already exists
             existing_user = session.get(self.DbUser, username)
             if existing_user is not None:
                 if not update_user:
-                    raise InvalidInputError(400, "Username already exists", f"User '{username}' already exists")
+                    raise InvalidInputError(400, "username_already_exists", f"User '{username}' already exists")
                 
-                if username == c.ADMIN_USERNAME and user_data.get("is_admin") is False:
-                    raise InvalidInputError(403, "Non-admin 'admin' user not permitted", "Setting the admin user to non-admin is not permitted")
+                if username == c.ADMIN_USERNAME and access_level != "admin":
+                    raise InvalidInputError(403, "admin_cannot_be_non_admin", "Setting the admin user to non-admin is not permitted")
                 
-                new_user = self.DbUser(password_hash=existing_user.password_hash, **user_data)
-                session.delete(existing_user)
+                # Update existing user
+                existing_user.access_level = access_level
+                existing_user.custom_fields = custom_fields_json
             else:
                 if update_user:
-                    raise InvalidInputError(404, "No user found for username", f"No user found for username: {username}")
+                    raise InvalidInputError(404, "no_user_found_for_username", f"No user found for username: {username}")
                 
-                password = user_fields.get('password')
                 if password is None:
-                    raise InvalidInputError(400, "Missing required field 'password'", f"Missing required field 'password' when adding a new user")
+                    raise InvalidInputError(400, "missing_password", f"Missing required field 'password' when adding a new user")
+                
                 password_hash = pwd_context.hash(password)
-                new_user = self.DbUser(password_hash=password_hash, **user_data)
-            
-            # Add the user to the session
-            session.add(new_user)
+                new_user = self.DbUser(
+                    username=username,
+                    access_level=access_level,
+                    password_hash=password_hash,
+                    custom_fields=custom_fields_json
+                )
+                session.add(new_user)
             
             # Commit the transaction
             session.commit()
@@ -353,39 +301,45 @@ class Authenticator(_t.Generic[User]):
         finally:
             session.close()
     
-    def create_or_get_user_from_provider(self, provider_name: str, user_info: dict) -> User:
+    def create_or_get_user_from_provider(self, provider_name: str, user_info: dict) -> RegisteredUser:
         provider = next((p for p in self.auth_providers if p.name == provider_name), None)
         if provider is None:
-            raise InvalidInputError(404, "Provider not found", f"Provider '{provider_name}' not found")
+            raise InvalidInputError(404, "auth_provider_not_found", f"Provider '{provider_name}' not found")
         
         user = provider.provider_configs.get_user(user_info)
         session = self.Session()
         try:
-            db_user = session.get(self.DbUser, user.username)
-            if db_user is None:
-                # Create new user
-                user_data = user.model_dump()
-                password_hash = ""  # No hash makes it impossible to login with username and password
-                db_user = self.DbUser(password_hash=password_hash, **user_data)
+            # Convert user to database user
+            custom_fields_json = user.custom_fields.model_dump_json()
+            db_user = self.DbUser(
+                username=user.username,
+                access_level=user.access_level,
+                password_hash="",  # By omitting password_hash, it becomes impossible to login with username and password (OAuth only)
+                custom_fields=custom_fields_json
+            )
+
+            existing_db_user = session.get(self.DbUser, db_user.username)
+            if existing_db_user is None:
                 session.add(db_user)
-                session.commit()
         
-            return self.User.model_validate(db_user)
+            session.commit()
+
+            return self._convert_db_user_to_user(db_user)
         
         finally:
             session.close()
 
-    def get_user(self, username: str, password: str) -> User:
+    def get_user(self, username: str, password: str) -> RegisteredUser:
         session = self.Session()
         try:
             # Query for user by username
             db_user = session.get(self.DbUser, username)
             
             if db_user and pwd_context.verify(password, db_user.password_hash):
-                user = self.User.model_validate(db_user)
-                return user # type: ignore
+                user = self._convert_db_user_to_user(db_user)
+                return user
             else:
-                raise InvalidInputError(401, "Incorrect username or password", f"Incorrect username or password")
+                raise InvalidInputError(401, "incorrect_username_or_password", f"Incorrect username or password")
 
         finally:
             session.close()
@@ -395,39 +349,39 @@ class Authenticator(_t.Generic[User]):
         try:
             db_user = session.get(self.DbUser, username)
             if db_user is None:
-                raise InvalidInputError(401, "User not found", f"Username '{username}' not found for password change")
+                raise InvalidInputError(401, "user_not_found", f"Username '{username}' not found for password change")
             
             if db_user.password_hash and pwd_context.verify(old_password, db_user.password_hash):
                 db_user.password_hash = pwd_context.hash(new_password)
                 session.commit()
             else:
-                raise InvalidInputError(401, "Incorrect password", f"Incorrect password")
+                raise InvalidInputError(401, "incorrect_password", f"Incorrect password")
         finally:
             session.close()
 
     def delete_user(self, username: str) -> None:
         if username == c.ADMIN_USERNAME:
-            raise InvalidInputError(403, "Cannot delete admin user", "Cannot delete the admin user")
+            raise InvalidInputError(403, "cannot_delete_admin_user", "Cannot delete the admin user")
         
         session = self.Session()
         try:
             db_user = session.get(self.DbUser, username)
             if db_user is None:
-                raise InvalidInputError(404, "No user found for username", f"No user found for username: {username}")
+                raise InvalidInputError(404, "no_user_found_for_username", f"No user found for username: {username}")
             session.delete(db_user)
             session.commit()
         finally:
             session.close()
 
-    def get_all_users(self) -> list:
+    def get_all_users(self) -> list[RegisteredUser]:
         session = self.Session()
         try:
             db_users = session.query(self.DbUser).all()
-            return [self.User.model_validate(user) for user in db_users]
+            return [self._convert_db_user_to_user(user) for user in db_users]
         finally:
             session.close()
     
-    def create_access_token(self, user: User, expiry_minutes: int | None, *, title: str | None = None) -> tuple[str, datetime]:
+    def create_access_token(self, user: AbstractUser, expiry_minutes: int | None, *, title: str | None = None) -> tuple[str, datetime]:
         """
         Creates an API key if title is provided. Otherwise, creates a JWT token.
         """
@@ -440,7 +394,7 @@ class Authenticator(_t.Generic[User]):
         if title is not None:
             session = self.Session()
             try:
-                token_id = "sqrl-" + uuid.uuid4().hex
+                token_id = "sqrl-" + secrets.token_urlsafe(16)
                 hashed_key = u.hash_string(token_id, salt=self.secret_key)
                 api_key = self.DbApiKey(hashed_key=hashed_key, title=title, username=user.username, created_at=created_at, expires_at=expire_at)
                 session.add(api_key)
@@ -453,7 +407,7 @@ class Authenticator(_t.Generic[User]):
         
         return token_id, expire_at
     
-    def get_user_from_token(self, token: str | None) -> User | None:
+    def get_user_from_token(self, token: str | None) -> RegisteredUser | None:
         """
         Get a user from an access token (JWT, or API key if token starts with 'sqrl-')
         """
@@ -481,14 +435,15 @@ class Authenticator(_t.Generic[User]):
             db_user = session.get(self.DbUser, username)
             if db_user is None:
                 raise InvalidTokenError()
+            
+            user = self._convert_db_user_to_user(db_user)
         
         except InvalidTokenError:
-            raise InvalidInputError(401, "Invalid authorization token", "Invalid authorization token")
+            raise InvalidInputError(401, "invalid_authorization_token", "Invalid authorization token")
         finally:
             session.close()
         
-        user = self.User.model_validate(db_user)
-        return user # type: ignore
+        return user
     
     def get_all_api_keys(self, username: str) -> list[ApiKey]:
         """
@@ -518,19 +473,19 @@ class Authenticator(_t.Generic[User]):
             ).first()
             
             if api_key is None:
-                raise InvalidInputError(404, "API key not found", f"The API key could not be found: {api_key_id}")
+                raise InvalidInputError(404, "api_key_not_found", f"The API key could not be found: {api_key_id}")
             
             session.delete(api_key)
             session.commit()
         finally:
             session.close()
 
-    def can_user_access_scope(self, user: User | None, scope: PermissionScope) -> bool:
-        if user is None:
+    def can_user_access_scope(self, user: AbstractUser, scope: PermissionScope) -> bool:
+        if user.access_level == "guest":
             user_level = PermissionScope.PUBLIC
-        elif user.is_admin:
+        elif user.access_level == "admin":
             user_level = PermissionScope.PRIVATE
-        else:
+        else:  # member
             user_level = PermissionScope.PROTECTED
         
         return user_level.value >= scope.value
@@ -840,15 +795,15 @@ class Authenticator(_t.Generic[User]):
                 raise InvalidInputError(400, "invalid_grant", "Invalid code_verifier")
             
             # Get user
-            user = session.get(self.DbUser, auth_code.username)
-            if user is None:
+            db_user = session.get(self.DbUser, auth_code.username)
+            if db_user is None:
                 raise InvalidInputError(400, "invalid_grant", "User not found")
             
             # Mark authorization code as used
             auth_code.used = True
             
             # Generate tokens
-            user_obj = self.User.model_validate(user)
+            user_obj = self._convert_db_user_to_user(db_user)
             access_token, token_expires_at = self.create_access_token(user_obj, expiry_minutes=access_token_expiry_minutes)
             access_token_hash = pwd_context.hash(access_token)
             
@@ -902,8 +857,8 @@ class Authenticator(_t.Generic[User]):
                 raise InvalidInputError(400, "invalid_grant", "Invalid or expired refresh token")
             
             # Get user
-            user = session.get(self.DbUser, oauth_token.username)
-            if user is None:
+            db_user = session.get(self.DbUser, oauth_token.username)
+            if db_user is None:
                 raise InvalidInputError(400, "invalid_client", "User not found")
             
             # Check secret key is available
@@ -911,7 +866,7 @@ class Authenticator(_t.Generic[User]):
                 raise ConfigurationError(f"Environment variable '{c.SQRL_SECRET_KEY}' is required for OAuth token operations")
             
             # Generate new tokens
-            user_obj = self.User.model_validate(user)
+            user_obj = self._convert_db_user_to_user(db_user)
             access_token, token_expires_at = self.create_access_token(user_obj, expiry_minutes=access_token_expiry_minutes)
             access_token_hash = pwd_context.hash(access_token)