| """ |
| Notion Login Helper for MCPMark |
| ================================= |
| |
| This module provides a utility class and CLI script for logging into Notion |
| using Playwright. It saves the authenticated session state to a file, |
| which can be used for subsequent automated tasks. |
| """ |
|
|
| import argparse |
| from pathlib import Path |
| from typing import Optional |
|
|
| from playwright.sync_api import ( |
| BrowserContext, |
| Page, |
| TimeoutError as PlaywrightTimeoutError, |
| sync_playwright, |
| ) |
|
|
| from src.base.login_helper import BaseLoginHelper |
| from src.logger import get_logger |
|
|
| |
| logger = get_logger(__name__) |
|
|
|
|
| class NotionLoginHelper(BaseLoginHelper): |
| """ |
| Utility helper for logging into Notion using Playwright. |
| """ |
|
|
| SUPPORTED_BROWSERS = {"chromium", "firefox"} |
|
|
| def __init__( |
| self, |
| *, |
| url: Optional[str] = None, |
| headless: bool = True, |
| state_path: Optional[str | Path] = None, |
| browser: str = "firefox", |
| ) -> None: |
| """ |
| Initializes the Notion login helper. |
| |
| Args: |
| url: The Notion URL to open after launching the browser. |
| headless: Whether to run Playwright in headless mode. |
| state_path: The path to save the authenticated session state. |
| browser: The browser engine to use ('chromium' or 'firefox'). |
| """ |
| super().__init__() |
| if browser not in self.SUPPORTED_BROWSERS: |
| raise ValueError( |
| f"Unsupported browser '{browser}'. Supported browsers are: {', '.join(self.SUPPORTED_BROWSERS)}" |
| ) |
|
|
| self.url = url or "https://www.notion.so/login" |
| self.headless = headless |
| self.browser_name = browser |
| self.state_path = ( |
| Path(state_path or Path.cwd() / "notion_state.json").expanduser().resolve() |
| ) |
| self._browser_context: Optional[BrowserContext] = None |
| self._playwright = None |
| self._browser = None |
|
|
| def login(self) -> BrowserContext: |
| """ |
| Launches a browser, performs login, and saves the session state. |
| """ |
| if self.state_path.exists(): |
| try: |
| self.state_path.unlink() |
| except OSError as e: |
| logger.warning("Unable to remove existing state file: %s", e) |
|
|
| if self._playwright is None: |
| self._playwright = sync_playwright().start() |
|
|
| browser_type = getattr(self._playwright, self.browser_name) |
| self._browser = browser_type.launch(headless=self.headless) |
| context = self._browser.new_context() |
| page = context.new_page() |
|
|
| logger.info("Navigating to Notion URL: %s", self.url) |
| page.goto(self.url, wait_until="load") |
|
|
| if self.headless: |
| self._handle_headless_login(context) |
| else: |
| logger.info( |
| "A browser window has been opened. Please complete the Notion login." |
| ) |
| logger.info( |
| "After you see your workspace, return to this terminal and press <ENTER>." |
| ) |
| initial_url = page.url |
| input() |
| try: |
| page.wait_for_url(lambda u: u != initial_url, timeout=10_000) |
| except PlaywrightTimeoutError: |
| pass |
|
|
| try: |
| page.wait_for_load_state("domcontentloaded", timeout=5_000) |
| except PlaywrightTimeoutError: |
| pass |
|
|
| context.storage_state(path=str(self.state_path)) |
| logger.info("✅ Login successful! Session state saved to %s", self.state_path) |
|
|
| self._browser_context = context |
| return context |
|
|
| def close(self) -> None: |
| """Closes the underlying browser and Playwright instance.""" |
| if self._browser_context: |
| try: |
| self._browser_context.close() |
| finally: |
| self._browser_context = None |
| if self._browser: |
| try: |
| self._browser.close() |
| finally: |
| self._browser = None |
| if self._playwright: |
| self._playwright.stop() |
| self._playwright = None |
|
|
| def _handle_headless_login(self, context: BrowserContext) -> None: |
| """ |
| Guides the user through the login process in headless mode. |
| """ |
| page: Page = context.pages[0] |
| login_url = "https://www.notion.so/login" |
| page.goto(login_url, wait_until="domcontentloaded") |
|
|
| email = input("Enter your Notion email address: ").strip() |
| try: |
| email_input = page.locator( |
| 'input[placeholder="Enter your email address..."]' |
| ) |
| email_input.wait_for(state="visible", timeout=120_000) |
| email_input.fill(email) |
| email_input.press("Enter") |
| except PlaywrightTimeoutError: |
| raise RuntimeError("Timed out waiting for the email input field.") |
| except Exception: |
| page.get_by_role("button", name="Continue", exact=True).click() |
|
|
| try: |
| code_input = page.locator('input[placeholder="Enter code"]') |
| code_input.wait_for(state="visible", timeout=120_000) |
| code = input("Enter the verification code from your email: ").strip() |
| code_input.fill(code) |
| code_input.press("Enter") |
| except PlaywrightTimeoutError: |
| raise RuntimeError("Timed out waiting for the verification code input.") |
| except Exception: |
| page.get_by_role("button", name="Continue", exact=True).click() |
|
|
| try: |
| page.wait_for_url(lambda url: url != login_url, timeout=180_000) |
| except PlaywrightTimeoutError: |
| logger.warning("Login redirect timed out, but proceeding to save state.") |
|
|
| if self.url and self.url != login_url: |
| page.goto(self.url, wait_until="domcontentloaded") |
|
|
| def __enter__(self) -> "NotionLoginHelper": |
| self.login() |
| return self |
|
|
| def __exit__(self, exc_type, exc_val, exc_tb): |
| self.close() |
|
|
|
|
| def main(): |
| """Main entry point for the Notion login CLI script.""" |
| parser = argparse.ArgumentParser( |
| description="Authenticate to Notion and generate a session state file.", |
| formatter_class=argparse.ArgumentDefaultsHelpFormatter, |
| ) |
| parser.add_argument( |
| "--headless", |
| action="store_true", |
| help="Run the login flow in headless mode (prompts for credentials).", |
| ) |
| parser.add_argument( |
| "--browser", |
| default="firefox", |
| choices=["chromium", "firefox"], |
| help="The browser engine to use for Playwright.", |
| ) |
| args = parser.parse_args() |
|
|
| helper = NotionLoginHelper(headless=args.headless, browser=args.browser) |
| with helper: |
| logger.info("Login process completed.") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
