recovery_service.py 1.5 KB

12345678910111213141516171819202122232425262728293031323334
  1. import json
  2. import redis
  3. from pydantic import BaseModel
  4. import random
  5. class RedisRecoveryData(BaseModel):
  6. code: int
  7. RECOVERY_REDIS_KEY = "pin_recovery_{user_id}"
  8. def generate_recovery_key(user_id:int):
  9. pin_code = str(random.randint(100000, 999999))
  10. redis_client = redis.Redis(host='localhost', port=6379, db=0, decode_responses=True)
  11. redis_client.set(RECOVERY_REDIS_KEY.format(user_id=user_id), json.dumps({"code": pin_code}))
  12. redis_client.expire(RECOVERY_REDIS_KEY.format(user_id=user_id), 900) # Expire in 15 minutes
  13. return pin_code
  14. def add_token(user_id, token):
  15. redis_client = redis.Redis(host='localhost', port=6379, db=0, decode_responses=True)
  16. redis_client.delete(RECOVERY_REDIS_KEY.format(user_id=user_id))
  17. redis_client.set(RECOVERY_REDIS_KEY.format(user_id=user_id), json.dumps({"token": token}))
  18. def get_token(user_id):
  19. redis_client = redis.Redis(host='localhost', port=6379, db=0, decode_responses=True)
  20. data = redis_client.get(RECOVERY_REDIS_KEY.format(user_id=user_id))
  21. json_data = json.loads(str(data)) if data else None
  22. return json_data.get("token") if json_data else None
  23. def get_recovery_data(user_id: int) -> RedisRecoveryData:
  24. redis_client = redis.Redis(host='localhost', port=6379, db=0,decode_responses=True)
  25. data = redis_client.get(RECOVERY_REDIS_KEY.format(user_id=user_id))
  26. json_data = json.loads(str(data)) if data else None
  27. if json_data:
  28. return RedisRecoveryData(**json_data)
  29. return RedisRecoveryData(code=-1)