cookie_extractor.py

#!/usr/bin/env python3
"""
Cookie Extractor for GrokClient

This script extracts cookies from browsers for the grok.com domain and formats them
for use with the Swift GrokClient. It specifically focuses on required cookies like
x-anonuserid, x-challenge, x-signature, sso, and sso-rw.

Usage:
    python cookie_extractor.py [--domain DOMAIN] [--format FORMAT] [--output OUTPUT]
"""

import argparse
import json
import os
import sys
try:
    import browsercookie
except ImportError:
    print("Error: Required package 'browsercookie' not found.")
    print("Install it using: pip install browsercookie")
    sys.exit(1)

def find_project_root():
    """Find the project root directory from the script location."""
    # Start with the script's directory
    script_dir = os.path.dirname(os.path.abspath(__file__))
    
    # If the script is in a Scripts folder, go up one level
    if os.path.basename(script_dir) == "Scripts":
        return os.path.dirname(script_dir)
    
    # Otherwise, assume current directory
    return os.getcwd()

def get_default_output_path():
    """Get the default path for the GrokCookies.swift file."""
    project_root = find_project_root()
    default_path = os.path.join(project_root, "Sources", "GrokClient", "GrokCookies.swift")
    
    # Make sure the directory exists
    os.makedirs(os.path.dirname(default_path), exist_ok=True)
    
    return default_path

def extract_cookies(domain=".grok.com", required_cookies=None):
    """
    Extracts cookies for the specified domain from available browsers.
    
    Args:
        domain: The domain to filter cookies for (default: .grok.com)
        required_cookies: Optional list of cookie names to filter for
        
    Returns:
        A dictionary of cookie name-value pairs
    """
    cookies = {}
    browsers_checked = []
    
    # Try Chrome first
    try:
        print("Checking Chrome for cookies...")
        chrome_cookies = browsercookie.chrome()
        chrome_dict = {cookie.name: cookie.value for cookie in chrome_cookies if domain in cookie.domain}
        cookies.update(chrome_dict)
        browsers_checked.append(f"Chrome ({len(chrome_dict)} cookies)")
    except Exception as e:
        print(f"Could not extract cookies from Chrome: {e}")
    
    # Try Firefox
    if not cookies or (required_cookies and not all(c in cookies for c in required_cookies)):
        try:
            print("Checking Firefox for cookies...")
            firefox_cookies = browsercookie.firefox()
            firefox_dict = {cookie.name: cookie.value for cookie in firefox_cookies if domain in cookie.domain}
            cookies.update(firefox_dict)
            browsers_checked.append(f"Firefox ({len(firefox_dict)} cookies)")
        except Exception as e:
            print(f"Could not extract cookies from Firefox: {e}")
    
    # Filter for required cookies if specified
    if required_cookies:
        missing = [c for c in required_cookies if c not in cookies]
        if missing:
            print(f"Warning: Some required cookies were not found: {', '.join(missing)}")
    
    print(f"Checked browsers: {', '.join(browsers_checked)}")
    return cookies

def print_swift_dict(cookies):
    """
    Prints cookies in a Swift dictionary literal format.
    
    Args:
        cookies: Dictionary of cookie name-value pairs
    """
    if not cookies:
        print("No cookies found for the specified domain.")
        return None
    
    swift_dict = "[\n"
    for name, value in cookies.items():
        swift_dict += f'    "{name}": "{value}",\n'
    swift_dict = swift_dict.rstrip(",\n") + "\n]"
    print("\nSwift Dictionary Literal for Cookies:")
    print(swift_dict)
    
    return swift_dict

def write_swift_file(cookies, filename=None):
    """
    Writes cookies to a Swift file that can be imported into a project.
    
    Args:
        cookies: Dictionary of cookie name-value pairs
        filename: The Swift file to create (defaults to Sources/GrokClient/GrokCookies.swift)
    """
    if not cookies:
        return False
    
    # Use default path if filename not provided
    if filename is None:
        filename = get_default_output_path()
    
    # Ensure directory exists
    os.makedirs(os.path.dirname(filename), exist_ok=True)
        
    swift_code = """import Foundation

// Generated by cookie_extractor.py
// Contains cookies for GrokClient authentication

public struct GrokCookies {
    public static let cookies: [String: String] = [
"""
    
    for name, value in cookies.items():
        swift_code += f'        "{name}": "{value}",\n'
    
    swift_code = swift_code.rstrip(",\n") + "\n    ]\n}\n"
    
    with open(filename, "w") as f:
        f.write(swift_code)
    
    print(f"\nWrote Swift file: {filename}")
    return True

def main():
    parser = argparse.ArgumentParser(description="Extract cookies for GrokClient")
    parser.add_argument("--domain", default=".grok.com", help="Domain to filter cookies for")
    parser.add_argument("--format", choices=["swift", "json", "both"], default="swift", 
                        help="Output format (default: swift)")
    parser.add_argument("--output", help="Output file (for Swift or JSON)")
    parser.add_argument("--required", action="store_true", help="Check for required GrokClient cookies")
    args = parser.parse_args()
    
    # Define required cookies for GrokClient
    required_cookies = None
    if args.required:
        required_cookies = ["x-anonuserid", "x-challenge", "x-signature", "sso", "sso-rw"]
    
    print(f"Extracting cookies for domain: {args.domain}")
    cookies = extract_cookies(domain=args.domain, required_cookies=required_cookies)
    
    if not cookies:
        print("No cookies found. Make sure you've logged into Grok in your browser.")
        return 1
    
    count = len(cookies)
    print(f"\nFound {count} cookie{'s' if count != 1 else ''}.")
    
    if args.format in ["swift", "both"]:
        swift_str = print_swift_dict(cookies)
        if args.output and args.format == "swift":
            output_path = args.output
            with open(output_path, "w") as f:
                f.write(swift_str)
            print(f"Wrote Swift dictionary to: {output_path}")
    
    if args.format in ["json", "both"]:
        json_str = json.dumps(cookies, indent=2)
        print("\nJSON format:")
        print(json_str)
        if args.output and args.format == "json":
            output_path = args.output
            with open(output_path, "w") as f:
                f.write(json_str)
            print(f"Wrote JSON to: {output_path}")
    
    # Always create a Swift file with the cookies, use default path if not specified
    if not args.output or args.format != "swift":
        default_path = get_default_output_path()
        write_swift_file(cookies, filename=default_path)
        print(f"Note: GrokCookies.swift will be used automatically by GrokClient when included in your project.")
    
    return 0

if __name__ == "__main__":
    sys.exit(main())