Coverage for polar/oauth2/schemas.py: 90%

1import ipaddress 

2import re 

3from typing import Annotated, Any, Literal 

4 

5from fastapi.openapi.constants import REF_TEMPLATE 

6from pydantic import ( 

7 UUID4, 

8 AfterValidator, 

9 AnyUrl, 

10 BeforeValidator, 

11 Discriminator, 

12 EmailStr, 

13 Field, 

14 HttpUrl, 

15 TypeAdapter, 

16) 

17 

18from polar.auth.scope import SCOPES_SUPPORTED, Scope, scope_to_list 

19from polar.kit.schemas import Schema, TimestampedSchema 

20 

21from .sub_type import SubType 

22 

23_LOCALHOST_HOST_PATTERN = re.compile(r"([^\.]+\.)?localhost(\d+)?", flags=re.IGNORECASE) 

24 

25 

26def _is_localhost(host: str) -> bool: 

27 try: 

28 return ipaddress.IPv4Address(host).is_private 

29 except ValueError: 

30 return _LOCALHOST_HOST_PATTERN.match(host) is not None 

31 

32 

33def _is_https_or_localhost(value: HttpUrl) -> HttpUrl: 

34 if value.scheme == "http" and (value.host is None or not _is_localhost(value.host)): 34 ↛ 35line 34 didn't jump to line 35 because the condition on line 34 was never true

35 raise ValueError("An HTTPS URL is required.") 

36 return value 

37 

38 

39URIOrLocalhost = Annotated[AnyUrl, AfterValidator(_is_https_or_localhost)] 

40Scopes = Annotated[list[Scope], BeforeValidator(scope_to_list)] 

41 

42 

43class OAuth2ClientConfiguration(Schema): 

44 redirect_uris: list[URIOrLocalhost] 

45 token_endpoint_auth_method: Literal[ 

46 "client_secret_basic", "client_secret_post", "none" 

47 ] = "client_secret_post" 

48 grant_types: list[Literal["authorization_code", "refresh_token"]] = [ 

49 "authorization_code", 

50 "refresh_token", 

51 ] 

52 response_types: list[Literal["code"]] = ["code"] 

53 scope: str = " ".join(SCOPES_SUPPORTED) 

54 client_name: str 

55 client_uri: str | None = None 

56 logo_uri: HttpUrl | None = None 

57 tos_uri: HttpUrl | None = None 

58 policy_uri: HttpUrl | None = None 

59 default_sub_type: SubType = SubType.organization 

60 

61 

62class OAuth2ClientConfigurationUpdate(OAuth2ClientConfiguration): 

63 client_id: str 

64 

65 

66class OAuth2Client(TimestampedSchema, OAuth2ClientConfiguration): 

67 client_id: str 

68 client_secret: str 

69 client_id_issued_at: int 

70 client_secret_expires_at: int 

71 

72 

73class OAuth2ClientPublic(TimestampedSchema): 

74 client_id: str 

75 client_name: str | None 

76 client_uri: str | None 

77 logo_uri: str | None 

78 tos_uri: str | None 

79 policy_uri: str | None 

80 

81 

82class AuthorizeUser(Schema): 

83 id: UUID4 

84 email: EmailStr 

85 avatar_url: str | None 

86 

87 

88class AuthorizeOrganization(Schema): 

89 id: UUID4 

90 slug: str 

91 avatar_url: str | None 

92 

93 

94class AuthorizeResponseBase(Schema): 

95 client: OAuth2ClientPublic 

96 sub_type: SubType 

97 sub: AuthorizeUser | AuthorizeOrganization | None 

98 scopes: Scopes 

99 

100 

101class AuthorizeResponseUser(AuthorizeResponseBase): 

102 sub_type: Literal[SubType.user] 

103 sub: AuthorizeUser | None 

104 

105 

106class AuthorizeResponseOrganization(AuthorizeResponseBase): 

107 sub_type: Literal[SubType.organization] 

