Files

91 lines
3.5 KiB
Python

"""Fixed-provider OAuth refresh. No secret payloads in durable operation history."""
import json
import threading
import time
from urllib.parse import urlencode
from pigway_frame_flow.errors import ErrorCode as E
from pigway_frame_flow.errors import FlowError
from pigway_frame_flow.extensions.secrets import save_secret
class OAuthTokens:
ENDPOINTS = {
"google-drive": "https://oauth2.googleapis.com/token",
"alipan": "https://openapi.aliyundrive.com/oauth/access_token",
"baidu": "https://openapi.baidu.com/oauth/2.0/token",
}
def __init__(self, settings, assets):
self.settings, self.assets = settings, assets
self.lock = threading.RLock()
def refresh(self, kind, credentials, path=None):
with self.lock:
if not all(
credentials.get(key)
for key in ("refresh_token", "client_id", "client_secret")
):
raise FlowError(
E.CONFIG_REQUIRED,
"OAuth refresh requires client ID, client secret and refresh token",
409,
)
values = {
key: credentials[key]
for key in ("refresh_token", "client_id", "client_secret")
}
values["grant_type"] = "refresh_token"
for key in ("access_token", "refresh_token", "client_secret"):
if credentials.get(key):
self.settings.additional_secrets.append(credentials[key])
endpoint = self.ENDPOINTS[kind]
host = endpoint.split("/")[2]
if host not in self.settings.image_url_hosts:
self.settings.image_url_hosts.append(host)
as_json = kind == "alipan"
body = (
json.dumps(values).encode() if as_json else urlencode(values).encode()
)
try:
raw = self.assets.fetch_url(
endpoint,
method="POST",
headers={
"Content-Type": "application/json"
if as_json
else "application/x-www-form-urlencoded"
},
body=body,
limit=65536,
)
response = json.loads(raw)
if (
not isinstance(response.get("access_token"), str)
or not response["access_token"]
):
raise ValueError()
expiry = int(response["expires_in"])
if not 60 <= expiry <= 366 * 86400:
raise ValueError()
updated = {
**credentials,
"access_token": response["access_token"],
"expires_at": time.time() + expiry,
}
if response.get("refresh_token"):
updated["refresh_token"] = response["refresh_token"]
for key in ("access_token", "refresh_token"):
if updated.get(key):
self.settings.additional_secrets.append(updated[key])
if path:
save_secret(path, json.dumps(updated))
return updated
except Exception:
raise FlowError(
E.CLOUD_UNAVAILABLE,
"OAuth refresh failed; reconnect or refresh credentials manually",
502,
) from None