Upload 18 files

warp2protobuf/core/auth.py

@@ -15,22 +15,29 @@ import httpx
 import asyncio
 from dotenv import load_dotenv, set_key
 
-from ..config.settings import REFRESH_TOKEN_B64, REFRESH_URL, CLIENT_VERSION, OS_CATEGORY, OS_NAME, OS_VERSION
+from ..config.settings import (
+    REFRESH_TOKEN_B64,
+    REFRESH_URL,
+    CLIENT_VERSION,
+    OS_CATEGORY,
+    OS_NAME,
+    OS_VERSION,
+)
 from .logging import logger, log
 
 
 def decode_jwt_payload(token: str) -> dict:
     """Decode JWT payload to check expiration"""
     try:
-        parts = token.split('.')
+        parts = token.split(".")
         if len(parts) != 3:
             return {}
         payload_b64 = parts[1]
         padding = 4 - len(payload_b64) % 4
         if padding != 4:
-            payload_b64 += '=' * padding
+            payload_b64 += "=" * padding
         payload_bytes = base64.urlsafe_b64decode(payload_b64)
-        payload = json.loads(payload_bytes.decode('utf-8'))
+        payload = json.loads(payload_bytes.decode("utf-8"))
         return payload
     except Exception as e:
         logger.debug(f"Error decoding JWT: {e}")
@@ -39,9 +46,9 @@ def decode_jwt_payload(token: str) -> dict:
 
 def is_token_expired(token: str, buffer_minutes: int = 5) -> bool:
     payload = decode_jwt_payload(token)
-    if not payload or 'exp' not in payload:
+    if not payload or "exp" not in payload:
         return True
-    expiry_time = payload['exp']
+    expiry_time = payload["exp"]
     current_time = time.time()
     buffer_time = buffer_minutes * 60
     return (expiry_time - current_time) <= buffer_time
@@ -57,7 +64,9 @@ async def refresh_jwt_token() -> dict:
     # Prefer dynamic refresh token from environment if present
     env_refresh = os.getenv("WARP_REFRESH_TOKEN")
     if env_refresh:
-        payload = f"grant_type=refresh_token&refresh_token={env_refresh}".encode("utf-8")
+        payload = f"grant_type=refresh_token&refresh_token={env_refresh}".encode(
+            "utf-8"
+        )
     else:
         payload = base64.b64decode(REFRESH_TOKEN_B64)
     headers = {
@@ -68,15 +77,11 @@ async def refresh_jwt_token() -> dict:
         "content-type": "application/x-www-form-urlencoded",
         "accept": "*/*",
         "accept-encoding": "gzip, br",
-        "content-length": str(len(payload))
+        "content-length": str(len(payload)),
     }
     try:
         async with httpx.AsyncClient(timeout=30.0) as client:
-            response = await client.post(
-                REFRESH_URL,
-                headers=headers,
-                content=payload
-            )
+            response = await client.post(REFRESH_URL, headers=headers, content=payload)
             if response.status_code == 200:
                 token_data = response.json()
                 logger.info("Token refresh successful")
@@ -91,9 +96,11 @@ async def refresh_jwt_token() -> dict:
 
 
 def update_env_file(new_jwt: str) -> bool:
-    env_path = Path(".env")
+    # 如果在Hugging Face环境（/tmp/.env存在），就用它，否则用本地的.env
+    ENV_FILE_PATH = "/tmp/.env" if os.path.exists("/tmp/.env") else ".env"
+    # env_path = Path(".env")
     try:
-        set_key(str(env_path), "WARP_JWT", new_jwt)
+        set_key(str(ENV_FILE_PATH), "WARP_JWT", new_jwt)
         logger.info("Updated .env file with new JWT token")
         return True
     except Exception as e:
@@ -102,9 +109,12 @@ def update_env_file(new_jwt: str) -> bool:
 
 
 def update_env_refresh_token(refresh_token: str) -> bool:
