Coverage for chalicelib/core/mcp/authorizers.py: 32%

1import logging 

2 

3import jwt 

4from decouple import config 

5from fastapi import HTTPException 

6 

7from chalicelib.utils import pg_client 

8from chalicelib.utils.TimeUTC import TimeUTC 

9from chalicelib.utils.log import sanitize 

10from schemas import schemas, MCP 

11 

12logger = logging.getLogger(__name__) 

13 

14AUDIENCE = "mcp:OpenReplay" 

15MCP_LOGIN_TIMEOUT_S = config("MCP_LOGIN_TIMEOUT_S", cast=int, default=0) 

16 

17 

18def get_supported_audience(): 

19 return [AUDIENCE] 

20 

21 

22def is_mcp_token(token: str) -> bool: 

23 try: 

24 if len(token) < 5 or "." not in token: 

25 return False 

26 decoded_token = jwt.decode(token, options={"verify_signature": False, "verify_exp": False}) 

27 audience = decoded_token.get("aud") 

28 return audience == AUDIENCE 

29 except jwt.InvalidTokenError: 

30 logger.error(f"Invalid token for is_spot_token: {sanitize(token, max_length=16)}...") 

31 raise 

32 

33 

34def jwt_authorizer(scheme: str, token: str, leeway=0) -> dict | None: 

35 if scheme.lower() != "bearer" or len(token) < 5: 

36 return None 

37 try: 

38 payload = jwt.decode(jwt=token, 

39 key=config("JWT_MCP_SECRET"), 

40 algorithms=config("JWT_MCP_ALGORITHM"), 

41 audience=get_supported_audience(), 

42 leeway=leeway) 

43 except jwt.ExpiredSignatureError: 

44 logger.debug("! JWT Expired signature") 

45 return None 

46 except jwt.exceptions.InvalidSignatureError: 

47 logger.warning("! JWT Signature verification failed") 

48 return None 

49 except BaseException as e: 

50 logger.warning("! JWT Base Exception", exc_info=e) 

51 return None 

52 return payload 

53 

54 

55def generate_jwt(user_id, tenant_id, iat, jti, client_id): 

56 token = jwt.encode( 

57 payload={ 

58 "userId": user_id, 

59 "tenantId": tenant_id, 

60 "exp": iat + config("JWT_MCP_EXPIRATION", cast=int), 

61 "iss": config("JWT_ISSUER"), 

62 "iat": iat, 

63 "jti": jti, 

64 "aud": AUDIENCE, 

65 "clientId": client_id 

66 }, 

67 key=config("JWT_MCP_SECRET"), 

68 algorithm=config("JWT_MCP_ALGORITHM") 

69 ) 

70 return token 

71 

72 

73def store_token_request(data: MCP.AuthorizeSchema, cotext: schemas.CurrentContext): 

74 with pg_client.PostgresClient() as cur: 

75 query = cur.mogrify( 

76 """ 

77 INSERT INTO public.mcp_authentication_tokens(user_id, client_id, state, iat) 

78 VALUES (%(user_id)s, %(client_id)s, %(state)s, 

79 timezone('utc'::text, now() - INTERVAL '10s')) ON CONFLICT DO NOTHING;""", 

80 {"client_id": data.client_id, "state": data.state, 

81 "user_id": cotext.user_id, }, 

82 ) 

83 cur.execute(query=query) 

84 

85 

86def get_token_by_state(client_id, state): 

87 with pg_client.PostgresClient() as cur: 

88 query = cur.mogrify( 

89 """ 

90 UPDATE public.mcp_authentication_tokens 

91 SET generated= TRUE 

92 WHERE client_id = %(client_id)s 

93 AND state = %(state)s 

94 AND NOT generated  

95 RETURNING *,EXTRACT(epoch FROM iat)::BIGINT AS iat, 

96 (SELECT tenant_id  

97 FROM public.users  

98 WHERE users.user_id = mcp_authentication_tokens.user_id  

99 AND users.deleted_at IS NULL) AS tenant_id;""", 

100 {"client_id": client_id, "state": state}, 

101 ) 

102 cur.execute(query=query) 

103 row = cur.fetchone() 

104 

105 if row is None: 

106 raise HTTPException(status_code=404, detail="State not found or token already generated") 

107 if MCP_LOGIN_TIMEOUT_S > 0 and (row["iat"] + MCP_LOGIN_TIMEOUT_S) >= TimeUTC.now(): 

108 raise HTTPException(status_code=408, detail="Login timed out") 

109 

110 return generate_jwt(user_id=row["user_id"], tenant_id=row["tenant_id"], iat=row["iat"], client_id=client_id, jti=row["jti"])