Upload 53 files

warp2protobuf/core/auth.py

@@ -6,18 +6,29 @@ JWT Authentication for Warp API
 Handles JWT token management, refresh, and validation.
 Integrates functionality from refresh_jwt.py.
 """
-import base64
-import json
-import os
-import time
-from pathlib import Path
-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
+import base64
+import json
+import os
+import time
+from pathlib import Path
+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 .logging import logger, log
-from .proxy_manager import AsyncProxyManager  # 新增: 导入代理管理器
+from .proxy_manager import AsyncProxyManager  # 新增: 导入代理管理器
+
+
+DEFAULT_DATA_DIR = os.getenv("DATA_DIR", "/data")
+ENV_FILE_PATH = Path(os.getenv("ENV_FILE_PATH", os.path.join(DEFAULT_DATA_DIR, ".env")))
+
+try:
+    ENV_FILE_PATH.parent.mkdir(parents=True, exist_ok=True)
+except Exception:
+    pass
+
+load_dotenv(dotenv_path=ENV_FILE_PATH if ENV_FILE_PATH.exists() else None, override=True)
 
 
 def decode_jwt_payload(token: str) -> dict:
@@ -137,26 +148,32 @@ async def refresh_jwt_token() -> dict:
     return {}
 
 
-def update_env_file(new_jwt: str) -> bool:
-    env_path = Path(".env")
-    try:
-        set_key(str(env_path), "WARP_JWT", new_jwt)
-        logger.info("Updated .env file with new JWT token")
-        return True
-    except Exception as e:
-        logger.error(f"Error updating .env file: {e}")
-        return False
-
-
-def update_env_refresh_token(refresh_token: str) -> bool:
-    env_path = Path(".env")
-    try:
-        set_key(str(env_path), "WARP_REFRESH_TOKEN", refresh_token)
-        logger.info("Updated .env with WARP_REFRESH_TOKEN")
-        return True
-    except Exception as e:
-        logger.error(f"Error updating .env WARP_REFRESH_TOKEN: {e}")
-        return False
+def update_env_file(new_jwt: str) -> bool:
+    os.environ["WARP_JWT"] = new_jwt
+    if ENV_FILE_PATH is None:
+        return True
+
+    try:
+        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:
+        logger.warning(f"Unable to persist JWT to env file {ENV_FILE_PATH}: {e}")
+        return True
+
+
+def update_env_refresh_token(refresh_token: str) -> bool:
+    if refresh_token:
+        os.environ["WARP_REFRESH_TOKEN"] = refresh_token
+    if ENV_FILE_PATH is None:
+        return True
+    try:
+        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:
+        logger.warning(f"Unable to persist refresh token to env file {ENV_FILE_PATH}: {e}")
+        return True
 
 
 async def check_and_refresh_token() -> bool:
@@ -195,20 +212,29 @@ async def check_and_refresh_token() -> bool:
 
 
 async def get_valid_jwt() -> str:
-    from dotenv import load_dotenv as _load
-    _load(override=True)
+    from dotenv import load_dotenv as _load
+    if ENV_FILE_PATH is not None:
+        _load(dotenv_path=ENV_FILE_PATH, override=True)
+    else:
+        _load(override=True)
     jwt = os.getenv("WARP_JWT")
     if not jwt:
         logger.info("No JWT token found, attempting to refresh...")
         if await check_and_refresh_token():
-            _load(override=True)
+            if ENV_FILE_PATH is not None:
+                _load(dotenv_path=ENV_FILE_PATH, override=True)
+            else:
+                _load(override=True)
             jwt = os.getenv("WARP_JWT")
         if not jwt:
             raise RuntimeError("WARP_JWT is not set and refresh failed")
     if is_token_expired(jwt, buffer_minutes=2):
         logger.info("JWT token is expired or expiring soon, attempting to refresh...")
         if await check_and_refresh_token():
-            _load(override=True)
+            if ENV_FILE_PATH is not None:
+                _load(dotenv_path=ENV_FILE_PATH, override=True)
+            else:
+                _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")