Add automatic login functionality with optional credentials and selectors

scrapers/capture_with_cdp.py

@@ -33,6 +33,11 @@ except Exception:
     logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
     logger = logging.getLogger("capture_with_cdp")
 from playwright.sync_api import sync_playwright
+from dotenv import load_dotenv
+import os
+
+# load .env
+load_dotenv()
 
 
 def sanitize_fn(s: str) -> str:
@@ -83,6 +88,9 @@ def main():
     parser.add_argument("--headless", action="store_true", help="Run headless")
     parser.add_argument("--capture-wait", type=int, default=8, help="Seconds to wait after reload to capture XHRs")
     parser.add_argument("--keep-open", action="store_true", help="Keep browser open after capture")
+    parser.add_argument("--username", help="STATdx username (or use STATDX_USERNAME env var)")
+    parser.add_argument("--password", help="STATdx password (or use STATDX_PASSWORD env var)")
+    parser.add_argument("--post-login-selector", default="#ds-app", help="Selector to wait for after login")
     args = parser.parse_args()
 
     out_dir = Path(args.output_dir)
@@ -191,6 +199,63 @@ def main():
         logger.info("Opening page: {}. Please log in if needed, then press ENTER to start capture.", args.url)
         page.goto(args.url)
 
+        # Automatic login if credentials provided
+        username = args.username or os.getenv('STATDX_USERNAME')
+        password = args.password or os.getenv('STATDX_PASSWORD')
+        if username and password:
+            logger.info('Attempting automatic login using provided credentials')
+            user_selectors = ['input[name="username"]', 'input.usernameSelector', 'input[type="email"]', 'input[name*="email" i]', 'input[name*="user" i]', 'input[type="text"]']
+            pass_selectors = ['input[type="password"]', 'input[name="password"]', 'input.passwordSelector']
+
+            user_sel = None
+            for sel in user_selectors:
+                try:
+                    if page.query_selector(sel):
+                        user_sel = sel
+                        break
+                except Exception:
+                    continue
+
+            pass_sel = None
+            for sel in pass_selectors:
+                try:
+                    if page.query_selector(sel):
+                        pass_sel = sel
+                        break
+                except Exception:
+                    continue
+
+            if user_sel and pass_sel:
+                try:
+                    page.fill(user_sel, username)
+                    time.sleep(0.1)
+                    page.fill(pass_sel, password)
+                    time.sleep(0.1)
+                    # try submit methods
+                    submit_selectors = ['button[type=submit]', 'input[type=submit]', 'button:has-text("Sign in")', 'button:has-text("Sign In")', 'button:has-text("Log in")']
+                    clicked = False
+                    for s in submit_selectors:
+                        try:
+                            btn = page.query_selector(s)
+                            if btn:
+                                btn.click()
+                                clicked = True
+                                logger.info('Clicked submit button using selector: %s', s)
+                                break
+                        except Exception:
+                            continue
+                    if not clicked:
+                        page.press(pass_sel, 'Enter')
+
+                    # wait for post-login selector
+                    try:
+                        if args.post_login_selector:
+                            page.wait_for_selector(args.post_login_selector, timeout=10000)
+                    except Exception:
+                        logger.info('Post-login selector not found within timeout')
+                except Exception:
+                    logger.exception('Automatic login attempt failed')
+
         input("After logging in and navigating to the page you want to capture, press Enter to begin capture...\n")
 
         logger.info("Reloading page to trigger XHRs and network activity")