feat: implement forgot password feature (#5534)

api/libs/helper.py

@@ -1,18 +1,23 @@
 import json
+import logging
 import random
 import re
 import string
 import subprocess
+import time
 import uuid
 from collections.abc import Generator
 from datetime import datetime
 from hashlib import sha256
-from typing import Union
+from typing import Any, Optional, Union
 from zoneinfo import available_timezones
 
-from flask import Response, stream_with_context
+from flask import Response, current_app, stream_with_context
 from flask_restful import fields
 
+from extensions.ext_redis import redis_client
+from models.account import Account
+
 
 def run(script):
     return subprocess.getstatusoutput('source /root/.bashrc && ' + script)
@@ -46,12 +51,12 @@ def uuid_value(value):
         error = ('{value} is not a valid uuid.'
                  .format(value=value))
         raise ValueError(error)
-    
+
 def alphanumeric(value: str):
     # check if the value is alphanumeric and underlined
     if re.match(r'^[a-zA-Z0-9_]+$', value):
         return value
-    
+
     raise ValueError(f'{value} is not a valid alphanumeric value')
 
 def timestamp_value(timestamp):
@@ -163,3 +168,97 @@ def compact_generate_response(response: Union[dict, Generator]) -> Response:
 
         return Response(stream_with_context(generate()), status=200,
                         mimetype='text/event-stream')
+
+
+class TokenManager:
+
+    @classmethod
+    def generate_token(cls, account: Account, token_type: str, additional_data: dict = None) -> str:
+        old_token = cls._get_current_token_for_account(account.id, token_type)
+        if old_token:
+            if isinstance(old_token, bytes):
+                old_token = old_token.decode('utf-8')
+            cls.revoke_token(old_token, token_type)
+
+        token = str(uuid.uuid4())
+        token_data = {
+            'account_id': account.id,
+            'email': account.email,
+            'token_type': token_type
+        }
+        if additional_data:
+            token_data.update(additional_data)
+
+        expiry_hours = current_app.config[f'{token_type.upper()}_TOKEN_EXPIRY_HOURS']
+        token_key = cls._get_token_key(token, token_type)
+        redis_client.setex(
+            token_key,
+            expiry_hours * 60 * 60,
+            json.dumps(token_data)
+        )
+
+        cls._set_current_token_for_account(account.id, token, token_type, expiry_hours)
+        return token
+
+    @classmethod
+    def _get_token_key(cls, token: str, token_type: str) -> str:
+        return f'{token_type}:token:{token}'
+
+    @classmethod
+    def revoke_token(cls, token: str, token_type: str):
+        token_key = cls._get_token_key(token, token_type)
+        redis_client.delete(token_key)
+
+    @classmethod
+    def get_token_data(cls, token: str, token_type: str) -> Optional[dict[str, Any]]:
+        key = cls._get_token_key(token, token_type)
+        token_data_json = redis_client.get(key)
+        if token_data_json is None:
+            logging.warning(f"{token_type} token {token} not found with key {key}")
+            return None
+        token_data = json.loads(token_data_json)
+        return token_data
+
+    @classmethod
+    def _get_current_token_for_account(cls, account_id: str, token_type: str) -> Optional[str]:
+        key = cls._get_account_token_key(account_id, token_type)
+        current_token = redis_client.get(key)
+        return current_token
+
+    @classmethod
+    def _set_current_token_for_account(cls, account_id: str, token: str, token_type: str, expiry_hours: int):
+        key = cls._get_account_token_key(account_id, token_type)
+        redis_client.setex(key, expiry_hours * 60 * 60, token)
+
+    @classmethod
+    def _get_account_token_key(cls, account_id: str, token_type: str) -> str:
+        return f'{token_type}:account:{account_id}'
+
+
+class RateLimiter:
+    def __init__(self, prefix: str, max_attempts: int, time_window: int):
+        self.prefix = prefix
+        self.max_attempts = max_attempts
+        self.time_window = time_window
+
+    def _get_key(self, email: str) -> str:
+        return f"{self.prefix}:{email}"
+
+    def is_rate_limited(self, email: str) -> bool:
+        key = self._get_key(email)
+        current_time = int(time.time())
+        window_start_time = current_time - self.time_window
+
+        redis_client.zremrangebyscore(key, '-inf', window_start_time)
+        attempts = redis_client.zcard(key)
+
+        if attempts and int(attempts) >= self.max_attempts:
+            return True
+        return False
+
+    def increment_rate_limit(self, email: str):
+        key = self._get_key(email)
+        current_time = int(time.time())
+
+        redis_client.zadd(key, {current_time: current_time})
+        redis_client.expire(key, self.time_window * 2)