test / auth /jwt_handler.py
gaoqilan's picture
Upload 103 files
1f1b4db verified
from utils.http import HTTPClient
from config.constants import APP_LANGUAGE, BASE_URL, APP_NAME, APP_VERSION, DISPLAY_NAME
from protos import jwt_pb2
class JWTHandler:
def __init__(self, api_key: str, server_url: str = BASE_URL, http_client: HTTPClient = HTTPClient()):
self.server_url = server_url
self.api_key = api_key
self.jwt_token = None
self.http_client = http_client
async def get_jwt_token(self) -> str:
"""Get JWT token from server"""
request = jwt_pb2.jwt_request(
surfwind_jwt_request=jwt_pb2.surfwind_jwt_request(
app_name=APP_NAME,
version=APP_VERSION,
api_key=self.api_key,
language=APP_LANGUAGE,
display_name=DISPLAY_NAME,
app_identifier=APP_NAME,
)
)
response = await self.http_client.post(
url=f"{self.server_url}/exa.auth_pb.AuthService/GetUserJwt",
headers={
"Content-Type": "application/proto",
},
data=request.SerializeToString(),
)
if response.status_code == 200:
jwt_response = jwt_pb2.jwt_response()
jwt_response.ParseFromString(response.content)
self.jwt_token = jwt_response.jwt_token
return self.jwt_token
print(response.content)
raise Exception(f"Failed to get JWT token: {response.status_code}")