user.py 2.6 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758
  1. class user_util():
  2. def get_user_id(token):
  3. db = dataset.connect('mysql://choozmo:pAssw0rd@db.ptt.cx:3306/AI_anchor?charset=utf8mb4')
  4. credentials_exception = HTTPException(
  5. status_code=status.HTTP_401_UNAUTHORIZED,
  6. detail="Could not validate credentials",
  7. headers={"WWW-Authenticate": "Bearer"},
  8. )
  9. try:
  10. payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
  11. username: str = payload.get("sub")
  12. if username is None:
  13. raise credentials_exception
  14. token_data = models.TokenData(username=username)
  15. except JWTError:
  16. raise credentials_exception
  17. user = get_user(username=token_data.username)
  18. if user is None:
  19. raise credentials_exception
  20. user_id = first(db.query('SELECT * FROM users where username="' + user.username+'"'))['id']
  21. return user_id
  22. def check_user_exists(username):
  23. db = dataset.connect('mysql://choozmo:pAssw0rd@db.ptt.cx:3306/AI_anchor?charset=utf8mb4')
  24. if int(next(iter(db.query('SELECT COUNT(*) FROM AI_anchor.users WHERE username = "'+username+'"')))['COUNT(*)']) > 0:
  25. return True
  26. else:
  27. return False
  28. def get_user(username: str):
  29. db = dataset.connect('mysql://choozmo:pAssw0rd@db.ptt.cx:3306/AI_anchor?charset=utf8mb4')
  30. if not check_user_exists(username): # if user don't exist
  31. return False
  32. user_dict = next(
  33. iter(db.query('SELECT * FROM AI_anchor.users where username ="'+username+'"')))
  34. user = models.User(**user_dict)
  35. return user
  36. def user_register(user):
  37. db = dataset.connect('mysql://choozmo:pAssw0rd@db.ptt.cx:3306/AI_anchor?charset=utf8mb4')
  38. table = db['users']
  39. user.password = get_password_hash(user.password)
  40. table.insert(dict(user))
  41. def get_password_hash(password):
  42. return pwd_context.hash(password)
  43. def verify_password(plain_password, hashed_password):
  44. return pwd_context.verify(plain_password, hashed_password)
  45. def authenticate_user(username: str, password: str):
  46. db = dataset.connect('mysql://choozmo:pAssw0rd@db.ptt.cx:3306/AI_anchor?charset=utf8mb4')
  47. if not check_user_exists(username): # if user don't exist
  48. return False
  49. user_dict = next(iter(db.query('SELECT * FROM AI_anchor.users where username ="'+username+'"')))
  50. user = models.User(**user_dict)
  51. if not verify_password(password, user.password):
  52. return False
  53. return user