google_tools/oauth_manager.py

"""OAuth2 Manager for Google APIs (Gmail, Calendar, Contacts)."""

import json
import os
import webbrowser
from datetime import datetime, timedelta
from http.server import BaseHTTPRequestHandler, HTTPServer
from pathlib import Path
from typing import Optional
from urllib.parse import parse_qs, urlparse

from google.auth.transport.requests import Request
from google.oauth2.credentials import Credentials
from google_auth_oauthlib.flow import InstalledAppFlow

# OAuth scopes
SCOPES = [
    "https://www.googleapis.com/auth/gmail.modify",
    "https://www.googleapis.com/auth/calendar",
    "https://www.googleapis.com/auth/contacts",
]

# Token file location
TOKEN_FILE = Path("config/google_oauth_token.json")
CREDENTIALS_FILE = Path("config/google_credentials.yaml")


class OAuthCallbackHandler(BaseHTTPRequestHandler):
    """Handles OAuth2 callback from browser."""

    authorization_code: Optional[str] = None

    def do_GET(self) -> None:
        """Handle GET request from OAuth callback."""
        query = urlparse(self.path).query
        params = parse_qs(query)

        if "code" in params:
            OAuthCallbackHandler.authorization_code = params["code"][0]
            self.send_response(200)
            self.send_header("Content-type", "text/html")
            self.end_headers()
            self.wfile.write(
                b"<html><body><h1>Authorization successful!</h1>"
                b"<p>You can close this window and return to the terminal.</p>"
                b"</body></html>"
            )
        else:
            self.send_response(400)
            self.send_header("Content-type", "text/html")
            self.end_headers()
            self.wfile.write(
                b"<html><body><h1>Authorization failed</h1>"
                b"<p>No authorization code received.</p></body></html>"
            )

    def log_message(self, format: str, *args) -> None:
        """Suppress logging to keep console clean."""
        pass


class GoogleOAuthManager:
    """Manages OAuth2 authentication for Google APIs."""

    def __init__(
        self,
        token_file: Path = TOKEN_FILE,
        credentials_file: Path = CREDENTIALS_FILE,
    ) -> None:
        self.token_file = token_file
        self.credentials_file = credentials_file
        self.credentials: Optional[Credentials] = None

    def is_authorized(self) -> bool:
        """Check if valid OAuth tokens exist."""
        if not self.token_file.exists():
            return False

        try:
            self.credentials = Credentials.from_authorized_user_file(
                str(self.token_file), SCOPES
            )
            return self.credentials and self.credentials.valid
        except Exception:
            return False

    def get_credentials(self) -> Optional[Credentials]:
        """Get valid OAuth credentials, refreshing if needed.

        Returns:
            Credentials object if authorized, None otherwise.
        """
        # Load existing tokens
        if self.token_file.exists():
            try:
                self.credentials = Credentials.from_authorized_user_file(
                    str(self.token_file), SCOPES
                )
            except Exception as e:
                print(f"[OAuth] Error loading tokens: {e}")
                return None

        # Refresh if expired
        if self.credentials and not self.credentials.valid:
            if self.credentials.expired and self.credentials.refresh_token:
                try:
                    print("[OAuth] Refreshing access token...")
                    self.credentials.refresh(Request())
                    self._save_credentials()
                    print("[OAuth] Token refreshed successfully")
                except Exception as e:
                    print(f"[OAuth] Token refresh failed: {e}")
                    print("[OAuth] Please run: python bot_runner.py --setup-google")
                    return None
            else:
                print("[OAuth] Credentials invalid, re-authorization needed")
                print("[OAuth] Please run: python bot_runner.py --setup-google")
                return None

        return self.credentials

    def run_oauth_flow(self, manual: bool = False) -> bool:
        """Run OAuth2 authorization flow.

        Args:
            manual: If True, requires manual code paste instead of browser callback.

        Returns:
            True if authorization successful, False otherwise.
        """
        if not self.credentials_file.exists():
            print(f"[OAuth] Error: {self.credentials_file} not found")
            print("[OAuth] Please create config/google_credentials.yaml with:")
            print("  client_id: YOUR_CLIENT_ID")
            print("  client_secret: YOUR_CLIENT_SECRET")
            return False

        try:
            # Load client credentials
            import yaml
            with open(self.credentials_file) as f:
                creds_config = yaml.safe_load(f)

            client_config = {
                "installed": {
                    "client_id": creds_config["client_id"],
                    "client_secret": creds_config["client_secret"],
                    "auth_uri": "https://accounts.google.com/o/oauth2/auth",
                    "token_uri": "https://oauth2.googleapis.com/token",
                    "redirect_uris": ["http://localhost:8081"],
                }
            }

            # Both manual and automatic modes use local server
            # (Google deprecated OOB flow in 2022)
            flow = InstalledAppFlow.from_client_config(client_config, SCOPES)
            flow.redirect_uri = "http://localhost:8081"

            auth_url, _ = flow.authorization_url(prompt="consent")

            print("\nPlease visit this URL to authorize:")
            print(f"\n{auth_url}\n")

            # Auto-open browser unless manual mode
            if not manual:
                try:
                    webbrowser.open(auth_url)
                    print("Opening browser automatically...\n")
                except Exception:
                    print("Could not open browser automatically. Please visit the URL above.\n")
            else:
                print("Please open the URL above in your browser manually.\n")

            # Start local server to receive callback
            server = HTTPServer(("localhost", 8081), OAuthCallbackHandler)
            print("Waiting for authorization... (press Ctrl+C to cancel)\n")

            # Wait for callback
            server.handle_request()

            if OAuthCallbackHandler.authorization_code:
                flow.fetch_token(code=OAuthCallbackHandler.authorization_code)
                self.credentials = flow.credentials
            else:
                print("[OAuth] Authorization failed - no code received")
                return False

            # Save tokens
            self._save_credentials()
            print("\n✓ Authorization successful!")
            print(f"✓ Tokens saved to {self.token_file}\n")
            return True

        except Exception as e:
            print(f"[OAuth] Authorization failed: {e}")
            return False

    def _save_credentials(self) -> None:
        """Save credentials to token file with secure permissions."""
        if not self.credentials:
            return

        # Ensure config directory exists
        self.token_file.parent.mkdir(parents=True, exist_ok=True)

        # Write to temp file first (atomic write)
        temp_file = self.token_file.with_suffix(".tmp")
        with open(temp_file, "w") as f:
            f.write(self.credentials.to_json())

        # Set file permissions to 600 (owner read/write only) on Unix systems
        if os.name != "nt":  # Not Windows
            os.chmod(temp_file, 0o600)

        # Atomic rename
        temp_file.replace(self.token_file)