You can close this window and return to F.R.I.D.A.Y.
" else: html = f"Error: {self.server.auth_error}
" self.wfile.write(html.encode('utf-8')) def log_message(self, format, *args): pass # Suppress HTTP logs class OAuthManager: def __init__(self, vault: CredentialVault, master_password: str, local_port: int = 8080): self.vault = vault self.master_password = master_password self.local_port = local_port self.redirect_uri = f"http://localhost:{self.local_port}/callback" self._refresh_task = None self._running = False def _generate_pkce(self): code_verifier = base64.urlsafe_b64encode(os.urandom(32)).decode('utf-8').rstrip('=') digest = hashlib.sha256(code_verifier.encode('utf-8')).digest() code_challenge = base64.urlsafe_b64encode(digest).decode('utf-8').rstrip('=') return code_verifier, code_challenge def authenticate(self, service: str, client_id: str, client_secret: str, account_identifier: str, scopes: list = None): """Perform full PKCE OAuth 2.0 flow via local browser.""" if service not in OAUTH_CONFIGS: raise ValueError(f"Unsupported service: {service}") config = OAUTH_CONFIGS[service] scopes = scopes or config["default_scopes"] code_verifier, code_challenge = self._generate_pkce() state = base64.urlsafe_b64encode(os.urandom(16)).decode('utf-8').rstrip('=') auth_params = { "client_id": client_id, "redirect_uri": self.redirect_uri, "response_type": "code", "scope": " ".join(scopes), "state": state, "code_challenge": code_challenge, "code_challenge_method": "S256" } # Some services (like Spotify) need specific prompt parameters if service == "google": auth_params["access_type"] = "offline" auth_params["prompt"] = "consent" auth_url = f"{config['auth_url']}?{urllib.parse.urlencode(auth_params)}" # Start local server to catch the callback server = HTTPServer(('localhost', self.local_port), OAuthCallbackHandler) server.auth_code = None server.auth_state = None server.auth_error = None logging.info(f"Opening browser for {service} authentication...") webbrowser.open(auth_url) # Wait for callback (blocks) server.handle_request() server.server_close() if server.auth_error: raise Exception(f"OAuth Error: {server.auth_error}") if server.auth_state != state: raise Exception("OAuth Error: State mismatch (possible CSRF)") if not server.auth_code: raise Exception("OAuth Error: No code returned") # Exchange code for token logging.info("Exchanging code for tokens...") token_data = self._exchange_token(service, client_id, client_secret, server.auth_code, code_verifier) # Calculate absolute expiry time if "expires_in" in token_data: token_data["expires_at"] = time.time() + int(token_data["expires_in"]) # Store client info alongside token for background refresh token_data["_client_id"] = client_id token_data["_client_secret"] = client_secret token_data["_service"] = service self.vault.store(service, account_identifier, token_data, self.master_password) logging.info(f"Successfully authenticated and stored tokens for {service}/{account_identifier}") return token_data def _exchange_token(self, service: str, client_id: str, client_secret: str, code: str, code_verifier: str = None) -> dict: config = OAUTH_CONFIGS[service] data = { "client_id": client_id, "client_secret": client_secret, "grant_type": "authorization_code", "code": code, "redirect_uri": self.redirect_uri } if code_verifier: data["code_verifier"] = code_verifier data = urllib.parse.urlencode(data).encode('utf-8') req = urllib.request.Request(config["token_url"], data=data, method="POST") req.add_header('Content-Type', 'application/x-www-form-urlencoded') req.add_header('Accept', 'application/json') try: with urllib.request.urlopen(req, timeout=10) as response: return json.loads(response.read().decode('utf-8')) except HTTPError as e: err_body = e.read().decode('utf-8') raise Exception(f"Token exchange failed: {e.code} - {err_body}") def refresh_token(self, service: str, account_identifier: str) -> dict: """Force refresh a token.""" token_data = self.vault.retrieve(service, account_identifier, self.master_password) refresh_token = token_data.get("refresh_token") if not refresh_token: logging.warning(f"No refresh token available for {service}/{account_identifier}") return token_data client_id = token_data.get("_client_id") client_secret = token_data.get("_client_secret") config = OAUTH_CONFIGS[service] data = { "client_id": client_id, "client_secret": client_secret, "grant_type": "refresh_token", "refresh_token": refresh_token } data_encoded = urllib.parse.urlencode(data).encode('utf-8') req = urllib.request.Request(config["token_url"], data=data_encoded, method="POST") req.add_header('Content-Type', 'application/x-www-form-urlencoded') req.add_header('Accept', 'application/json') try: with urllib.request.urlopen(req, timeout=10) as response: new_data = json.loads(response.read().decode('utf-8')) # Update token data (retain original refresh token if a new one isn't provided) if "refresh_token" not in new_data: new_data["refresh_token"] = refresh_token if "expires_in" in new_data: new_data["expires_at"] = time.time() + int(new_data["expires_in"]) new_data["_client_id"] = client_id new_data["_client_secret"] = client_secret new_data["_service"] = service self.vault.store(service, account_identifier, new_data, self.master_password) logging.info(f"Successfully refreshed token for {service}/{account_identifier}") return new_data except Exception as e: logging.error(f"Failed to refresh token for {service}/{account_identifier}: {e}") raise async def _background_refresh_loop(self): """Background loop to check token expiry every 5 minutes.""" self._running = True logging.info("Starting OAuth token refresh background loop...") while self._running: try: # We need a way to enumerate credentials to check expiry # This requires exposing a list method on the vault, or tracking them in memory # For this implementation, we assume the agent will call refresh_token on demand # or we implement an enumerate method on CredentialVault # Fetch all service/accounts from the vault rows = self.vault.list_accounts() now = time.time() for service, account in rows: try: token_data = self.vault.retrieve(service, account, self.master_password) expires_at = token_data.get("expires_at", 0) # Refresh if expiring within the next 10 minutes if expires_at and (expires_at - now) < 600: logging.info(f"Token for {service}/{account} expiring soon. Refreshing...") self.refresh_token(service, account) except Exception as e: logging.error(f"Background refresh error for {service}/{account}: {e}") except Exception as e: logging.error(f"OAuth refresh loop error: {e}") # Sleep 5 minutes await asyncio.sleep(300) def start_background_refresh(self): if self._refresh_task is None: self._refresh_task = asyncio.create_task(self._background_refresh_loop()) def stop_background_refresh(self): self._running = False if self._refresh_task: self._refresh_task.cancel() self._refresh_task = None