from typing import List, Optional import bcrypt from sqlalchemy import Column, Table, ForeignKey, UniqueConstraint from sqlalchemy.dialects.mysql import insert from sqlalchemy.engine import Row from sqlalchemy.sql import func, select, and_, or_ from sqlalchemy.types import TIMESTAMP, Integer, String, BIGINT, INTEGER from datetime import datetime from core.data.schema.base import BaseData, metadata, campaign, campaign_game aime_user: Table = Table( "aime_user", metadata, Column("id", Integer, nullable=False, primary_key=True, autoincrement=True), Column("username", String(25), unique=True), Column("email", String(255), unique=True), Column("password", String(255)), Column("permissions", Integer), Column("created_date", TIMESTAMP, server_default=func.now()), Column("last_login_date", TIMESTAMP, onupdate=func.now()), Column("suspend_expire_time", TIMESTAMP), mysql_charset="utf8mb4", ) campaign_progress: Table = Table( "campaign_progress", metadata, Column("id", BIGINT, nullable=False, primary_key=True, autoincrement=True), Column("user_id", Integer, ForeignKey("aime_user.id", ondelete="cascade", onupdate="cascade"), nullable=False), Column("campaign_id", INTEGER, ForeignKey("campaign.id", ondelete="cascade", onupdate="cascade"), nullable=False), Column("is_participating", INTEGER, nullable=False, server_default="0"), Column("progress", INTEGER, nullable=False, server_default="0"), UniqueConstraint("campaign_id", "user_id", name="campaign_progress_uk"), mysql_charset="utf8mb4", ) class UserData(BaseData): async def create_user( self, id: Optional[int] = None, username: Optional[str] = None, email: Optional[str] = None, password: Optional[str] = None, permission: int = 1, ) -> Optional[int]: if id is None: sql = insert(aime_user).values( username=username, email=email, password=password, permissions=permission, ) else: sql = insert(aime_user).values( id=id, username=username, email=email, password=password, permissions=permission, ) conflict = sql.on_duplicate_key_update( username=username, email=email, password=password, permissions=permission ) result = await self.execute(conflict) if result is None: return None return result.lastrowid async def get_user(self, user_id: int) -> Optional[Row]: sql = select(aime_user).where(aime_user.c.id == user_id) result = await self.execute(sql) if result is None: return False return result.fetchone() async def check_password(self, user_id: int, passwd: bytes = None) -> bool: usr = await self.get_user(user_id) if usr is None: return False if usr["password"] is None: return False if passwd is None or not passwd: return False return bcrypt.checkpw(passwd, usr["password"].encode()) async def delete_user(self, user_id: int) -> None: sql = aime_user.delete(aime_user.c.id == user_id) result = await self.execute(sql) if result is None: self.logger.error(f"Failed to delete user with id {user_id}") async def get_unregistered_users(self) -> List[Row]: """ Returns a list of users who have not registered with the webui. They may or may not have cards. """ sql = select(aime_user).where(aime_user.c.password == None) result = await self.execute(sql) if result is None: return None return result.fetchall() async def find_user_by_email(self, email: str) -> Row: sql = select(aime_user).where(aime_user.c.email == email) result = await self.execute(sql) if result is None: return False return result.fetchone() async def find_user_by_username(self, username: str) -> List[Row]: sql = aime_user.select(aime_user.c.username.like(f"%{username}%")) result = await self.execute(sql) if result is None: return False return result.fetchall() async def change_password(self, user_id: int, new_passwd: str) -> bool: sql = aime_user.update(aime_user.c.id == user_id).values(password = new_passwd) result = await self.execute(sql) return result is not None async def change_username(self, user_id: int, new_name: str) -> bool: sql = aime_user.update(aime_user.c.id == user_id).values(username = new_name) result = await self.execute(sql) return result is not None async def get_user_by_username(self, username: str) -> Optional[Row]: result = await self.execute(aime_user.select(aime_user.c.username == username)) if result: return result.fetchone() async def change_permission(self, user_id: int, new_perms: int) -> Optional[bool]: sql = aime_user.update(aime_user.c.id == user_id).values(permissions = new_perms) result = await self.execute(sql) return result is not None async def change_email(self, user_id: int, new_email: int) -> Optional[bool]: sql = aime_user.update(aime_user.c.id == user_id).values(email = new_email) result = await self.execute(sql) return result is not None async def get_user_campaign_progress_by_campaign(self, user_id: int, campaign_id: int) -> Optional[Row]: result = await self.execute(campaign_progress .join(campaign, campaign_progress.c.campaign_id == campaign.c.id) .select(and_(campaign_progress.c.user_id == user_id, campaign_progress.c.campaign_id == campaign_id))) if result is not None: return result.fetchone() async def get_user_all_campaign_progress(self, user_id: int) -> Optional[List[Row]]: result = await self.execute(campaign_progress .join(campaign, campaign_progress.c.campaign_id == campaign.c.id) .select(campaign_progress.c.user_id == user_id)) if result is not None: return result.fetchall() async def get_user_active_campaign_progress(self, user_id: int) -> Optional[List[Row]]: # distrib_end_date instead of end_date to give the games time to distribute the rewards result = await self.execute(campaign_progress .join(campaign, campaign_progress.c.campaign_id == campaign.c.id) .select(and_(campaign_progress.c.user_id == user_id, and_( or_(campaign.c.start_date >= datetime.now(), campaign.c.start_date == None), or_(campaign.c.distrib_end_date < datetime.now(), campaign.c.distrib_end_date == None), )))) if result is not None: return result.fetchall() async def get_user_all_campaign_progress_by_game(self, user_id: int, game_id: str) -> Optional[List[Row]]: result = await self.execute(campaign_progress .join(campaign, campaign_progress.c.campaign_id == campaign.c.id) .join(campaign_game, campaign_progress.c.campaign_id == campaign_game.c.campaign_id) .select(and_(campaign_progress.c.user_id == user_id, campaign_game.c.game_id == game_id))) if result is not None: return result.fetchall() async def get_user_active_campaign_progress_by_game(self, user_id: int, game_id: str) -> Optional[List[Row]]: result = await self.execute(campaign_progress .join(campaign, campaign_progress.c.campaign_id == campaign.c.id) .join(campaign_game, campaign_progress.c.campaign_id == campaign_game.c.campaign_id) .select(and_(and_(campaign_progress.c.user_id == user_id, campaign_game.c.game_id == game_id), and_( or_(campaign.c.start_date >= datetime.now(), campaign.c.start_date == None), or_(campaign.c.distrib_end_date < datetime.now(), campaign.c.distrib_end_date == None), )))) if result is not None: return result.fetchall() async def set_user_campaign_is_participating(self, user_id: int, campaign_id: int, is_participating: bool) -> Optional[bool]: sql = insert(campaign_progress).values( user_id = user_id, campaign_id = campaign_id, is_participating = is_participating ) conflict = sql.on_duplicate_key_update(is_participating = is_participating) result = await self.execute(conflict) if result is None: self.logger.error(f"Failed to update campaign participation status for user {user_id} on campaign {campaign_id}!") return None return result.lastrowid async def set_user_campaign_progress(self, user_id: int, campaign_id: int, progress: int) -> Optional[bool]: sql = insert(campaign_progress).values( user_id = user_id, campaign_id = campaign_id, is_participating = 1, progress = progress ) conflict = sql.on_duplicate_key_update(progress = progress) result = await self.execute(conflict) if result is None: self.logger.error(f"Failed to update campaign progress for user {user_id} on campaign {campaign_id} to {progress}!") return None return result.lastrowid