Coverage for apps/kwai-api/src/kwai_api/v1/auth/cookies.py: 100%

1"""Module that defines methods for handling cookies.""" 

2 

3import jwt 

4 

5from fastapi import Request 

6from kwai_bc_identity.tokens.refresh_token import RefreshTokenEntity 

7from kwai_core.settings import Settings 

8from starlette.responses import Response 

9 

10 

11COOKIE_ACCESS_TOKEN = "access_token" 

12COOKIE_REFRESH_TOKEN = "refresh_token" 

13COOKIE_KWAI = "kwai" 

14 

15 

16def use_access_token(request: Request) -> str | None: 

17 """Use access token from memory if it exists, otherwise try to get it from the cookie.""" 

18 access_token = getattr(request.state, COOKIE_ACCESS_TOKEN, None) 

19 if access_token is None: 

20 access_token = request.cookies.get(COOKIE_ACCESS_TOKEN, None) 

21 return access_token 

22 

23 

24def delete_cookies(response: Response): 

25 """Delete all cookies.""" 

26 response.delete_cookie(key=COOKIE_KWAI) 

27 response.delete_cookie(key=COOKIE_ACCESS_TOKEN) 

28 response.delete_cookie(key=COOKIE_REFRESH_TOKEN) 

29 

30 

31def create_cookies( 

32 request: Request, 

33 response: Response, 

34 refresh_token: RefreshTokenEntity, 

35 settings: Settings, 

36) -> None: 

37 """Create cookies for access en refresh token. 

38 

39 To make the access token reusable for other dependencies, it will also be stored in the state of the request. 

40 Use the use_access_token method as dependency to get the access token from the state or the cookie. 

41 """ 

42 encoded_access_token = jwt.encode( 

43 { 

44 "iat": refresh_token.access_token.traceable_time.created_at.timestamp, 

45 "exp": refresh_token.access_token.expiration.timestamp, 

46 "jti": str(refresh_token.access_token.identifier), 

47 "sub": str(refresh_token.access_token.user_account.user.uuid), 

48 "scope": [], 

49 }, 

50 settings.security.jwt_secret, 

51 settings.security.jwt_algorithm, 

52 ) 

53 setattr(request.state, COOKIE_ACCESS_TOKEN, encoded_access_token) 

54 

55 encoded_refresh_token = jwt.encode( 

56 { 

57 "iat": refresh_token.traceable_time.created_at.timestamp, 

58 "exp": refresh_token.expiration.timestamp, 

59 "jti": str(refresh_token.identifier), 

60 }, 

61 settings.security.jwt_refresh_secret, 

62 settings.security.jwt_algorithm, 

63 ) 

64 response.set_cookie( 

65 key=COOKIE_KWAI, 

66 value="Y", 

67 expires=refresh_token.expiration.timestamp, 

68 secure=settings.frontend.test, 

69 ) 

70 response.set_cookie( 

71 key=COOKIE_ACCESS_TOKEN, 

72 value=encoded_access_token, 

73 expires=refresh_token.access_token.expiration.timestamp, 

74 httponly=True, 

75 secure=not settings.frontend.test, 

76 ) 

77 response.set_cookie( 

78 key=COOKIE_REFRESH_TOKEN, 

79 value=encoded_refresh_token, 

80 expires=refresh_token.expiration.timestamp, 

81 httponly=True, 

82 secure=not settings.frontend.test, 

83 )