Initial sanitized code sync

2026-03-15 20:48:19 +08:00
commit 17ee51ba04
126 changed files with 22546 additions and 0 deletions
--- a/pkg/auth/codex.go
+++ b/pkg/auth/codex.go
@@ -0,0 +1,154 @@
+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
+}