smallbusinessinventorymanagementsystem-r5ffzu/app/crud/user.py

from datetime import datetime, timedelta
from typing import Any, Dict, Optional, Union

from sqlalchemy.orm import Session

from app.core.config import settings
from app.core.security import (
    generate_password_reset_token,
    generate_verification_token,
    get_password_hash,
    verify_password,
)
from app.crud.base import CRUDBase
from app.models.user import User
from app.schemas.user import UserCreate, UserUpdate


class CRUDUser(CRUDBase[User, UserCreate, UserUpdate]):
    """
    CRUD operations for User model.
    """
    def get_by_email(self, db: Session, *, email: str) -> Optional[User]:
        """
        Get a user by email.
        """
        return db.query(User).filter(User.email == email).first()

    def get_by_verification_token(self, db: Session, *, token: str) -> Optional[User]:
        """
        Get a user by verification token.
        """
        return db.query(User).filter(
            User.verification_token == token,
            User.verification_token_expires > datetime.utcnow()
        ).first()

    def get_by_password_reset_token(self, db: Session, *, token: str) -> Optional[User]:
        """
        Get a user by password reset token.
        """
        return db.query(User).filter(
            User.password_reset_token == token,
            User.password_reset_token_expires > datetime.utcnow()
        ).first()

    def create(self, db: Session, *, obj_in: UserCreate) -> User:
        """
        Create a new user.
        """
        # Create verification token that expires in 48 hours
        verification_token = generate_verification_token()
        verification_token_expires = datetime.utcnow() + timedelta(
            hours=settings.VERIFICATION_TOKEN_EXPIRE_HOURS
        )

        # Create the user
        db_obj = User(
            email=obj_in.email,
            hashed_password=get_password_hash(obj_in.password),
            full_name=obj_in.full_name,
            is_superuser=obj_in.is_superuser,
            is_active=obj_in.is_active,
            is_verified=False,  # New users are not verified by default
            verification_token=verification_token,
            verification_token_expires=verification_token_expires,
        )
        db.add(db_obj)
        db.commit()
        db.refresh(db_obj)
        return db_obj

    def update(
        self, db: Session, *, db_obj: User, obj_in: Union[UserUpdate, Dict[str, Any]]
    ) -> User:
        """
        Update a user.
        """
        if isinstance(obj_in, dict):
            update_data = obj_in
        else:
            update_data = obj_in.model_dump(exclude_unset=True)
        
        if update_data.get("password"):
            hashed_password = get_password_hash(update_data["password"])
            del update_data["password"]
            update_data["hashed_password"] = hashed_password
        
        return super().update(db, db_obj=db_obj, obj_in=update_data)

    def authenticate(self, db: Session, *, email: str, password: str) -> Optional[User]:
        """
        Authenticate a user.
        """
        user = self.get_by_email(db, email=email)
        if not user:
            return None
        if not verify_password(password, user.hashed_password):
            return None
        return user

    def is_active(self, user: User) -> bool:
        """
        Check if a user is active.
        """
        return user.is_active

    def is_superuser(self, user: User) -> bool:
        """
        Check if a user is a superuser.
        """
        return user.is_superuser

    def is_verified(self, user: User) -> bool:
        """
        Check if a user has verified their email.
        """
        return user.is_verified

    def verify_email(self, db: Session, *, token: str) -> Optional[User]:
        """
        Verify a user's email with the provided token.
        """
        user = self.get_by_verification_token(db, token=token)
        if not user:
            return None
        
        # Mark the user as verified and clear the token
        user_data = {"is_verified": True, "verification_token": None, "verification_token_expires": None}
        user_updated = super().update(db, db_obj=user, obj_in=user_data)
        return user_updated

    def create_password_reset_token(self, db: Session, *, email: str) -> Optional[str]:
        """
        Create a password reset token for a user.
        """
        user = self.get_by_email(db, email=email)
        if not user:
            return None
        
        # Create a token that expires in 24 hours
        reset_token = generate_password_reset_token()
        reset_token_expires = datetime.utcnow() + timedelta(
            hours=settings.PASSWORD_RESET_TOKEN_EXPIRE_HOURS
        )
        
        # Update the user with the reset token
        user_data = {
            "password_reset_token": reset_token,
            "password_reset_token_expires": reset_token_expires
        }
        self.update(db, db_obj=user, obj_in=user_data)
        
        return reset_token

    def reset_password(self, db: Session, *, token: str, new_password: str) -> Optional[User]:
        """
        Reset a user's password using a valid reset token.
        """
        user = self.get_by_password_reset_token(db, token=token)
        if not user:
            return None
        
        # Update the password and clear the token
        user_data = {
            "password": new_password, 
            "password_reset_token": None,
            "password_reset_token_expires": None
        }
        user_updated = self.update(db, db_obj=user, obj_in=user_data)
        return user_updated


# Create a singleton instance
user = CRUDUser(User)