init: whatever

src/tokens.py

@@ -0,0 +1,96 @@
+import json
+import time
+import os
+import base64
+from pathlib import Path
+from dataclasses import dataclass
+import aiohttp
+
+DATA_DIR = Path(os.environ.get("DATA_DIR", "./data"))
+TOKENS_FILE = DATA_DIR / "tokens.json"
+
+CLIENT_ID = "app_EMoamEEZ73f0CkXaXp7hrann"
+TOKEN_URL = "https://auth.openai.com/oauth/token"
+
+
+@dataclass
+class Tokens:
+    access_token: str
+    refresh_token: str
+    expires_at: float  # unix timestamp
+
+    @property
+    def is_expired(self) -> bool:
+        return time.time() >= self.expires_at - 10
+
+
+def load_tokens() -> Tokens | None:
+    if not TOKENS_FILE.exists():
+        return None
+    try:
+        with open(TOKENS_FILE) as f:
+            data = json.load(f)
+        access_token = data["access_token"]
+        return Tokens(
+            access_token=access_token,
+            refresh_token=data["refresh_token"],
+            expires_at=data["expires_at"],
+        )
+    except (json.JSONDecodeError, KeyError):
+        return None
+
+
+def save_tokens(tokens: Tokens):
+    TOKENS_FILE.parent.mkdir(parents=True, exist_ok=True)
+    with open(TOKENS_FILE, "w") as f:
+        json.dump(
+            {
+                "access_token": tokens.access_token,
+                "refresh_token": tokens.refresh_token,
+                "expires_at": tokens.expires_at,
+            },
+            f,
+            indent=2,
+        )
+
+
+async def refresh_tokens(refresh_token: str) -> Tokens | None:
+    async with aiohttp.ClientSession() as session:
+        data = {
+            "grant_type": "refresh_token",
+            "refresh_token": refresh_token,
+            "client_id": CLIENT_ID,
+        }
+        async with session.post(TOKEN_URL, data=data) as resp:
+            if not resp.ok:
+                text = await resp.text()
+                print(f"Token refresh failed: {resp.status} {text}")
+                return None
+            json_resp = await resp.json()
+            expires_in = json_resp["expires_in"]
+            return Tokens(
+                access_token=json_resp["access_token"],
+                refresh_token=json_resp["refresh_token"],
+                expires_at=time.time() + expires_in,
+            )
+
+
+async def get_valid_tokens() -> Tokens | None:
+    tokens = load_tokens()
+    if not tokens:
+        print("No tokens found")
+        return None
+
+    if tokens.is_expired:
+        print("Token expired, refreshing...")
+        if not tokens.refresh_token:
+            print("No refresh token available")
+            return None
+        new_tokens = await refresh_tokens(tokens.refresh_token)
+        if not new_tokens:
+            print("Failed to refresh token")
+            return None
+        save_tokens(new_tokens)
+        return new_tokens
+
+    return tokens