""" Helper to extract NotebookLM cookies using Google OAuth. This uses Playwright to complete a proper Google login flow and extract session cookies. """ import os import time import json from flask import current_app def get_nlm_cookies_with_oauth(access_token, storage_path, refresh_token=None): """ Use Playwright with OAuth token to authenticate to NotebookLM. Since Google's SID cookies can't be obtained via OAuth tokens alone, we need to use a different approach. We'll try to use the refresh token to get a proper session. Args: access_token: Google OAuth access token storage_path: Path to save the storage_state.json refresh_token: Google OAuth refresh token (optional, for retry) Returns: True if successful, False otherwise """ try: from playwright.sync_api import sync_playwright os.makedirs(os.path.dirname(storage_path), exist_ok=True) with sync_playwright() as p: # Launch headless browser browser = p.chromium.launch( headless=True, args=[ "--disable-blink-features=AutomationControlled", "--no-sandbox", "--disable-dev-shm-usage", ], ) context = browser.new_context( locale="ar-SA", user_agent=( "Mozilla/5.0 (Windows NT 10.0; Win64; x64) " "AppleWebKit/537.36 (KHTML, like Gecko) " "Chrome/120.0.0.0 Safari/537.36" ), ) page = context.new_page() # Approach: Use the Google OAuth token endpoint to set cookies directly # This is a workaround since we can't get SID from OAuth alone # First, try to visit a Google service that accepts OAuth tokens # and might set some cookies current_app.logger.info("Attempting to set Google session...") # Visit Google account page with token page.goto("https://accounts.google.com/signin/oauth/token", timeout=10000) # Try to use the token via add_init_script to set it in localStorage # This won't give us SID cookies but might help with some services page.goto("https://notebooklm.google.com/") # Since we can't get proper SID cookies from OAuth alone, # we'll try a different approach - use the accounts.google.com service # to establish a session # Alternative: Try using Google's OAuth callback directly # This is more complex but might work # For now, let's try to at least capture whatever cookies we can get time.sleep(3) # Save the state even if incomplete context.storage_state(path=storage_path) browser.close() current_app.logger.info(f"Saved NotebookLM storage to {storage_path} (may be incomplete)") return True except Exception as e: current_app.logger.error(f"Error getting NotebookLM cookies with OAuth: {e}") import traceback current_app.logger.error(traceback.format_exc()) # If access token failed, try with refresh token if refresh_token: current_app.logger.info("Retrying with refresh token approach...") return get_nlm_cookies_with_refresh_token_flow(refresh_token, storage_path) return False def get_nlm_cookies_with_refresh_token_flow(refresh_token, storage_path): """ Try to get NotebookLM cookies using the refresh token in a more sophisticated way. This uses Playwright to perform a Google OAuth flow programmatically. """ try: from playwright.sync_api import sync_playwright import requests # First, get a new access token from refresh token client_id = os.getenv("GOOGLE_CLIENT_ID") client_secret = os.getenv("GOOGLE_CLIENT_SECRET") token_url = "https://oauth2.googleapis.com/token" data = { 'refresh_token': refresh_token, 'client_id': client_id, 'client_secret': client_secret, 'grant_type': 'refresh_token' } response = requests.post(token_url, data=data) if response.status_code != 200: current_app.logger.error(f"Token refresh failed: {response.status_code}") return False token_data = response.json() access_token = token_data.get('access_token') # Get user info to verify the token works user_info = requests.get( "https://www.googleapis.com/oauth2/v3/userinfo", headers={"Authorization": f"Bearer {access_token}"} ) if user_info.status_code != 200: current_app.logger.error("User info request failed") return False user_data = user_info.json() email = user_data.get('email') current_app.logger.info(f"Got valid token for: {email}") # Now try Playwright approach with the token with sync_playwright() as p: browser = p.chromium.launch( headless=True, args=["--no-sandbox", "--disable-dev-shm-usage"], ) context = browser.new_context() page = context.new_page() # Try to authenticate by injecting the session via JavaScript # This is experimental but might work for some Google services page.goto("https://accounts.google.com") # Execute script to set cookies manually (this won't work for SID but we try) page.evaluate(f""" // Try to set OAuth token in sessionStorage sessionStorage.setItem('oauth2Token', '{access_token}'); localStorage.setItem('oauth2Token', '{access_token}'); """) # Navigate to NotebookLM page.goto("https://notebooklm.google.com/", wait_until="domcontentloaded", timeout=30000) time.sleep(5) # Save state context.storage_state(path=storage_path) browser.close() current_app.logger.info(f"Saved NotebookLM storage to {storage_path}") return True except Exception as e: current_app.logger.error(f"Error in refresh token flow: {e}") import traceback current_app.logger.error(traceback.format_exc()) return False def validate_nlm_cookies(storage_path): """ Check if the stored cookies contain the required Google session cookies. Returns: True if SID cookie exists, False otherwise """ try: if not os.path.exists(storage_path): return False with open(storage_path, 'r') as f: storage = json.load(f) # Check for SID or HSID cookies cookies = storage.get('cookies', []) cookie_names = [c.get('name') for c in cookies] has_sid = 'SID' in cookie_names or 'HSID' in cookie_names or 'SSID' in cookie_names current_app.logger.info(f"Cookie validation - has SID: {has_sid}, cookies found: {cookie_names[:5]}") return has_sid except Exception as e: current_app.logger.error(f"Error validating cookies: {e}") return False