import hashlib import logging import secrets from datetime import datetime, timedelta, timezone import bcrypt from sqlalchemy import func, select, update from fabledassistant.models import async_session from fabledassistant.models.conversation import Conversation from fabledassistant.models.note import Note from fabledassistant.models.invitation import InvitationToken from fabledassistant.models.password_reset import PasswordResetToken from fabledassistant.models.setting import Setting from fabledassistant.models.user import User logger = logging.getLogger(__name__) def hash_password(password: str) -> str: return bcrypt.hashpw(password.encode(), bcrypt.gensalt()).decode() def verify_password(password: str, password_hash: str) -> bool: return bcrypt.checkpw(password.encode(), password_hash.encode()) async def get_user_count() -> int: async with async_session() as session: return await session.scalar(select(func.count(User.id))) or 0 async def create_user( username: str, password: str | None, email: str | None = None, oauth_sub: str | None = None, ) -> User: user_count = await get_user_count() role = "admin" if user_count == 0 else "user" async with async_session() as session: user = User( username=username, email=email, password_hash=hash_password(password) if password is not None else None, oauth_sub=oauth_sub, role=role, ) session.add(user) await session.commit() await session.refresh(user) # First user claims all pre-existing data (migration assigns user_id=1, # which matches SERIAL first insert; also handles any NULLs) if role == "admin": await session.execute( update(Note) .where((Note.user_id.is_(None)) | (Note.user_id == user.id)) .values(user_id=user.id) ) await session.execute( update(Conversation) .where( (Conversation.user_id.is_(None)) | (Conversation.user_id == user.id) ) .values(user_id=user.id) ) await session.execute( update(Setting) .where(Setting.user_id.is_(None)) .values(user_id=user.id) ) # Auto-close registration after first user setup session.add(Setting(user_id=user.id, key="registration_open", value="false")) await session.commit() logger.info("First user '%s' created as admin, claimed orphaned data", username) return user async def authenticate(username: str, password: str) -> User | None: async with async_session() as session: result = await session.execute( select(User).where(User.username == username) ) user = result.scalars().first() if user is None: return None if user.password_hash is None: # OAuth-only user — cannot use password login return None if verify_password(password, user.password_hash): return user return None async def get_user_by_id(user_id: int) -> User | None: async with async_session() as session: return await session.get(User, user_id) async def get_user_by_username(username: str) -> User | None: async with async_session() as session: result = await session.execute( select(User).where(User.username == username) ) return result.scalars().first() async def change_password(user_id: int, current_password: str, new_password: str) -> int | None: """Change a user's password. Returns new session_version on success, None if current password is wrong.""" async with async_session() as session: user = await session.get(User, user_id) if not user: return None if not verify_password(current_password, user.password_hash): return None user.password_hash = hash_password(new_password) user.session_version += 1 await session.commit() logger.info("Password changed for user %d (%s)", user_id, user.username) return user.session_version async def invalidate_other_sessions(user_id: int) -> int: """Bump session_version so all sessions except the current one are invalidated. Returns the new session_version so the caller can update the active session cookie. """ async with async_session() as session: user = await session.get(User, user_id) user.session_version += 1 await session.commit() logger.info("Sessions invalidated for user %d (%s)", user_id, user.username) return user.session_version async def is_registration_open() -> bool: """Check if new user registration is allowed. Always open when no users exist (first-user setup). Otherwise reads the admin's 'registration_open' setting (default closed). """ user_count = await get_user_count() if user_count == 0: return True async with async_session() as session: # Find the admin user's registration_open setting result = await session.execute( select(Setting) .join(User, Setting.user_id == User.id) .where(User.role == "admin", Setting.key == "registration_open") ) setting = result.scalar_one_or_none() return setting.value == "true" if setting else False async def list_users() -> list[User]: async with async_session() as session: result = await session.execute(select(User).order_by(User.created_at)) return list(result.scalars().all()) async def delete_user(user_id: int) -> bool: async with async_session() as session: user = await session.get(User, user_id) if not user: return False await session.delete(user) await session.commit() logger.info("Deleted user %d (%s)", user_id, user.username) return True async def set_registration_open(admin_user_id: int, open: bool) -> None: from fabledassistant.services.settings import set_setting await set_setting(admin_user_id, "registration_open", "true" if open else "false") async def update_user_email(user_id: int, email: str | None) -> None: async with async_session() as session: user = await session.get(User, user_id) if user: user.email = email await session.commit() async def get_user_by_oauth_sub(sub: str) -> User | None: async with async_session() as session: result = await session.execute( select(User).where(User.oauth_sub == sub) ) return result.scalars().first() async def link_oauth_sub(user_id: int, sub: str) -> None: async with async_session() as session: await session.execute( update(User).where(User.id == user_id).values(oauth_sub=sub) ) await session.commit() async def get_user_by_email(email: str) -> User | None: async with async_session() as session: result = await session.execute( select(User).where(func.lower(User.email) == email.lower()) ) return result.scalars().first() async def create_password_reset_token(user_id: int) -> str: """Generate a password reset token. Returns the raw token (for the email link).""" raw_token = secrets.token_urlsafe(32) token_hash = hashlib.sha256(raw_token.encode()).hexdigest() expires_at = datetime.now(timezone.utc) + timedelta(hours=1) async with async_session() as session: # Invalidate any existing unused tokens for this user result = await session.execute( select(PasswordResetToken).where( PasswordResetToken.user_id == user_id, PasswordResetToken.used.is_(False), ) ) for old_token in result.scalars().all(): old_token.used = True token = PasswordResetToken( user_id=user_id, token_hash=token_hash, expires_at=expires_at, ) session.add(token) await session.commit() logger.info("Password reset token created for user %d", user_id) return raw_token async def reset_password_with_token(raw_token: str, new_password: str) -> int | None: """Validate a reset token and update the user's password. Returns user_id on success.""" token_hash = hashlib.sha256(raw_token.encode()).hexdigest() async with async_session() as session: result = await session.execute( select(PasswordResetToken).where(PasswordResetToken.token_hash == token_hash) ) reset_token = result.scalars().first() if not reset_token: return None if reset_token.used: return None if reset_token.expires_at < datetime.now(timezone.utc): return None user = await session.get(User, reset_token.user_id) if not user: return None user.password_hash = hash_password(new_password) user.session_version += 1 reset_token.used = True await session.commit() logger.info("Password reset via token for user %d (%s)", user.id, user.username) return user.id async def create_invitation(email: str, invited_by: int) -> str: """Generate an invitation token. Returns the raw token (for the email link).""" raw_token = secrets.token_urlsafe(32) token_hash = hashlib.sha256(raw_token.encode()).hexdigest() expires_at = datetime.now(timezone.utc) + timedelta(days=7) async with async_session() as session: # Invalidate previous unused invitations for same email result = await session.execute( select(InvitationToken).where( func.lower(InvitationToken.email) == email.lower(), InvitationToken.used.is_(False), ) ) for old_token in result.scalars().all(): old_token.used = True token = InvitationToken( email=email.lower(), token_hash=token_hash, invited_by=invited_by, expires_at=expires_at, ) session.add(token) await session.commit() logger.info("Invitation created for %s by user %d", email, invited_by) return raw_token async def validate_invitation_token(raw_token: str) -> InvitationToken | None: """Look up by hash, check not used/expired. Returns the token record with email.""" token_hash = hashlib.sha256(raw_token.encode()).hexdigest() async with async_session() as session: result = await session.execute( select(InvitationToken).where(InvitationToken.token_hash == token_hash) ) invitation = result.scalars().first() if not invitation: return None if invitation.used: return None if invitation.expires_at < datetime.now(timezone.utc): return None return invitation async def register_with_invitation(raw_token: str, username: str, password: str) -> User | None: """Validate token, create user with the invitation's email, mark token used.""" token_hash = hashlib.sha256(raw_token.encode()).hexdigest() async with async_session() as session: result = await session.execute( select(InvitationToken).where(InvitationToken.token_hash == token_hash) ) invitation = result.scalars().first() if not invitation or invitation.used or invitation.expires_at < datetime.now(timezone.utc): return None invitation.used = True await session.commit() # Create user outside the invitation session user = await create_user(username, password, invitation.email) logger.info("User '%s' registered via invitation for %s", username, invitation.email) return user async def list_pending_invitations() -> list[InvitationToken]: """List unused, non-expired invitations.""" async with async_session() as session: result = await session.execute( select(InvitationToken).where( InvitationToken.used.is_(False), InvitationToken.expires_at > datetime.now(timezone.utc), ).order_by(InvitationToken.created_at.desc()) ) return list(result.scalars().all()) async def revoke_invitation(invitation_id: int) -> bool: """Mark an invitation as used (revoked).""" async with async_session() as session: invitation = await session.get(InvitationToken, invitation_id) if not invitation or invitation.used: return False invitation.used = True await session.commit() logger.info("Invitation %d revoked", invitation_id) return True