1#!/usr/bin/env python3
2"""
3Authentication Manager for NotebookLM
4Handles Google login and browser state persistence
5Based on the MCP server implementation
6 
7Implements hybrid auth approach:
8- Persistent browser profile (user_data_dir) for fingerprint consistency
9- Manual cookie injection from state.json for session cookies (Playwright bug workaround)
10See: https://github.com/microsoft/playwright/issues/36139
11"""
12 
13import json
14import time
15import argparse
16import shutil
17import re
18import sys
19from pathlib import Path
20from typing import Optional, Dict, Any
21 
22from patchright.sync_api import sync_playwright, BrowserContext
23 
24# Add parent directory to path
25sys.path.insert(0, str(Path(__file__).parent))
26 
27from config import BROWSER_STATE_DIR, STATE_FILE, AUTH_INFO_FILE, DATA_DIR
28from browser_utils import BrowserFactory
29 
30 
31class AuthManager:
32    """
33    Manages authentication and browser state for NotebookLM
34 
35    Features:
36    - Interactive Google login
37    - Browser state persistence
38    - Session restoration
39    - Account switching
40    """
41 
42    def __init__(self):
43        """Initialize the authentication manager"""
44        # Ensure directories exist
45        DATA_DIR.mkdir(parents=True, exist_ok=True)
46        BROWSER_STATE_DIR.mkdir(parents=True, exist_ok=True)
47 
48        self.state_file = STATE_FILE
49        self.auth_info_file = AUTH_INFO_FILE
50        self.browser_state_dir = BROWSER_STATE_DIR
51 
52    def is_authenticated(self) -> bool:
53        """Check if valid authentication exists"""
54        if not self.state_file.exists():
55            return False
56 
57        # Check if state file is not too old (7 days)
58        age_days = (time.time() - self.state_file.stat().st_mtime) / 86400
59        if age_days > 7:
60            print(f"⚠️ Browser state is {age_days:.1f} days old, may need re-authentication")
61 
62        return True
63 
64    def get_auth_info(self) -> Dict[str, Any]:
65        """Get authentication information"""
66        info = {
67            'authenticated': self.is_authenticated(),
68            'state_file': str(self.state_file),
69            'state_exists': self.state_file.exists()
70        }
71 
72        if self.auth_info_file.exists():
73            try:
74                with open(self.auth_info_file, 'r') as f:
75                    saved_info = json.load(f)
76                    info.update(saved_info)
77            except Exception:
78                pass
79 
80        if info['state_exists']:
81            age_hours = (time.time() - self.state_file.stat().st_mtime) / 3600
82            info['state_age_hours'] = age_hours
83 
84        return info
85 
86    def setup_auth(self, headless: bool = False, timeout_minutes: int = 10) -> bool:
87        """
88        Perform interactive authentication setup
89 
90        Args:
91            headless: Run browser in headless mode (False for login)
92            timeout_minutes: Maximum time to wait for login
93 
94        Returns:
95            True if authentication successful
96        """
97        print("🔐 Starting authentication setup...")
98        print(f"  Timeout: {timeout_minutes} minutes")
99 
100        playwright = None
101        context = None
102 
103        try:
104            playwright = sync_playwright().start()
105 
106            # Launch using factory
107            context = BrowserFactory.launch_persistent_context(
108                playwright,
109                headless=headless
110            )
111 
112            # Navigate to NotebookLM
113            page = context.new_page()
114            page.goto("https://notebooklm.google.com", wait_until="domcontentloaded")
115 
116            # Check if already authenticated
117            if "notebooklm.google.com" in page.url and "accounts.google.com" not in page.url:
118                print("  ✅ Already authenticated!")
119                self._save_browser_state(context)
120                return True
121 
122            # Wait for manual login
123            print("\n  ⏳ Please log in to your Google account...")
124            print(f"  ⏱️  Waiting up to {timeout_minutes} minutes for login...")
125 
126            try:
127                # Wait for URL to change to NotebookLM (regex ensures it's the actual domain, not a parameter)
128                timeout_ms = int(timeout_minutes * 60 * 1000)
129                page.wait_for_url(re.compile(r"^https://notebooklm\.google\.com/"), timeout=timeout_ms)
130 
131                print(f"  ✅ Login successful!")
132 
133                # Save authentication state
134                self._save_browser_state(context)
135                self._save_auth_info()
136                return True
137 
138            except Exception as e:
139                print(f"  ❌ Authentication timeout: {e}")
140                return False
141 
142        except Exception as e:
143            print(f"  ❌ Error: {e}")
144            return False
145 
146        finally:
147            # Clean up browser resources
148            if context:
149                try:
150                    context.close()
151                except Exception:
152                    pass
153 
154            if playwright:
155                try:
156                    playwright.stop()
157                except Exception:
158                    pass
159 
160    def _save_browser_state(self, context: BrowserContext):
161        """Save browser state to disk"""
162        try:
163            # Save storage state (cookies, localStorage)
164            context.storage_state(path=str(self.state_file))
165            print(f"  💾 Saved browser state to: {self.state_file}")
166        except Exception as e:
167            print(f"  ❌ Failed to save browser state: {e}")
168            raise
169 
170    def _save_auth_info(self):
171        """Save authentication metadata"""
172        try:
173            info = {
174                'authenticated_at': time.time(),
175                'authenticated_at_iso': time.strftime('%Y-%m-%d %H:%M:%S')
176            }
177            with open(self.auth_info_file, 'w') as f:
178                json.dump(info, f, indent=2)
179        except Exception:
180            pass  # Non-critical
181 
182    def clear_auth(self) -> bool:
183        """
184        Clear all authentication data
185 
186        Returns:
187            True if cleared successfully
188        """
189        print("🗑️ Clearing authentication data...")
190 
191        try:
192            # Remove browser state
193            if self.state_file.exists():
194                self.state_file.unlink()
195                print("  ✅ Removed browser state")
196 
197            # Remove auth info
198            if self.auth_info_file.exists():
199                self.auth_info_file.unlink()
200                print("  ✅ Removed auth info")
201 
202            # Clear entire browser state directory
203            if self.browser_state_dir.exists():
204                shutil.rmtree(self.browser_state_dir)
205                self.browser_state_dir.mkdir(parents=True, exist_ok=True)
206                print("  ✅ Cleared browser data")
207 
208            return True
209 
210        except Exception as e:
211            print(f"  ❌ Error clearing auth: {e}")
212            return False
213 
214    def re_auth(self, headless: bool = False, timeout_minutes: int = 10) -> bool:
215        """
216        Perform re-authentication (clear and setup)
217 
218        Args:
219            headless: Run browser in headless mode
220            timeout_minutes: Login timeout in minutes
221 
222        Returns:
223            True if successful
224        """
225        print("🔄 Starting re-authentication...")
226 
227        # Clear existing auth
228        self.clear_auth()
229 
230        # Setup new auth
231        return self.setup_auth(headless, timeout_minutes)
232 
233    def validate_auth(self) -> bool:
234        """
235        Validate that stored authentication works
236        Uses persistent context to match actual usage pattern
237 
238        Returns:
239            True if authentication is valid
240        """
241        if not self.is_authenticated():
242            return False
243 
244        print("🔍 Validating authentication...")
245 
246        playwright = None
247        context = None
248 
249        try:
250            playwright = sync_playwright().start()
251 
252            # Launch using factory
253            context = BrowserFactory.launch_persistent_context(
254                playwright,
255                headless=True
256            )
257 
258            # Try to access NotebookLM
259            page = context.new_page()
260            page.goto("https://notebooklm.google.com", wait_until="domcontentloaded", timeout=30000)
261 
262            # Check if we can access NotebookLM
263            if "notebooklm.google.com" in page.url and "accounts.google.com" not in page.url:
264                print("  ✅ Authentication is valid")
265                return True
266            else:
267                print("  ❌ Authentication is invalid (redirected to login)")
268                return False
269 
270        except Exception as e:
271            print(f"  ❌ Validation failed: {e}")
272            return False
273 
274        finally:
275            if context:
276                try:
277                    context.close()
278                except Exception:
279                    pass
280            if playwright:
281                try:
282                    playwright.stop()
283                except Exception:
284                    pass
285 
286 
287def main():
288    """Command-line interface for authentication management"""
289    parser = argparse.ArgumentParser(description='Manage NotebookLM authentication')
290 
291    subparsers = parser.add_subparsers(dest='command', help='Commands')
292 
293    # Setup command
294    setup_parser = subparsers.add_parser('setup', help='Setup authentication')
295    setup_parser.add_argument('--headless', action='store_true', help='Run in headless mode')
296    setup_parser.add_argument('--timeout', type=float, default=10, help='Login timeout in minutes (default: 10)')
297 
298    # Status command
299    subparsers.add_parser('status', help='Check authentication status')
300 
301    # Validate command
302    subparsers.add_parser('validate', help='Validate authentication')
303 
304    # Clear command
305    subparsers.add_parser('clear', help='Clear authentication')
306 
307    # Re-auth command
308    reauth_parser = subparsers.add_parser('reauth', help='Re-authenticate (clear + setup)')
309    reauth_parser.add_argument('--timeout', type=float, default=10, help='Login timeout in minutes (default: 10)')
310 
311    args = parser.parse_args()
312 
313    # Initialize manager
314    auth = AuthManager()
315 
316    # Execute command
317    if args.command == 'setup':
318        if auth.setup_auth(headless=args.headless, timeout_minutes=args.timeout):
319            print("\n✅ Authentication setup complete!")
320            print("You can now use ask_question.py to query NotebookLM")
321        else:
322            print("\n❌ Authentication setup failed")
323            exit(1)
324 
325    elif args.command == 'status':
326        info = auth.get_auth_info()
327        print("\n🔐 Authentication Status:")
328        print(f"  Authenticated: {'Yes' if info['authenticated'] else 'No'}")
329        if info.get('state_age_hours'):
330            print(f"  State age: {info['state_age_hours']:.1f} hours")
331        if info.get('authenticated_at_iso'):
332            print(f"  Last auth: {info['authenticated_at_iso']}")
333        print(f"  State file: {info['state_file']}")
334 
335    elif args.command == 'validate':
336        if auth.validate_auth():
337            print("Authentication is valid and working")
338        else:
339            print("Authentication is invalid or expired")
340            print("Run: auth_manager.py setup")
341 
342    elif args.command == 'clear':
343        if auth.clear_auth():
344            print("Authentication cleared")
345 
346    elif args.command == 'reauth':
347        if auth.re_auth(timeout_minutes=args.timeout):
348            print("\n✅ Re-authentication complete!")
349        else:
350            print("\n❌ Re-authentication failed")
351            exit(1)
352 
353    else:
354        parser.print_help()
355 
356 
357if __name__ == "__main__":
358    main()