-    env_path = Path(".env")
+    # 如果在Hugging Face环境（/tmp/.env存在），就用它，否则用本地的.env
+    ENV_FILE_PATH = "/tmp/.env" if os.path.exists("/tmp/.env") else ".env"
+    # env_path = Path(".env")
+
     try:
-        set_key(str(env_path), "WARP_REFRESH_TOKEN", refresh_token)
+        set_key(str(ENV_FILE_PATH), "WARP_REFRESH_TOKEN", refresh_token)
         logger.info("Updated .env with WARP_REFRESH_TOKEN")
         return True
     except Exception as e:
@@ -137,11 +147,13 @@ async def check_and_refresh_token() -> bool:
             return False
     else:
         payload = decode_jwt_payload(current_jwt)
-        if payload and 'exp' in payload:
-            expiry_time = payload['exp']
+        if payload and "exp" in payload:
+            expiry_time = payload["exp"]
             time_left = expiry_time - time.time()
             hours_left = time_left / 3600
-            logger.debug(f"Current token is still valid ({hours_left:.1f} hours remaining)")
+            logger.debug(
+                f"Current token is still valid ({hours_left:.1f} hours remaining)"
+            )
         else:
             logger.debug("Current token appears valid")
         return True
@@ -149,6 +161,7 @@ async def check_and_refresh_token() -> bool:
 
 async def get_valid_jwt() -> str:
     from dotenv import load_dotenv as _load
+
     _load(override=True)
     jwt = os.getenv("WARP_JWT")
     if not jwt:
@@ -164,14 +177,19 @@ async def get_valid_jwt() -> str:
             _load(override=True)
             jwt = os.getenv("WARP_JWT")
             if not jwt or is_token_expired(jwt, buffer_minutes=0):
-                logger.warning("Warning: New token has short expiry but proceeding anyway")
+                logger.warning(
+                    "Warning: New token has short expiry but proceeding anyway"
+                )
         else:
-            logger.warning("Warning: JWT token refresh failed, trying to use existing token")
+            logger.warning(
+                "Warning: JWT token refresh failed, trying to use existing token"
+            )
     return jwt
 
 
 def get_jwt_token() -> str:
     from dotenv import load_dotenv as _load
+
     _load()
     return os.getenv("WARP_JWT", "")
 
@@ -187,13 +205,16 @@ async def refresh_jwt_if_needed() -> bool:
 # ============ Anonymous token acquisition (quota refresh) ============
 
 _ANON_GQL_URL = "https://app.warp.dev/graphql/v2?op=CreateAnonymousUser"
-_IDENTITY_TOOLKIT_BASE = "https://identitytoolkit.googleapis.com/v1/accounts:signInWithCustomToken"
+_IDENTITY_TOOLKIT_BASE = (
+    "https://identitytoolkit.googleapis.com/v1/accounts:signInWithCustomToken"
+)
 
 
 def _extract_google_api_key_from_refresh_url() -> str:
     try:
         # REFRESH_URL like: https://app.warp.dev/proxy/token?key=API_KEY
         from urllib.parse import urlparse, parse_qs
+
         parsed = urlparse(REFRESH_URL)
         qs = parse_qs(parsed.query)
         key = qs.get("key", [""])[0]
@@ -235,7 +256,7 @@ async def _create_anonymous_user() -> dict:
         "input": {
             "anonymousUserType": "NATIVE_CLIENT_ANONYMOUS_USER_FEATURE_GATED",
             "expirationType": "NO_EXPIRATION",
-            "referralCode": None
+            "referralCode": None,
         },
         "requestContext": {
             "clientContext": {"version": CLIENT_VERSION},
@@ -244,21 +265,31 @@ async def _create_anonymous_user() -> dict:
                 "linuxKernelVersion": None,
                 "name": OS_NAME,
                 "version": OS_VERSION,
-            }
-        }
+            },
+        },
+    }
+    body = {
+        "query": query,
+        "variables": variables,
+        "operationName": "CreateAnonymousUser",
     }
