Преглед на файлове

feat: migrate temp tokens to Redis and add cleanup script

unknown преди 2 дни
родител
ревизия
c0944daa09
променени са 3 файла, в които са добавени 100 реда и са изтрити 22 реда
  1. 31 0
      backend/cleanup_obsolete_tables.py
  2. 15 22
      backend/routers/auth.py
  3. 54 0
      backend/services/token_service.py

+ 31 - 0
backend/cleanup_obsolete_tables.py

@@ -0,0 +1,31 @@
+import db
+import sys
+
+def cleanup():
+    print("Starting database cleanup...")
+    
+    tables_to_drop = [
+        "email_verification_tokens",
+        "password_reset_tokens"
+    ]
+    
+    try:
+        # Check if tables exist before dropping (optional but safer)
+        for table in tables_to_drop:
+            print(f"Dropping table: {table}...")
+            db.execute_commit(f"DROP TABLE IF EXISTS {table}")
+            print(f"Successfully dropped {table}.")
+            
+        print("\nDatabase cleanup completed successfully!")
+        print("All temporary tokens are now managed by Redis.")
+        
+    except Exception as e:
+        print(f"Error during cleanup: {e}", file=sys.stderr)
+        sys.exit(1)
+
+if __name__ == "__main__":
+    confirm = input("This will permanently delete token tables from MySQL. Are you sure? (y/N): ")
+    if confirm.lower() == 'y':
+        cleanup()
+    else:
+        print("Cleanup cancelled.")

+ 15 - 22
backend/routers/auth.py

@@ -13,6 +13,7 @@ from dependencies import get_current_user, require_admin
 import config
 import secrets
 from services.email_service import send_verification_email, send_password_reset_email
+from services.token_service import token_service
 
 try:
     from google.oauth2 import id_token
@@ -40,10 +41,8 @@ async def register(request: Request, user: schemas.UserCreate, lang: str = "en")
     
     user_id = db.execute_commit(query, params)
     
-    # Generate Verification Token
-    token = secrets.token_urlsafe(32)
-    expires_at = datetime.utcnow() + timedelta(hours=24)
-    db.execute_commit("INSERT INTO email_verification_tokens (user_id, token, expires_at) VALUES (%s, %s, %s)", (user_id, token, expires_at))
+    # Generate Verification Token (Redis)
+    token = token_service.create_verification_token(user_id)
     
     # Send Email
     send_verification_email(user.email, token, user.preferred_language or lang)
@@ -53,16 +52,12 @@ async def register(request: Request, user: schemas.UserCreate, lang: str = "en")
 
 @router.get("/verify-email")
 async def verify_email(token: str, lang: str = "en"):
-    res = db.execute_query("SELECT user_id, expires_at FROM email_verification_tokens WHERE token = %s", (token,))
-    if not res:
-        raise HTTPException(status_code=400, detail="Invalid verification token")
-    
-    user_id, expires_at = res[0]['user_id'], res[0]['expires_at']
-    if expires_at < datetime.utcnow():
-        raise HTTPException(status_code=400, detail="Verification token expired")
+    user_id = token_service.verify_email_token(token)
+    if not user_id:
+        raise HTTPException(status_code=400, detail="Invalid or expired verification token")
     
     db.execute_commit("UPDATE users SET is_active = 1 WHERE id = %s", (user_id,))
-    db.execute_commit("DELETE FROM email_verification_tokens WHERE user_id = %s", (user_id,))
+    token_service.delete_verification_token(token)
     
     return {"message": "Email verified successfully. You can now log in."}
 
@@ -182,9 +177,8 @@ async def forgot_password(request: schemas.ForgotPassword, lang: str = "en"):
     user = db.execute_query("SELECT id, preferred_language FROM users WHERE email = %s", (request.email,))
     if not user: raise HTTPException(status_code=404, detail="Email not found")
     
