Add openai codex support

internal/cmd/login.go

@@ -5,27 +5,33 @@ package cmd
 
 import (
 	"context"
-	"github.com/luispater/CLIProxyAPI/internal/auth"
+	"os"
+
+	"github.com/luispater/CLIProxyAPI/internal/auth/gemini"
 	"github.com/luispater/CLIProxyAPI/internal/client"
 	"github.com/luispater/CLIProxyAPI/internal/config"
 	log "github.com/sirupsen/logrus"
-	"os"
 )
 
 // DoLogin handles the entire user login and setup process.
 // It authenticates the user, sets up the user's project, checks API enablement,
 // and saves the token for future use.
-func DoLogin(cfg *config.Config, projectID string) {
+func DoLogin(cfg *config.Config, projectID string, options *LoginOptions) {
+	if options == nil {
+		options = &LoginOptions{}
+	}
+
 	var err error
-	var ts auth.TokenStorage
+	var ts gemini.GeminiTokenStorage
 	if projectID != "" {
 		ts.ProjectID = projectID
 	}
 
 	// Initialize an authenticated HTTP client. This will trigger the OAuth flow if necessary.
 	clientCtx := context.Background()
-	log.Info("Initializing authentication...")
-	httpClient, errGetClient := auth.GetAuthenticatedClient(clientCtx, &ts, cfg)
+	log.Info("Initializing Google authentication...")
+	geminiAuth := gemini.NewGeminiAuth()
+	httpClient, errGetClient := geminiAuth.GetAuthenticatedClient(clientCtx, &ts, cfg, options.NoBrowser)
 	if errGetClient != nil {
 		log.Fatalf("failed to get authenticated client: %v", errGetClient)
 		return
@@ -33,7 +39,7 @@ func DoLogin(cfg *config.Config, projectID string) {
 	log.Info("Authentication successful.")
 
 	// Initialize the API client.
-	cliClient := client.NewClient(httpClient, &ts, cfg)
+	cliClient := client.NewGeminiClient(httpClient, &ts, cfg)
 
 	// Perform the user setup process.
 	err = cliClient.SetupUser(clientCtx, ts.Email, projectID)

internal/cmd/openai_login.go

@@ -0,0 +1,173 @@
+package cmd
+
+import (
+	"context"
+	"crypto/rand"
+	"encoding/hex"
+	"fmt"
+	"net/http"
+	"os"
+	"strings"
+	"time"
+
+	"github.com/luispater/CLIProxyAPI/internal/auth/codex"
+	"github.com/luispater/CLIProxyAPI/internal/browser"
+	"github.com/luispater/CLIProxyAPI/internal/client"
+	"github.com/luispater/CLIProxyAPI/internal/config"
+	log "github.com/sirupsen/logrus"
+)
+
+// LoginOptions contains options for login
+type LoginOptions struct {
+	NoBrowser bool
+}
+
+// DoCodexLogin handles the Codex OAuth login process
+func DoCodexLogin(cfg *config.Config, options *LoginOptions) {
+	if options == nil {
+		options = &LoginOptions{}
+	}
+
+	ctx := context.Background()
+
+	log.Info("Initializing Codex authentication...")
+
+	// Generate PKCE codes
+	pkceCodes, err := codex.GeneratePKCECodes()
+	if err != nil {
+		log.Fatalf("Failed to generate PKCE codes: %v", err)
+		return
+	}
+
+	// Generate random state parameter
+	state, err := generateRandomState()
+	if err != nil {
+		log.Fatalf("Failed to generate state parameter: %v", err)
+		return
+	}
+
+	// Initialize OAuth server
+	oauthServer := codex.NewOAuthServer(1455)
+
+	// Start OAuth callback server
+	if err = oauthServer.Start(ctx); err != nil {
+		if strings.Contains(err.Error(), "already in use") {
+			authErr := codex.NewAuthenticationError(codex.ErrPortInUse, err)
+			log.Error(codex.GetUserFriendlyMessage(authErr))
+			os.Exit(13) // Exit code 13 for port-in-use error
+		}
+		authErr := codex.NewAuthenticationError(codex.ErrServerStartFailed, err)
+		log.Fatalf("Failed to start OAuth callback server: %v", authErr)
+		return
+	}
+	defer func() {
+		if err = oauthServer.Stop(ctx); err != nil {
+			log.Warnf("Failed to stop OAuth server: %v", err)
+		}
+	}()
+
+	// Initialize Codex auth service
+	openaiAuth := codex.NewCodexAuth(cfg)
+
+	// Generate authorization URL
+	authURL, err := openaiAuth.GenerateAuthURL(state, pkceCodes)
+	if err != nil {
+		log.Fatalf("Failed to generate authorization URL: %v", err)
+		return
+	}
+
+	// Open browser or display URL
+	if !options.NoBrowser {
+		log.Info("Opening browser for authentication...")
+
+		// Check if browser is available
+		if !browser.IsAvailable() {
+			log.Warn("No browser available on this system")
+			log.Infof("Please manually open this URL in your browser:\n\n%s\n", authURL)
+		} else {
+			if err = browser.OpenURL(authURL); err != nil {
+				authErr := codex.NewAuthenticationError(codex.ErrBrowserOpenFailed, err)
+				log.Warn(codex.GetUserFriendlyMessage(authErr))
+				log.Infof("Please manually open this URL in your browser:\n\n%s\n", authURL)
+
+				// Log platform info for debugging
+				platformInfo := browser.GetPlatformInfo()
+				log.Debugf("Browser platform info: %+v", platformInfo)
+			} else {
+				log.Debug("Browser opened successfully")
+			}
+		}
+	} else {
+		log.Infof("Please open this URL in your browser:\n\n%s\n", authURL)
+	}
+
+	log.Info("Waiting for authentication callback...")
+
+	// Wait for OAuth callback
+	result, err := oauthServer.WaitForCallback(5 * time.Minute)
+	if err != nil {
+		if strings.Contains(err.Error(), "timeout") {
+			authErr := codex.NewAuthenticationError(codex.ErrCallbackTimeout, err)
+			log.Error(codex.GetUserFriendlyMessage(authErr))
+		} else {
+			log.Errorf("Authentication failed: %v", err)
+		}
+		return
+	}
+
+	if result.Error != "" {
+		oauthErr := codex.NewOAuthError(result.Error, "", http.StatusBadRequest)
+		log.Error(codex.GetUserFriendlyMessage(oauthErr))
+		return
+	}
+
+	// Validate state parameter
+	if result.State != state {
+		authErr := codex.NewAuthenticationError(codex.ErrInvalidState, fmt.Errorf("expected %s, got %s", state, result.State))
+		log.Error(codex.GetUserFriendlyMessage(authErr))
+		return
+	}
+
+	log.Debug("Authorization code received, exchanging for tokens...")
+
+	// Exchange authorization code for tokens
+	authBundle, err := openaiAuth.ExchangeCodeForTokens(ctx, result.Code, pkceCodes)
+	if err != nil {
+		authErr := codex.NewAuthenticationError(codex.ErrCodeExchangeFailed, err)
+		log.Errorf("Failed to exchange authorization code for tokens: %v", authErr)
+		log.Debug("This may be due to network issues or invalid authorization code")
+		return
+	}
+
+	// Create token storage
+	tokenStorage := openaiAuth.CreateTokenStorage(authBundle)
+
+	// Initialize Codex client
+	openaiClient, err := client.NewCodexClient(cfg, tokenStorage)
+	if err != nil {
+		log.Fatalf("Failed to initialize Codex client: %v", err)
+		return
+	}
+
+	// Save token storage
+	if err = openaiClient.SaveTokenToFile(); err != nil {
+		log.Fatalf("Failed to save authentication tokens: %v", err)
+		return
+	}
+
+	log.Info("Authentication successful!")
+	if authBundle.APIKey != "" {
+		log.Info("API key obtained and saved")
+	}
+
+	log.Info("You can now use Codex services through this CLI")
+}
+
+// generateRandomState generates a cryptographically secure random state parameter
+func generateRandomState() (string, error) {
+	bytes := make([]byte, 16)
+	if _, err := rand.Read(bytes); err != nil {
+		return "", fmt.Errorf("failed to generate random bytes: %w", err)
+	}
+	return hex.EncodeToString(bytes), nil
+}

internal/cmd/run.go

@@ -8,29 +8,37 @@ package cmd
 import (
 	"context"
 	"encoding/json"
-	"github.com/luispater/CLIProxyAPI/internal/api"
-	"github.com/luispater/CLIProxyAPI/internal/auth"
-	"github.com/luispater/CLIProxyAPI/internal/client"
-	"github.com/luispater/CLIProxyAPI/internal/config"
-	"github.com/luispater/CLIProxyAPI/internal/util"
-	"github.com/luispater/CLIProxyAPI/internal/watcher"
-	log "github.com/sirupsen/logrus"
 	"io/fs"
 	"net/http"
 	"os"
 	"os/signal"
 	"path/filepath"
 	"strings"
+	"sync"
 	"syscall"
 	"time"
+
+	"github.com/luispater/CLIProxyAPI/internal/api"
+	"github.com/luispater/CLIProxyAPI/internal/auth/codex"
+	"github.com/luispater/CLIProxyAPI/internal/auth/gemini"
+	"github.com/luispater/CLIProxyAPI/internal/client"
+	"github.com/luispater/CLIProxyAPI/internal/config"
+	"github.com/luispater/CLIProxyAPI/internal/util"
+	"github.com/luispater/CLIProxyAPI/internal/watcher"
+	log "github.com/sirupsen/logrus"
+	"github.com/tidwall/gjson"
 )
 
 // StartService initializes and starts the main API proxy service.
 // It loads all available authentication tokens, creates a pool of clients,
 // starts the API server, and handles graceful shutdown signals.
+//
+// Parameters:
+//   - cfg: The application configuration
+//   - configPath: The path to the configuration file
 func StartService(cfg *config.Config, configPath string) {
 	// Create a pool of API clients, one for each token file found.
-	cliClients := make([]*client.Client, 0)
+	cliClients := make([]client.Client, 0)
 	err := filepath.Walk(cfg.AuthDir, func(path string, info fs.FileInfo, err error) error {
 		if err != nil {
 			return err
@@ -39,31 +47,51 @@ func StartService(cfg *config.Config, configPath string) {
 		// Process only JSON files in the auth directory.
 		if !info.IsDir() && strings.HasSuffix(info.Name(), ".json") {
 			log.Debugf("Loading token from: %s", path)
-			f, errOpen := os.Open(path)
-			if errOpen != nil {
-				return errOpen
+			data, errReadFile := os.ReadFile(path)
+			if errReadFile != nil {
+				return errReadFile
 			}
-			defer func() {
-				_ = f.Close()
-			}()
 
-			// Decode the token storage file.
-			var ts auth.TokenStorage
-			if err = json.NewDecoder(f).Decode(&ts); err == nil {
-				// For each valid token, create an authenticated client.
-				clientCtx := context.Background()
-				log.Info("Initializing authentication for token...")
-				httpClient, errGetClient := auth.GetAuthenticatedClient(clientCtx, &ts, cfg)
-				if errGetClient != nil {
-					// Log fatal will exit, but we return the error for completeness.
-					log.Fatalf("failed to get authenticated client for token %s: %v", path, errGetClient)
-					return errGetClient
+			tokenType := "gemini"
+			typeResult := gjson.GetBytes(data, "type")
+			if typeResult.Exists() {
+				tokenType = typeResult.String()
+			}
+
+			clientCtx := context.Background()
+
+			if tokenType == "gemini" {
+				var ts gemini.GeminiTokenStorage
+				if err = json.Unmarshal(data, &ts); err == nil {
+					// For each valid token, create an authenticated client.
+					log.Info("Initializing gemini authentication for token...")
+					geminiAuth := gemini.NewGeminiAuth()
+					httpClient, errGetClient := geminiAuth.GetAuthenticatedClient(clientCtx, &ts, cfg)
+					if errGetClient != nil {
+						// Log fatal will exit, but we return the error for completeness.
+						log.Fatalf("failed to get authenticated client for token %s: %v", path, errGetClient)
+						return errGetClient
+					}
+					log.Info("Authentication successful.")
+
+					// Add the new client to the pool.
+					cliClient := client.NewGeminiClient(httpClient, &ts, cfg)
+					cliClients = append(cliClients, cliClient)
+				}
+			} else if tokenType == "codex" {
+				var ts codex.CodexTokenStorage
+				if err = json.Unmarshal(data, &ts); err == nil {
+					// For each valid token, create an authenticated client.
+					log.Info("Initializing codex authentication for token...")
+					codexClient, errGetClient := client.NewCodexClient(cfg, &ts)
+					if errGetClient != nil {
+						// Log fatal will exit, but we return the error for completeness.
+						log.Fatalf("failed to get authenticated client for token %s: %v", path, errGetClient)
+						return errGetClient
+					}
+					log.Info("Authentication successful.")
+					cliClients = append(cliClients, codexClient)
 				}
-				log.Info("Authentication successful.")
-
-				// Add the new client to the pool.
-				cliClient := client.NewClient(httpClient, &ts, cfg)
-				cliClients = append(cliClients, cliClient)
 			}
 		}
 		return nil
@@ -74,13 +102,10 @@ func StartService(cfg *config.Config, configPath string) {
 
 	if len(cfg.GlAPIKey) > 0 {
 		for i := 0; i < len(cfg.GlAPIKey); i++ {
-			httpClient, errSetProxy := util.SetProxy(cfg, &http.Client{})
-			if errSetProxy != nil {
-				log.Fatalf("set proxy failed: %v", errSetProxy)
-			}
+			httpClient := util.SetProxy(cfg, &http.Client{})
 
 			log.Debug("Initializing with Generative Language API key...")
-			cliClient := client.NewClient(httpClient, nil, cfg, cfg.GlAPIKey[i])
+			cliClient := client.NewGeminiClient(httpClient, nil, cfg, cfg.GlAPIKey[i])
 			cliClients = append(cliClients, cliClient)
 		}
 	}
@@ -101,7 +126,7 @@ func StartService(cfg *config.Config, configPath string) {
 	log.Info("API server started successfully")
 
 	// Setup file watcher for config and auth directory changes
-	fileWatcher, errNewWatcher := watcher.NewWatcher(configPath, cfg.AuthDir, func(newClients []*client.Client, newCfg *config.Config) {
+	fileWatcher, errNewWatcher := watcher.NewWatcher(configPath, cfg.AuthDir, func(newClients []client.Client, newCfg *config.Config) {
 		// Update the API server with new clients and configuration
 		apiServer.UpdateClients(newClients, newCfg)
 	})
@@ -132,12 +157,50 @@ func StartService(cfg *config.Config, configPath string) {
 	sigChan := make(chan os.Signal, 1)
 	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
 
+	// Background token refresh ticker for Codex clients
+	ctxRefresh, cancelRefresh := context.WithCancel(context.Background())
+	var wgRefresh sync.WaitGroup
+	wgRefresh.Add(1)
+	go func() {
+		defer wgRefresh.Done()
+		ticker := time.NewTicker(1 * time.Hour)
+		defer ticker.Stop()
+		checkAndRefresh := func() {
+			for i := 0; i < len(cliClients); i++ {
+				if codexCli, ok := cliClients[i].(*client.CodexClient); ok {
+					ts := codexCli.TokenStorage().(*codex.CodexTokenStorage)
+					if ts != nil && ts.Expire != "" {
+						if expTime, errParse := time.Parse(time.RFC3339, ts.Expire); errParse == nil {
+							if time.Until(expTime) <= 5*24*time.Hour {
+								log.Debugf("refreshing codex tokens for %s", codexCli.GetEmail())
+								_ = codexCli.RefreshTokens(ctxRefresh)
+							}
+						}
+					}
+				}
+			}
+		}
+		// Initial check on start
+		checkAndRefresh()
+		for {
+			select {
+			case <-ctxRefresh.Done():
+				return
+			case <-ticker.C:
+				checkAndRefresh()
+			}
+		}
+	}()
+
 	// Main loop to wait for shutdown signal.
 	for {
 		select {
 		case <-sigChan:
 			log.Debugf("Received shutdown signal. Cleaning up...")
 
+			cancelRefresh()
+			wgRefresh.Wait()
+
 			// Create a context with a timeout for the shutdown process.
 			ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
 			_ = cancel
@@ -150,8 +213,6 @@ func StartService(cfg *config.Config, configPath string) {
 			log.Debugf("Cleanup completed. Exiting...")
 			os.Exit(0)
 		case <-time.After(5 * time.Second):
-			// This case is currently empty and acts as a periodic check.
-			// It could be used for periodic tasks in the future.
 		}
 	}
 }