Coverage for src/kwai/modules/identity/refresh_access_token.py: 68%

1"""Module that defines the use case for refreshing an access token.""" 

2 

3from dataclasses import dataclass 

4 

5from kwai.modules.identity.authenticate_user import AuthenticationException 

6from kwai.modules.identity.tokens.access_token_repository import AccessTokenRepository 

7from kwai.modules.identity.tokens.log_user_login_service import LogUserLoginService 

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

9from kwai.modules.identity.tokens.refresh_token_repository import RefreshTokenRepository 

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

11 

12 

13@dataclass(kw_only=True, frozen=True, slots=True) 

14class RefreshAccessTokenCommand: 

15 """Input for the refresh access token use case. 

16 

17 Attributes: 

18 identifier: The identifier of the refresh token. 

19 access_token_expiry_minutes: Minutes before expiring the access token. 

20 Default is 2 hours. 

21 refresh_token_expiry_minutes: Minutes before expiring the refresh token. 

22 Default is 2 months. 

23 """ 

24 

25 identifier: str # The identifier of the refresh token. 

26 access_token_expiry_minutes: int = 60 * 2 # 2 hours 

27 refresh_token_expiry_minutes: int = 60 * 24 * 60 # 2 months 

28 

29 

30class RefreshAccessToken: 

31 """Use case for refreshing an access token. 

32 

33 Attributes: 

34 _refresh_token_repo (RefreshTokenRepository): The repo for getting and creating 

35 a new refresh token. 

36 _access_token_repo (AccessTokenRepository): The repo for updating and creating 

37 an access token. 

38 

39 Note: 

40 A new access token will also result in a new refresh token. 

41 """ 

42 

43 def __init__( 

44 self, 

45 refresh_token_repo: RefreshTokenRepository, 

46 access_token_repo: AccessTokenRepository, 

47 log_user_login_service: LogUserLoginService, 

48 ): 

49 self._refresh_token_repo = refresh_token_repo 

50 self._access_token_repo = access_token_repo 

51 self._log_user_login_service = log_user_login_service 

52 

53 async def execute(self, command: RefreshAccessTokenCommand) -> RefreshTokenEntity: 

54 """Execute the use case. 

55 

56 Args: 

57 command: The input for this use case. 

58 

59 Raises: 

60 RefreshTokenNotFoundException: Raised when the refresh token does not exist. 

61 AuthenticationException: Raised when the refresh token is expired, the 

62 refresh token is revoked or the user is revoked. 

63 """ 

64 refresh_token = await self._refresh_token_repo.get_by_token_identifier( 

65 TokenIdentifier(hex_string=command.identifier) 

66 ) 

67 

68 if refresh_token.expired: 

69 message = "Refresh token is expired" 

70 await self._log_user_login_service.notify_failure( 

71 message, 

72 user_account=refresh_token.access_token.user_account, 

73 refresh_token=refresh_token, 

74 ) 

75 raise AuthenticationException(message) 

76 

77 if refresh_token.revoked: 

78 message = "Refresh token is revoked" 

79 await self._log_user_login_service.notify_failure( 

80 message, 

81 user_account=refresh_token.access_token.user_account, 

82 refresh_token=refresh_token, 

83 ) 

84 raise AuthenticationException(message) 

85 

86 # When the user is revoked, revoke the access and refresh tokens. 

87 if refresh_token.access_token.user_account.revoked: 

88 refresh_token.revoke() 

89 await self._refresh_token_repo.update(refresh_token) 

90 # The access token is also revoked, so update it 

91 await self._access_token_repo.update(refresh_token.access_token) 

92 

93 message = "User is revoked" 

94 await self._log_user_login_service.notify_failure( 

95 message, 

96 user_account=refresh_token.access_token.user_account, 

97 refresh_token=refresh_token, 

98 ) 

99 raise AuthenticationException(message) 

100 

101 # Renew the refresh token 

102 refresh_token = refresh_token.renew( 

103 command.refresh_token_expiry_minutes, command.access_token_expiry_minutes 

104 ) 

105 

106 await self._refresh_token_repo.update(refresh_token) 

107 await self._access_token_repo.update(refresh_token.access_token) 

108 

109 return refresh_token