google-cookie-retrieval/selenium_cookie_extractor_json.py

#!/usr/bin/env python3
import argparse
import json
import subprocess
import sys
import time

from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.common.by import By
from selenium.webdriver.common.keys import Keys
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.support.ui import WebDriverWait


def check_rbw_unlocked():
    """Check if rbw vault is unlocked."""
    result = subprocess.run(["rbw", "unlocked"], capture_output=True)
    return result.returncode == 0


def get_rbw_credentials(entry: str):
    """Get username and password from rbw."""
    # Get password
    result = subprocess.run(
        ["rbw", "get", entry],
        capture_output=True,
        text=True,
    )
    if result.returncode != 0:
        return None, None

    password = result.stdout.strip()

    # Get username
    result = subprocess.run(
        ["rbw", "get", "--field", "username", entry],
        capture_output=True,
        text=True,
    )
    username = result.stdout.strip() if result.returncode == 0 else None

    return username, password


def copy_to_clipboard(text: str):
    """Copy text to clipboard using wl-copy or xclip."""
    # Try wl-copy first (Wayland)
    try:
        subprocess.run(
            ["wl-copy"],
            input=text.encode(),
            check=True,
            capture_output=True,
        )
        return True
    except (subprocess.CalledProcessError, FileNotFoundError):
        pass

    # Fall back to xclip (X11)
    try:
        subprocess.run(
            ["xclip", "-selection", "clipboard"],
            input=text.encode(),
            check=True,
            capture_output=True,
        )
        return True
    except (subprocess.CalledProcessError, FileNotFoundError):
        pass

    return False


def automate_google_login(driver, username: str, password: str):
    """Attempt to automate Google login. Returns True if successful."""
    wait = WebDriverWait(driver, 10)

    try:
        # Wait for and fill email field
        email_field = wait.until(
            EC.presence_of_element_located((By.CSS_SELECTOR, 'input[type="email"]'))
        )
        email_field.send_keys(username)
        email_field.send_keys(Keys.RETURN)

        # Wait for password field
        time.sleep(2)  # Google has animation between screens
        password_field = wait.until(
            EC.presence_of_element_located((By.CSS_SELECTOR, 'input[type="password"]'))
        )
        password_field.send_keys(password)
        password_field.send_keys(Keys.RETURN)

        return True
    except Exception as e:
        print(f"Auto-login failed: {e}")
        return False


def extract_cookies(driver):
    """Extract the required cookies from the browser."""
    # Navigate to a chat.google.com path to ensure we get the right cookies
    driver.get("https://chat.google.com/u/0/mole/world")

    all_cookies = driver.get_cookies()
    target_names = {"COMPASS", "SSID", "SID", "OSID", "HSID"}
    extracted = {}
    compass_cookie = None

    for cookie in all_cookies:
        name_upper = cookie["name"].upper()
        if name_upper not in target_names:
            continue
        if name_upper == "COMPASS":
            if cookie.get("path", "") == "/":
                compass_cookie = cookie
            elif compass_cookie is None:
                compass_cookie = cookie
        else:
            extracted[name_upper] = cookie["value"]

    if compass_cookie:
        extracted["COMPASS"] = compass_cookie["value"]

    return extracted


def main():
    parser = argparse.ArgumentParser(
        description="Extract Google Chat cookies for Mautrix bridge authentication"
    )
    parser.add_argument(
        "--entry",
        default="google.com",
        help="rbw entry name for Google credentials (default: google.com)",
    )
    parser.add_argument(
        "--copy",
        action="store_true",
        help="Copy the cookie JSON to clipboard",
    )
    parser.add_argument(
        "--manual",
        action="store_true",
        help="Skip auto-login, do manual login only",
    )
    args = parser.parse_args()

    username = None
    password = None

    # Get credentials from rbw if not manual mode
    if not args.manual:
        if not check_rbw_unlocked():
            print("rbw vault is locked. Please run: rbw unlock")
            sys.exit(1)

        username, password = get_rbw_credentials(args.entry)
        if not username or not password:
            print(f"Could not get credentials from rbw entry '{args.entry}'")
            print("Falling back to manual login...")

    # Configure Chrome options
    options = webdriver.ChromeOptions()
    options.add_argument("--incognito")

    # Initialize the Chrome driver
    service = Service()
    driver = webdriver.Chrome(service=service, options=options)

    try:
        print("Opening https://chat.google.com ...")
        driver.get("https://chat.google.com")

        if username and password and not args.manual:
            print("Attempting automatic login...")
            if automate_google_login(driver, username, password):
                print("Credentials entered. Complete any 2FA if prompted.")
            else:
                print("Auto-login failed, please log in manually.")

        print("\nOnce logged in, press Enter to extract cookies...")
        input()

        cookies = extract_cookies(driver)
        json_data = json.dumps(cookies, indent=2)

        print("\nExtracted Cookie JSON:")
        print(json_data)

        if args.copy:
            if copy_to_clipboard(json_data):
                print("\nCopied to clipboard!")
            else:
                print("\nFailed to copy to clipboard (wl-copy/xclip not found)")

    finally:
        print("\nClosing browser...")
        driver.quit()


if __name__ == "__main__":
    main()