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

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

2 

3import jwt 

4 

5from starlette.responses import Response 

6 

7from kwai.core.settings import Settings 

8from kwai.modules.identity.tokens.refresh_token import RefreshTokenEntity 

9 

10 

11COOKIE_ACCESS_TOKEN = "access_token" 

12COOKIE_REFRESH_TOKEN = "refresh_token" 

13COOKIE_KWAI = "kwai" 

14 

15 

16def delete_cookies(response: Response): 

17 """Delete all cookies.""" 

18 response.delete_cookie(key=COOKIE_KWAI) 

19 response.delete_cookie(key=COOKIE_ACCESS_TOKEN) 

20 response.delete_cookie(key=COOKIE_REFRESH_TOKEN) 

21 

22 

23def create_cookies( 

24 response: Response, refresh_token: RefreshTokenEntity, settings: Settings 

25) -> None: 

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

27 encoded_access_token = jwt.encode( 

28 { 

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

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

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

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

33 "scope": [], 

34 }, 

35 settings.security.jwt_secret, 

36 settings.security.jwt_algorithm, 

37 ) 

38 encoded_refresh_token = jwt.encode( 

39 { 

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

41 "exp": refresh_token.expiration.timestamp, 

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

43 }, 

44 settings.security.jwt_refresh_secret, 

45 settings.security.jwt_algorithm, 

46 ) 

47 response.set_cookie( 

48 key=COOKIE_KWAI, 

49 value="Y", 

50 expires=refresh_token.expiration.timestamp, 

51 secure=settings.frontend.test, 

52 ) 

53 response.set_cookie( 

54 key=COOKIE_ACCESS_TOKEN, 

55 value=encoded_access_token, 

56 expires=refresh_token.access_token.expiration.timestamp, 

57 httponly=True, 

58 secure=not settings.frontend.test, 

59 ) 

60 response.set_cookie( 

61 key=COOKIE_REFRESH_TOKEN, 

62 value=encoded_refresh_token, 

63 expires=refresh_token.expiration.timestamp, 

64 httponly=True, 

65 secure=not settings.frontend.test, 

66 )