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 kwai_bc_identity.tokens.refresh_token import RefreshTokenEntity 

6from kwai_core.settings import Settings 

7from starlette.responses import Response 

8 

9 

10COOKIE_ACCESS_TOKEN = "access_token" 

11COOKIE_REFRESH_TOKEN = "refresh_token" 

12COOKIE_KWAI = "kwai" 

13 

14 

15def delete_cookies(response: Response): 

16 """Delete all cookies.""" 

17 response.delete_cookie(key=COOKIE_KWAI) 

18 response.delete_cookie(key=COOKIE_ACCESS_TOKEN) 

19 response.delete_cookie(key=COOKIE_REFRESH_TOKEN) 

20 

21 

22def create_cookies( 

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

24) -> None: 

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

26 encoded_access_token = jwt.encode( 

27 { 

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

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

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

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

32 "scope": [], 

33 }, 

34 settings.security.jwt_secret, 

35 settings.security.jwt_algorithm, 

36 ) 

37 encoded_refresh_token = jwt.encode( 

38 { 

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

40 "exp": refresh_token.expiration.timestamp, 

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

42 }, 

43 settings.security.jwt_refresh_secret, 

44 settings.security.jwt_algorithm, 

45 ) 

46 response.set_cookie( 

47 key=COOKIE_KWAI, 

48 value="Y", 

49 expires=refresh_token.expiration.timestamp, 

50 secure=settings.frontend.test, 

51 ) 

52 response.set_cookie( 

53 key=COOKIE_ACCESS_TOKEN, 

54 value=encoded_access_token, 

55 expires=refresh_token.access_token.expiration.timestamp, 

56 httponly=True, 

57 secure=not settings.frontend.test, 

58 ) 

59 response.set_cookie( 

60 key=COOKIE_REFRESH_TOKEN, 

61 value=encoded_refresh_token, 

62 expires=refresh_token.expiration.timestamp, 

63 httponly=True, 

64 secure=not settings.frontend.test, 

65 )