Coverage for kwai/api/dependencies.py: 79%

1"""Module that integrates the dependencies in FastAPI.""" 

2import jwt 

3from fastapi import Depends, HTTPException, status 

4from fastapi.security import OAuth2PasswordBearer 

5from lagom.integrations.fast_api import FastApiIntegration 

6 

7from kwai.core.db.database import Database 

8from kwai.core.dependencies import container 

9from kwai.core.settings import Settings 

10from kwai.modules.identity.tokens.access_token_db_repository import ( 

11 AccessTokenDbRepository, 

12) 

13from kwai.modules.identity.tokens.access_token_repository import ( 

14 AccessTokenNotFoundException, 

15) 

16from kwai.modules.identity.tokens.token_identifier import TokenIdentifier 

17from kwai.modules.identity.users.user import UserEntity 

18 

19deps = FastApiIntegration(container) 

20 

21oauth = OAuth2PasswordBearer(tokenUrl="/api/v1/auth/login") 

22 

23 

24async def get_current_user( 

25 settings=deps.depends(Settings), 

26 db=deps.depends(Database), 

27 token: str = Depends(oauth), 

28) -> UserEntity: 

29 """Try to get the current user from the access token. 

30 

31 Not authorized will be raised when the access token is not found, expired, revoked 

32 or when the user is revoked. 

33 """ 

34 payload = jwt.decode( 

35 token, 

36 settings.security.jwt_secret, 

37 algorithms=[settings.security.jwt_algorithm], 

38 ) 

39 access_token_repo = AccessTokenDbRepository(db) 

40 try: 

41 access_token = await access_token_repo.get_by_identifier( 

42 TokenIdentifier(hex_string=payload["jti"]) 

43 ) 

44 except AccessTokenNotFoundException as exc: 

45 raise HTTPException(status.HTTP_401_UNAUTHORIZED) from exc 

46 

47 # Check if the access token is assigned to the user we have in the subject of JWT. 

48 if not access_token.user_account.user.uuid == payload["sub"]: 

49 raise HTTPException(status.HTTP_401_UNAUTHORIZED) 

50 

51 if access_token.revoked: 

52 raise HTTPException(status.HTTP_401_UNAUTHORIZED) 

53 

54 if access_token.user_account.revoked: 

55 raise HTTPException(status.HTTP_401_UNAUTHORIZED) 

56 

57 if access_token.expired: 

58 raise HTTPException(status.HTTP_401_UNAUTHORIZED) 

59 

60 return access_token.user_account.user