security.py 2.5 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879
  1. """
  2. 安全相关功能
  3. """
  4. try:
  5. from passlib.context import CryptContext
  6. pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
  7. USE_PASSLIB = True
  8. except Exception:
  9. # 如果 passlib 有问题,使用 bcrypt 直接验证
  10. import bcrypt
  11. USE_PASSLIB = False
  12. try:
  13. from jose import JWTError, jwt
  14. except ImportError:
  15. # 如果没有jose,使用简单的实现
  16. import jwt
  17. JWTError = Exception
  18. from datetime import datetime, timedelta
  19. from typing import Optional
  20. from app.config import settings
  21. def verify_password(plain_password: str, hashed_password: str) -> bool:
  22. """验证密码"""
  23. if USE_PASSLIB:
  24. try:
  25. return pwd_context.verify(plain_password, hashed_password)
  26. except Exception:
  27. # 如果 passlib 验证失败,尝试使用 bcrypt 直接验证
  28. import bcrypt
  29. try:
  30. return bcrypt.checkpw(plain_password.encode('utf-8'), hashed_password.encode('utf-8'))
  31. except Exception:
  32. return False
  33. else:
  34. # 直接使用 bcrypt
  35. import bcrypt
  36. try:
  37. return bcrypt.checkpw(plain_password.encode('utf-8'), hashed_password.encode('utf-8'))
  38. except Exception:
  39. return False
  40. def get_password_hash(password: str) -> str:
  41. """获取密码哈希"""
  42. if USE_PASSLIB:
  43. try:
  44. return pwd_context.hash(password)
  45. except Exception:
  46. # 如果 passlib 失败,使用 bcrypt 直接生成
  47. import bcrypt
  48. salt = bcrypt.gensalt()
  49. hashed = bcrypt.hashpw(password.encode('utf-8'), salt)
  50. return hashed.decode('utf-8')
  51. else:
  52. # 直接使用 bcrypt
  53. import bcrypt
  54. salt = bcrypt.gensalt()
  55. hashed = bcrypt.hashpw(password.encode('utf-8'), salt)
  56. return hashed.decode('utf-8')
  57. def create_access_token(data: dict, expires_delta: Optional[timedelta] = None) -> str:
  58. """创建访问令牌"""
  59. to_encode = data.copy()
  60. if expires_delta:
  61. expire = datetime.utcnow() + expires_delta
  62. else:
  63. expire = datetime.utcnow() + timedelta(hours=24)
  64. to_encode.update({"exp": expire})
  65. encoded_jwt = jwt.encode(to_encode, settings.SECRET_KEY, algorithm="HS256")
  66. return encoded_jwt
  67. def verify_token(token: str) -> Optional[dict]:
  68. """验证令牌"""
  69. try:
  70. payload = jwt.decode(token, settings.SECRET_KEY, algorithms=["HS256"])
  71. return payload
  72. except JWTError:
  73. return None