oauth_session.py

import json
import pathlib
import webbrowser
import threading
import queue
import os
from http.server import HTTPServer, BaseHTTPRequestHandler
from typing import Dict, List, Optional
from urllib.parse import urlparse, urlunparse
from requests_oauthlib import OAuth2Session

# Allow OAuth2 over HTTP for localhost (development only)
# This is safe because localhost traffic never leaves your machine
os.environ['OAUTHLIB_INSECURE_TRANSPORT'] = '1'




class OAuthSession:
    """
    Encapsulated OAuth2 authentication and session management for API access.

    This class handles OAuth2 configuration, token management, and provides
    an authenticated session for making API requests.
    """
    
    def __init__(
        self,
        authorization_base_url: str,
        token_url: str,
        scope: List[str],
        config_dir: str = "./.api-blaster",
        auth_file: str = "api-blaster-auth.json",
        conf_file: str = "api-blaster-conf.json"
    ):
        """
        Initialize the OAuth2 client.
        
        Args:
            authorization_base_url: OAuth2 authorization endpoint
            token_url: OAuth2 token endpoint  
            scope: List of OAuth2 scopes to request
            config_dir: Directory to store configuration and token files
            auth_file: Filename for storing authentication tokens
            conf_file: Filename for storing client configuration
        """
        self.authorization_base_url = authorization_base_url
        self.token_url = token_url
        self.refresh_url = token_url
        self.scope = scope
        
        # Setup file paths
        self.config_dir = pathlib.Path(config_dir)
        self.config_dir.mkdir(parents=True, exist_ok=True)
        self.auth_file_path = self.config_dir / auth_file
        self.conf_file_path = self.config_dir / conf_file
        
        # Initialize session
        self.session: Optional[OAuth2Session] = None
        self.config: Optional[Dict] = None
        
    def _token_updater(self, token: Dict) -> None:
        """Update stored token when it's refreshed."""
        with open(self.auth_file_path, "w") as f:
            json.dump(token, f)
    
    def _load_config(self) -> Dict:
        """Load configuration from file or prompt user for credentials."""
        try:
            with open(self.conf_file_path, "r") as f:
                config = json.load(f)
            
            # Validate required fields
            required_fields = ["client_id", "client_secret", "redirect_uri"]
            for field in required_fields:
                if field not in config:
                    raise KeyError(f"Missing required field: {field}")
                    
            return config
            
        except (FileNotFoundError, json.JSONDecodeError, KeyError):
            return self._prompt_for_config()
    
    def _prompt_for_config(self) -> Dict:
        """Prompt user for OAuth2 configuration and save it."""
        print("No config saved to file yet, lets get set up!")
        print("\nFor the redirect URL, use http://localhost:8080/cb (recommended for OAuth)")
        config = {
            "client_id": input("Enter your Client ID: "),
            "client_secret": input("Enter your Client Secret: "),
            "redirect_uri": input("Enter your Redirect URL (default: http://localhost:8080/cb): ") or "http://localhost:8080/cb"
        }

        # Save config
        with open(self.conf_file_path, "w") as f:
            json.dump(config, f)

        return config
    
    def _load_existing_token(self) -> Optional[Dict]:
        """Load existing token from file if available and valid."""
        try:
            with open(self.auth_file_path, "r") as f:
                token = json.load(f)

            # Validate required fields
            required_fields = ["refresh_token", "expires_at"]
            for field in required_fields:
                if field not in token:
                    raise KeyError(f"Missing required field: {field}")

            return token

        except (FileNotFoundError, json.JSONDecodeError, KeyError):
            return None

    def _start_callback_server(self, host: str, port: int) -> tuple[Optional[HTTPServer], queue.Queue, threading.Event]:
        """Start a local server to capture OAuth callback.

        Returns:
            Tuple of (server, response_queue, server_ready_event)
        """
        # Create queue for this server instance (thread-safe)
        response_queue = queue.Queue()

        # Create handler class with bound queue (via closure)
        # This ensures each server instance has its own isolated queue
        class BoundCallbackHandler(BaseHTTPRequestHandler):
            def do_GET(self):
                """Handle GET request with OAuth callback."""
                # Store the response in the queue (captured from closure)
                response_queue.put(self.path)

                # Send success response
                self.send_response(200)
                self.send_header('Content-type', 'text/html')
                self.end_headers()

                # Show success message
                success_html = """
                <!DOCTYPE html>
                <html>
                <head>
                    <title>Authentication Successful</title>
                    <style>
                        body {
                            font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, Cantarell, sans-serif;
                            display: flex;
                            align-items: center;
                            justify-content: center;
                            min-height: 100vh;
                            margin: 0;
                            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
                        }
                        .container {
                            text-align: center;
                            background: white;
                            padding: 60px 80px;
                            border-radius: 20px;
                            box-shadow: 0 20px 60px rgba(0,0,0,0.3);
                        }
                        .icon {
                            width: 80px;
                            height: 80px;
                            margin: 0 auto 30px;
                            animation: zap 0.6s ease-out;
                        }
                        h1 {
                            color: #2d3748;
                            margin: 0 0 15px 0;
                            font-size: 32px;
                            font-weight: 600;
                        }
                        p {
                            color: #718096;
                            margin: 0;
                            font-size: 16px;
                        }
                        @keyframes zap {
                            0% {
                                transform: scale(0) rotate(-180deg);
                                opacity: 0;
                            }
                            50% {
                                transform: scale(1.2) rotate(10deg);
                            }
                            100% {
                                transform: scale(1) rotate(0deg);
                                opacity: 1;
                            }
                        }
                    </style>
                </head>
                <body>
                    <div class="container">
                        <svg class="icon" viewBox="0 0 200 200" xmlns="http://www.w3.org/2000/svg">
                            <path d="M 110 20 L 70 100 L 95 100 L 80 180 L 140 90 L 110 90 Z"
                                  fill="#000000"
                                  stroke="#000000"
                                  stroke-width="3"
                                  stroke-linejoin="round"/>
                        </svg>
                        <h1>Authentication Successful!</h1>
                        <p>You can close this window and return to your terminal.</p>
                    </div>
                </body>
                </html>
                """
                self.wfile.write(success_html.encode())

            def log_message(self, format, *args):
                """Suppress server log messages."""
                pass

        # Create server with the bound handler
        # Event to signal when server is ready
        server_ready = threading.Event()

        try:
            server = HTTPServer((host, port), BoundCallbackHandler)
        except OSError as e:
            # Port already in use or other binding error
            print(f"Warning: Could not bind to {host}:{port} ({e}), falling back to manual flow")
            return None, response_queue, server_ready

        try:
            # Start server in background thread
            def run_server():
                try:
                    server_ready.set()  # Signal that server is ready
                    server.handle_request()  # Handle one request then stop
                finally:
                    server.server_close()  # Clean up

            server_thread = threading.Thread(target=run_server)
            server_thread.daemon = True
            server_thread.start()

            # Wait for server to be ready (with short timeout)
            if not server_ready.wait(timeout=5):
                print("Warning: Server failed to start, falling back to manual flow")
                server.server_close()
                return None, response_queue, server_ready

            # Give the server a moment to start listening
            import time
            time.sleep(0.1)

            # Return server and queue - caller will wait for response
            return server, response_queue, server_ready
        except Exception as e:
            # Always clean up on error
            server.server_close()
            raise
    
    def _perform_oauth_flow(self) -> Dict:
        """Perform OAuth2 authorization flow to get initial token."""
        print("No auth saved to file yet, lets get logged in!")

        # Create session for authorization
        oauth_session = OAuth2Session(
            client_id=self.config["client_id"],
            scope=self.scope,
            redirect_uri=self.config["redirect_uri"]
        )

        # Get authorization URL
        authorization_url, state = oauth_session.authorization_url(self.authorization_base_url)

        # Parse redirect URI to get host and port
        redirect_uri_parsed = urlparse(self.config["redirect_uri"])
        host = redirect_uri_parsed.hostname or "localhost"
        port = redirect_uri_parsed.port or 8080

        # Start local server to capture callback first
        server, response_queue, server_ready = self._start_callback_server(host, port)

        # Check if server started successfully
        if server is None or not server_ready.is_set():
            # Server failed to start, fall back to manual
            print("\nAutomatic flow failed - server could not start.")
            print(f"Opening browser for authorization...")
            print(f"Please visit: {authorization_url}")
            redirect_response = input("Paste the full redirect URL here: ")
        else:
            # Server started successfully, open browser for OAuth
            print(f"\nOpening browser for authorization...")
            print(f"If browser doesn't open automatically, visit: {authorization_url}")
            webbrowser.open(authorization_url)

            # Wait for callback from queue (with timeout)
            try:
                callback_path = response_queue.get(timeout=120)

                # Success! Construct full redirect URL from path
                # Parse the redirect URI to get base components
                parsed_redirect = urlparse(self.config['redirect_uri'])
                # Parse the callback path (includes query string)
                parsed_callback = urlparse(callback_path)
                # Combine: use base from redirect_uri, path and query from callback
                redirect_response = urlunparse((
                    parsed_redirect.scheme,
                    parsed_redirect.netloc,
                    parsed_callback.path or parsed_redirect.path,
                    '',  # params
                    parsed_callback.query,
                    ''   # fragment
                ))
                print("✓ Authorization received!")
            except queue.Empty:
                # Timeout waiting for callback
                print("\nAutomatic flow timed out waiting for authorization.")
                print(f"Please visit: {authorization_url}")
                redirect_response = input("Paste the full redirect URL here: ")

        # Fetch the access token
        token = oauth_session.fetch_token(
            self.token_url,
            client_secret=self.config["client_secret"],
            authorization_response=redirect_response,
        )

        # Save the token
        self._token_updater(token)
        return token
    
    def authenticate(self) -> OAuth2Session:
        """
        Authenticate and return an OAuth2Session ready for API calls.
        
        This method handles the complete authentication flow:
        1. Load configuration (prompt if needed)
        2. Load existing token or perform OAuth flow
        3. Create authenticated session with auto-refresh
        
        Returns:
            OAuth2Session: Authenticated session for making API requests
        """
        # Load configuration
        self.config = self._load_config()
        
        # Setup client credentials for auto-refresh
        client_creds = {
            "client_id": self.config["client_id"],
            "client_secret": self.config["client_secret"]
        }
        
        # Try to load existing token
        token = self._load_existing_token()
        
        if token is None:
            # No existing token, perform OAuth flow
            token = self._perform_oauth_flow()
        
        # Create authenticated session with auto-refresh
        self.session = OAuth2Session(
            client_id=self.config["client_id"],
            token=token,
            auto_refresh_kwargs=client_creds,
            auto_refresh_url=self.refresh_url,
            token_updater=self._token_updater,
        )
        
        return self.session
    
    def get_session(self) -> OAuth2Session:
        """
        Get the authenticated session.
        
        Returns:
            OAuth2Session: The authenticated session
            
        Raises:
            RuntimeError: If authenticate() hasn't been called yet
        """
        if self.session is None:
            raise RuntimeError("Must call authenticate() first")
        return self.session


def create_oauth_session(
    authorization_base_url: str,
    token_url: str,
    scope: List[str],
    **kwargs
) -> OAuth2Session:
    """
    Convenience function to create and authenticate an OAuth2 session.

    Args:
        authorization_base_url: OAuth2 authorization endpoint
        token_url: OAuth2 token endpoint
        scope: List of OAuth2 scopes to request
        **kwargs: Additional arguments passed to OAuthSession constructor

    Returns:
        OAuth2Session: Authenticated session ready for making API requests
    """
    oauth_client = OAuthSession(authorization_base_url, token_url, scope, **kwargs)
    return oauth_client.authenticate()