|
| 1 | +import os |
| 2 | +from _sha256 import sha256 |
| 3 | +from time import time |
| 4 | + |
| 5 | +from redis import Redis |
| 6 | + |
| 7 | +# 会话默认过期时间,一个月 |
| 8 | +DEFAULT_TIMEOUT = 30 * 24 * 3600 |
| 9 | + |
| 10 | +# 会话 token 及过期时间的 key |
| 11 | +SESSION_TOKEN_KEY = 'SESSION:TOKEN' |
| 12 | +SESSION_EXPIRE_TS_KEY = 'SESSION:EXPIRE' |
| 13 | + |
| 14 | +# 会话状态 |
| 15 | +SESSION_NOT_LOGIN = 'SESSION_NOT_LOGIN' |
| 16 | +SESSION_EXPIRE = 'SESSION_EXPIRE' |
| 17 | +SESSION_TOKEN_CORRECT = 'SESSION_TOKEN_CORRECT' |
| 18 | +SESSION_TOKEN_INCORRECT = 'SESSION_TOKEN_INCORRECT' |
| 19 | + |
| 20 | + |
| 21 | +def generate_token(): |
| 22 | + """生成一个随机的会话令牌""" |
| 23 | + return sha256(os.urandom(256)).hexdigest() |
| 24 | + |
| 25 | + |
| 26 | +class LoginSession: |
| 27 | + def __init__(self, client: Redis, user_id: str): |
| 28 | + self.client = client |
| 29 | + self.user_id = user_id |
| 30 | + |
| 31 | + def create(self, timeout=DEFAULT_TIMEOUT) -> str: |
| 32 | + """创建新的会话,并返回会话token""" |
| 33 | + session_token = generate_token() |
| 34 | + |
| 35 | + # 设置过期时间 |
| 36 | + expire_time = time() + timeout |
| 37 | + self.client.hset(SESSION_TOKEN_KEY, self.user_id, session_token) |
| 38 | + self.client.hset(SESSION_EXPIRE_TS_KEY, self.user_id, expire_time) |
| 39 | + return session_token |
| 40 | + |
| 41 | + def validate(self, token) -> str: |
| 42 | + """校验token""" |
| 43 | + session_token = self.client.hget(SESSION_TOKEN_KEY, self.user_id) |
| 44 | + expire_time = self.client.hget(SESSION_EXPIRE_TS_KEY, self.user_id) |
| 45 | + |
| 46 | + if (session_token is None) or (expire_time is None): |
| 47 | + return SESSION_NOT_LOGIN |
| 48 | + |
| 49 | + # 将字符串类型的时间转换为浮点数类型 |
| 50 | + if time() > float(expire_time): |
| 51 | + return SESSION_EXPIRE |
| 52 | + |
| 53 | + if session_token == token: |
| 54 | + return SESSION_TOKEN_CORRECT |
| 55 | + |
| 56 | + return SESSION_TOKEN_INCORRECT |
| 57 | + |
| 58 | + def destroy(self): |
| 59 | + """销毁会话""" |
| 60 | + self.client.hdel(SESSION_TOKEN_KEY, self.user_id) |
| 61 | + self.client.hdel(SESSION_EXPIRE_TS_KEY, self.user_id) |
0 commit comments