-    token = secrets.token_urlsafe(32)
-    expires_at = datetime.utcnow() + timedelta(minutes=10)
-    db.execute_commit("INSERT INTO password_reset_tokens (user_id, token, expires_at) VALUES (%s, %s, %s)", (user[0]['id'], token, expires_at))
+    # Generate Reset Token (Redis - 10 min TTL)
+    token = token_service.create_reset_token(user[0]['id'])
     
     # Send Email
     user_lang = user[0]['preferred_language'] or lang
@@ -194,18 +188,17 @@ async def forgot_password(request: schemas.ForgotPassword, lang: str = "en"):
 
 @router.post("/reset-password")
 async def reset_password(request: schemas.ResetPassword):
-    reset_data = db.execute_query("SELECT user_id, expires_at FROM password_reset_tokens WHERE token = %s", (request.token,))
-    if not reset_data: raise HTTPException(status_code=400, detail="Invalid token")
-    if reset_data[0]['expires_at'] < datetime.utcnow(): raise HTTPException(status_code=400, detail="Token expired")
+    user_id = token_service.verify_reset_token(request.token)
+    if not user_id:
+        raise HTTPException(status_code=400, detail="Invalid or expired reset token")
     
-    user_id = reset_data[0]['user_id']
     hashed_password = auth_utils.get_password_hash(request.new_password)
     db.execute_commit("UPDATE users SET password_hash = %s WHERE id = %s", (hashed_password, user_id))
     
-    # IMPORTANT: Delete ALL tokens for this user after successful reset
-    db.execute_commit("DELETE FROM password_reset_tokens WHERE user_id = %s", (user_id,))
+    # Successful reset - Cleanup ALL reset tokens for this user
+    token_service.cleanup_reset_tokens(user_id)
     
-    return {"message": "Password reset successfully"}
+    return {"message": "Password updated successfully"}
 
 @router.get("/me", response_model=schemas.UserResponse)
 async def get_me(user: dict = Depends(get_current_user)):

+ 54 - 0
backend/services/token_service.py

@@ -0,0 +1,54 @@
+import secrets
+from session_utils import r
+
+class TokenService:
+    @staticmethod
+    def generate_token(length: int = 32) -> str:
+        return secrets.token_urlsafe(length)
+
+    def create_verification_token(self, user_id: int, expires_seconds: int = 86400) -> str:
+        """Create an email verification token (default 24h)"""
+        token = self.generate_token()
+        # Key: verify_token:TOKEN -> user_id
+        r.setex(f"verify_token:{token}", expires_seconds, str(user_id))
+        return token
+
+    def verify_email_token(self, token: str) -> int:
+        """Verify token and return user_id, or None if invalid/expired"""
+        user_id = r.get(f"verify_token:{token}")
+        if user_id:
+            # We don't delete immediately here? Usually yes, after successful verification.
+            return int(user_id)
+        return None
+
+    def delete_verification_token(self, token: str):
+        r.delete(f"verify_token:{token}")
+
+    def create_reset_token(self, user_id: int, expires_seconds: int = 600) -> str:
+        """Create a password reset token (default 10 minutes)"""
+        token = self.generate_token()
+        # Key: reset_token:TOKEN -> user_id
+        r.setex(f"reset_token:{token}", expires_seconds, str(user_id))
+        
+        # Track all reset tokens for this user to allow bulk cleanup
+        r.sadd(f"user_reset_tokens:{user_id}", token)
+        r.expire(f"user_reset_tokens:{user_id}", expires_seconds)
+        
+        return token
+
+    def verify_reset_token(self, token: str) -> int:
+        """Verify password reset token"""
+        user_id = r.get(f"reset_token:{token}")
+        return int(user_id) if user_id else None
+
+    def cleanup_reset_tokens(self, user_id: int):
+        """Delete all reset tokens associated with a user"""
+        tokens = r.smembers(f"user_reset_tokens:{user_id}")
+        if tokens:
+            # Delete individual tokens
+            for token in tokens:
+                r.delete(f"reset_token:{token}")
+            # Delete the set itself
+            r.delete(f"user_reset_tokens:{user_id}")
+
+token_service = TokenService()