Coverage for fss\common\util\security.py: 94%

31 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-04-12 22:20 +0800

1"""Security util, such as encode decode and etc""" 

2 

3from datetime import timedelta, datetime 

4from typing import Any, Union 

5 

6from jose import jwt 

7from passlib.context import CryptContext 

8from starlette.responses import JSONResponse 

9 

10from fss.common.config import configs 

11 

12pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto") 

13DAYS_WITHOUT_LOGIN = 30 

14 

15 

16async def create_token( 

17 subject: Union[str, Any], expires_delta: timedelta = None, token_type: str = None 

18) -> str: 

19 if expires_delta: 

20 expire = datetime.now() + expires_delta 

21 else: 

22 expire = datetime.now() + timedelta(days=DAYS_WITHOUT_LOGIN) 

23 to_encode = {"exp": expire, "sub": str(subject), "type": token_type} 

24 encoded_jwt = jwt.encode(to_encode, configs.secret_key, algorithm=configs.algorithm) 

25 return encoded_jwt 

26 

27 

28async def verify_password(plain_password: str, hashed_password: str) -> bool: 

29 match: bool = pwd_context.verify(plain_password, hashed_password) 

30 return match 

31 

32 

33async def get_password_hash(password: str) -> str: 

34 return pwd_context.hash(password) 

35 

36 

37async def get_payload(token: str): 

38 return jwt.decode(token, configs.secret_key, algorithms=configs.algorithm) 

39 

40 

41def get_user_id(token: str): 

42 payload = jwt.decode(token, configs.secret_key, algorithms=configs.algorithm) 

43 return payload["sub"] 

44 

45 

46async def is_valid_token(token: str): 

47 payload = await get_payload(token) 

48 exp = payload.get("exp") 

49 now = datetime.now() 

50 if exp and datetime.fromtimestamp(exp) < now: 

51 return JSONResponse( 

52 status_code=401, 

53 content={"detail": "Token has expired"}, 

54 )