user.py 2.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172
  1. from core.db import WordDatabase
  2. from werkzeug.security import generate_password_hash, check_password_hash
  3. import os
  4. from configure import conf
  5. from flask_login import UserMixin, AnonymousUserMixin
  6. import shutil
  7. class AnonymousUser(AnonymousUserMixin):
  8. ...
  9. class UserWordDataBase(WordDatabase, UserMixin):
  10. def __init__(self, user: str, path: str):
  11. super().__init__(user, path)
  12. self.done(f'''
  13. CREATE TABLE IF NOT EXISTS User (
  14. id INTEGER PRIMARY KEY AUTOINCREMENT, -- 记录ID
  15. passwd TEXT NOT NULL -- 密码hash
  16. )''')
  17. self.insert(table="User", columns=["passwd"], values=f"'{generate_password_hash('88888888')}'") # 默认密码
  18. self.user = user
  19. def get_id(self):
  20. return self.user
  21. def check_passwd(self, passwd: str) -> bool:
  22. res = self.search(table="User", columns=["passwd"], limit=1, order_by=[("ID", "ASC")])
  23. if len(res) == 0:
  24. return False
  25. return check_password_hash(res[0][0], passwd)
  26. def set_passwd(self, passwd: str, record_id: int = 1):
  27. self.update(table="User", kw={"passwd": f"'{generate_password_hash(passwd)}'"}, where=f"id = {record_id}")
  28. def delete_user(self):
  29. self.delete_self()
  30. def check_base_db():
  31. if os.path.exists(os.path.join(conf["DB_TEMPLATE"], "base.db")):
  32. return
  33. WordDatabase("base", conf["DB_TEMPLATE"])
  34. def check_template(template: str) -> bool:
  35. check_base_db()
  36. return os.path.exists(os.path.join(conf["DB_TEMPLATE"], f"{template}.db"))
  37. def create_user(template: str, name: str, passwd: str):
  38. check_base_db()
  39. if not os.path.exists(os.path.join(conf["DB_TEMPLATE"], f"{template}.db")):
  40. return 0, None
  41. if os.path.exists(os.path.join(conf["DB_PATH"], f"{name}.db")):
  42. return -1, None
  43. shutil.copy(os.path.join(conf["DB_TEMPLATE"], f"{template}.db"), os.path.join(conf["DB_PATH"], f"{name}.db"))
  44. user = UserWordDataBase(name, conf["DB_PATH"])
  45. if len(passwd) > 0:
  46. user.set_passwd(passwd)
  47. return 1, user
  48. def load_user(name: str, passwd: str | None):
  49. if not os.path.exists(os.path.join(conf["DB_PATH"], f"{name}.db")):
  50. return None
  51. user = UserWordDataBase(name, conf["DB_PATH"])
  52. if passwd is None or user.check_passwd(passwd):
  53. return user
  54. return None