File size: 6,898 Bytes
a2ec7b6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 | """
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
# Initialize 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 # It's okay if the URL doesn't change
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()
|