gpt-plus-gpt/pkg/auth/codex.go

package auth

import (
	"context"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"strings"

	"gpt-plus/pkg/httpclient"
)

// ObtainCodexTokens performs a lightweight OAuth re-authorization using existing
// session cookies to obtain Codex CLI tokens (with refresh_token) for a specific workspace.
//
// This should be called AFTER activation is complete, when the cookie jar already
// contains valid session cookies from the initial login.
//
// targetWorkspaceID: the workspace/account ID to authorize for. Pass "" to use the default (first) workspace.
func ObtainCodexTokens(
	ctx context.Context,
	client *httpclient.Client,
	deviceID string,
	targetWorkspaceID string,
) (*LoginResult, error) {
	log.Printf("[codex] obtaining Codex CLI tokens for workspace=%s", targetWorkspaceID)

	// Ensure oai-did cookie is set
	cookieURL, _ := url.Parse(oauthIssuer)
	client.GetCookieJar().SetCookies(cookieURL, []*http.Cookie{
		{Name: "oai-did", Value: deviceID},
	})

	// Generate fresh PKCE pair and state for this authorization
	codeVerifier, codeChallenge := generatePKCE()
	state := generateState()

	// Build authorize URL with Codex CLI params
	authorizeParams := url.Values{
		"response_type":              {"code"},
		"client_id":                  {oauthClientID},
		"redirect_uri":              {oauthRedirectURI},
		"scope":                     {oauthScope},
		"code_challenge":            {codeChallenge},
		"code_challenge_method":     {"S256"},
		"state":                     {state},
		"id_token_add_organizations": {"true"},
		"codex_cli_simplified_flow": {"true"},
	}
	authorizeURL := oauthIssuer + "/oauth/authorize?" + authorizeParams.Encode()
	navH := navigateHeaders()

	if targetWorkspaceID == "" {
		// Default workspace: auto-follow redirects and use shortcut if code is returned directly.
		resp, err := client.Get(authorizeURL, navH)
		if err != nil {
			return nil, fmt.Errorf("codex authorize request: %w", err)
		}
		httpclient.ReadBody(resp)

		if resp.Request != nil {
			if code := extractCodeFromURL(resp.Request.URL.String()); code != "" {
				log.Printf("[codex] got code directly from authorize redirect (default workspace)")
				return exchangeCodeForTokens(ctx, client, code, codeVerifier)
			}
		}

		log.Printf("[codex] authorize response: status=%d, url=%s", resp.StatusCode, resp.Request.URL.String())

		// Fallback: go through consent flow
		consentURL := oauthIssuer + "/sign-in-with-chatgpt/codex/consent"
		authCode, err := followConsentRedirects(ctx, client, deviceID, consentURL, "")
		if err != nil {
			return nil, fmt.Errorf("codex consent flow: %w", err)
		}
		return exchangeCodeForTokens(ctx, client, authCode, codeVerifier)
	}

	// ===== Specific workspace requested (e.g. Team) =====
	log.Printf("[codex] specific workspace requested, using manual redirect flow")

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, authorizeURL, nil)
	if err != nil {
		return nil, fmt.Errorf("codex build authorize request: %w", err)
	}
	for k, v := range navH {
		req.Header.Set(k, v)
	}

	// Follow redirects manually. Stop when we see the consent page or a code for the wrong workspace.
	currentURL := authorizeURL
	for i := 0; i < 15; i++ {
		resp, err := client.DoNoRedirect(req)
		if err != nil {
			if code := extractCodeFromURL(currentURL); code != "" {
				log.Printf("[codex] got code from failed redirect (default workspace), ignoring — need workspace %s", targetWorkspaceID)
				break
			}
			return nil, fmt.Errorf("codex authorize redirect: %w", err)
		}
		httpclient.ReadBody(resp)

		if resp.StatusCode >= 300 && resp.StatusCode < 400 {
			loc := resp.Header.Get("Location")
			if loc == "" {
				break
			}
			if !strings.HasPrefix(loc, "http") {
				loc = oauthIssuer + loc
			}

			if code := extractCodeFromURL(loc); code != "" {
				log.Printf("[codex] auto-redirect has code (default workspace), ignoring — need workspace %s", targetWorkspaceID)
				break
			}

			log.Printf("[codex] redirect %d: %d → %s", i+1, resp.StatusCode, truncURL(loc))
			currentURL = loc
			req, _ = http.NewRequestWithContext(ctx, http.MethodGet, loc, nil)
			for k, v := range navH {
				req.Header.Set(k, v)
			}
			continue
		}

		log.Printf("[codex] reached page: status=%d, url=%s", resp.StatusCode, currentURL)
		break
	}

	// Now go through the consent flow with the target workspace.
	consentURL := oauthIssuer + "/sign-in-with-chatgpt/codex/consent"
	authCode, err := followConsentRedirects(ctx, client, deviceID, consentURL, targetWorkspaceID)
	if err != nil {
		return nil, fmt.Errorf("codex consent flow for workspace %s: %w", targetWorkspaceID, err)
	}

	tokens, err := exchangeCodeForTokens(ctx, client, authCode, codeVerifier)
	if err != nil {
		return nil, fmt.Errorf("codex token exchange: %w", err)
	}

	log.Printf("[codex] tokens obtained for workspace %s: has_refresh=%v, account_id=%s",
		targetWorkspaceID, tokens.RefreshToken != "", tokens.ChatGPTAccountID)
	return tokens, nil
}

// truncURL shortens a URL for logging.
func truncURL(u string) string {
	if len(u) > 100 {
		return u[:100] + "..."
	}
	return u
}