-    body = {"query": query, "variables": variables, "operationName": "CreateAnonymousUser"}
     async with httpx.AsyncClient(timeout=httpx.Timeout(30.0)) as client:
         resp = await client.post(_ANON_GQL_URL, headers=headers, json=body)
         if resp.status_code != 200:
-            raise RuntimeError(f"CreateAnonymousUser failed: HTTP {resp.status_code} {resp.text[:200]}")
+            raise RuntimeError(
+                f"CreateAnonymousUser failed: HTTP {resp.status_code} {resp.text[:200]}"
+            )
         data = resp.json()
         return data
 
 
 async def _exchange_id_token_for_refresh_token(id_token: str) -> dict:
     key = _extract_google_api_key_from_refresh_url()
-    url = f"{_IDENTITY_TOOLKIT_BASE}?key={key}" if key else f"{_IDENTITY_TOOLKIT_BASE}?key=AIzaSyBdy3O3S9hrdayLJxJ7mriBR4qgUaUygAs"
+    url = (
+        f"{_IDENTITY_TOOLKIT_BASE}?key={key}"
+        if key
+        else f"{_IDENTITY_TOOLKIT_BASE}?key=AIzaSyBdy3O3S9hrdayLJxJ7mriBR4qgUaUygAs"
+    )
     headers = {
         "accept-encoding": "gzip, br",
         "content-type": "application/x-www-form-urlencoded",
@@ -274,7 +305,9 @@ async def _exchange_id_token_for_refresh_token(id_token: str) -> dict:
     async with httpx.AsyncClient(timeout=httpx.Timeout(30.0)) as client:
         resp = await client.post(url, headers=headers, data=form)
         if resp.status_code != 200:
-            raise RuntimeError(f"signInWithCustomToken failed: HTTP {resp.status_code} {resp.text[:200]}")
+            raise RuntimeError(
+                f"signInWithCustomToken failed: HTTP {resp.status_code} {resp.text[:200]}"
+            )
         return resp.json()
 
 
@@ -296,7 +329,9 @@ async def acquire_anonymous_access_token() -> str:
     signin = await _exchange_id_token_for_refresh_token(id_token)
     refresh_token = signin.get("refreshToken")
     if not refresh_token:
-        raise RuntimeError(f"signInWithCustomToken did not return refreshToken: {signin}")
+        raise RuntimeError(
+            f"signInWithCustomToken did not return refreshToken: {signin}"
+        )
 
     # Persist refresh token for future time-based refreshes
     update_env_refresh_token(refresh_token)
@@ -311,12 +346,14 @@ async def acquire_anonymous_access_token() -> str:
         "content-type": "application/x-www-form-urlencoded",
         "accept": "*/*",
         "accept-encoding": "gzip, br",
-        "content-length": str(len(payload))
+        "content-length": str(len(payload)),
     }
     async with httpx.AsyncClient(timeout=httpx.Timeout(30.0)) as client:
         resp = await client.post(REFRESH_URL, headers=headers, content=payload)
         if resp.status_code != 200:
-            raise RuntimeError(f"Acquire access_token failed: HTTP {resp.status_code} {resp.text[:200]}")
+            raise RuntimeError(
+                f"Acquire access_token failed: HTTP {resp.status_code} {resp.text[:200]}"
+            )
         token_data = resp.json()
         access = token_data.get("access_token")
         if not access:
@@ -335,7 +372,7 @@ def print_token_info():
         logger.info("Cannot decode JWT token")
         return
     logger.info("=== JWT Token Information ===")
-    if 'email' in payload:
+    if "email" in payload:
         logger.info(f"Email: {payload['email']}")
-    if 'user_id' in payload:
-        logger.info(f"User ID: {payload['user_id']}") 
+    if "user_id" in payload:
+        logger.info(f"User ID: {payload['user_id']}")