108 sub: AuthorizeOrganization | None 

109 organizations: list[AuthorizeOrganization] 

110 

111 

112AuthorizeResponse = Annotated[ 

113 AuthorizeResponseUser | AuthorizeResponseOrganization, 

114 Discriminator(discriminator="sub_type"), 

115] 

116 

117authorize_response_adapter: TypeAdapter[AuthorizeResponse] = TypeAdapter( 

118 AuthorizeResponse 

119) 

120 

121 

122class TokenRequestBase(Schema): 

123 grant_type: Literal["authorization_code", "refresh_token", "web"] 

124 client_id: str 

125 client_secret: str 

126 

127 

128class AuthorizationCodeTokenRequest(TokenRequestBase): 

129 grant_type: Literal["authorization_code"] 

130 code: str 

131 redirect_uri: HttpUrl 

132 

133 

134class RefreshTokenRequest(TokenRequestBase): 

135 grant_type: Literal["refresh_token"] 

136 

137 refresh_token: str 

138 

139 

140class WebTokenRequest(TokenRequestBase): 

141 grant_type: Literal["web"] 

142 session_token: str 

143 sub_type: Literal["user", "organization"] = Field(default="user") 

144 sub: UUID4 | None = None 

145 scope: str | None = Field(default=None) 

146 

147 

148class TokenResponse(Schema): 

149 access_token: str 

150 token_type: Literal["Bearer"] 

151 expires_in: int 

152 refresh_token: str | None 

153 scope: str 

154 id_token: str 

155 

156 

157class RevokeTokenRequest(Schema): 

158 token: str 

159 token_type_hint: Literal["access_token", "refresh_token"] | None = None 

160 client_id: str 

161 client_secret: str 

162 

163 

164class RevokeTokenResponse(Schema): ... 164 ↛ 167line 164 didn't jump to line 167 because

165 

166 

167class IntrospectTokenRequest(Schema): 

168 token: str 

169 token_type_hint: Literal["access_token", "refresh_token"] | None = None 

170 client_id: str 

171 client_secret: str 

172 

173 

174class IntrospectTokenResponse(Schema): 

175 active: bool 

176 client_id: str 

177 token_type: Literal["access_token", "refresh_token"] 

178 scope: str 

179 sub_type: SubType 

180 sub: str 

181 aud: str 

182 iss: str 

183 exp: int 

184 iat: int 

185 

186 

187class UserInfoUser(Schema): 

188 sub: str 

189 name: str | None = None 

190 email: str | None = None 

191 email_verified: bool | None = None 

192 

193 

194class UserInfoOrganization(Schema): 

195 sub: str 

196 name: str | None = None 

197 

198 

199UserInfo = UserInfoUser | UserInfoOrganization 

200 

201 

202def add_oauth2_form_schemas(openapi_schema: dict[str, Any]) -> dict[str, Any]: 

203 openapi_schema["components"]["schemas"]["AuthorizationCodeTokenRequest"] = ( 

204 AuthorizationCodeTokenRequest.model_json_schema(ref_template=REF_TEMPLATE) 

205 ) 

206 openapi_schema["components"]["schemas"]["RefreshTokenRequest"] = ( 

207 RefreshTokenRequest.model_json_schema(ref_template=REF_TEMPLATE) 

208 ) 

209 openapi_schema["components"]["schemas"]["WebTokenRequest"] = ( 

210 WebTokenRequest.model_json_schema(ref_template=REF_TEMPLATE) 

211 ) 

212 openapi_schema["components"]["schemas"]["RevokeTokenRequest"] = ( 

213 RevokeTokenRequest.model_json_schema(ref_template=REF_TEMPLATE) 

214 ) 

215 openapi_schema["components"]["schemas"]["IntrospectTokenRequest"] = ( 

216 IntrospectTokenRequest.model_json_schema(ref_template=REF_TEMPLATE) 

217 ) 

218 return openapi_schema