codexautopool/backend/internal/auth/chromedp.go

package auth

import (
	"context"
	"fmt"
	"strings"
	"time"

	"codex-pool/internal/proxyutil"

	"github.com/chromedp/cdproto/fetch"
	"github.com/chromedp/cdproto/network"
	"github.com/chromedp/chromedp"
)

// CompleteWithChromedp 使用 chromedp 完成 S2A OAuth 授权
func CompleteWithChromedp(authURL, email, password, teamID string, headless bool, proxy string) (string, error) {
	var proxyServer string
	var proxyUser string
	var proxyPass string
	if proxy != "" {
		info, err := proxyutil.Parse(proxy)
		if err != nil {
			return "", fmt.Errorf("代理格式错误: %v", err)
		}
		if info.Server != nil {
			proxyServer = info.Server.String()
		}
		proxyUser = info.Username
		proxyPass = info.Password
	}

	opts := append(chromedp.DefaultExecAllocatorOptions[:],
		chromedp.Flag("headless", headless),
		chromedp.Flag("disable-gpu", true),
		chromedp.Flag("no-sandbox", true),
		chromedp.Flag("disable-dev-shm-usage", true),
		chromedp.Flag("disable-blink-features", "AutomationControlled"),
		chromedp.UserAgent("Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/133.0.0.0 Safari/537.36"),
	)

	if proxyServer != "" {
		opts = append(opts, chromedp.ProxyServer(proxyServer))
	}

	allocCtx, cancel := chromedp.NewExecAllocator(context.Background(), opts...)
	defer cancel()

	ctx, cancel := chromedp.NewContext(allocCtx)
	defer cancel()

	// 增加超时时间到 180 秒
	ctx, cancel = context.WithTimeout(ctx, 180*time.Second)
	defer cancel()

	var callbackURL string

	// 只在代理需要认证时才启用 Fetch 域
	if proxyServer != "" && proxyUser != "" {
		chromedp.ListenTarget(ctx, func(ev interface{}) {
			switch ev := ev.(type) {
			case *fetch.EventRequestPaused:
				// Fetch domain pauses requests; we must continue them to avoid stalling navigation.
				reqID := ev.RequestID
				go func() { _ = fetch.ContinueRequest(reqID).Do(ctx) }()

			case *fetch.EventAuthRequired:
				reqID := ev.RequestID
				source := fetch.AuthChallengeSourceServer
				if ev.AuthChallenge != nil {
					source = ev.AuthChallenge.Source
				}

				go func() {
					resp := &fetch.AuthChallengeResponse{Response: fetch.AuthChallengeResponseResponseDefault}
					if source == fetch.AuthChallengeSourceProxy {
						if proxyUser != "" {
							resp.Response = fetch.AuthChallengeResponseResponseProvideCredentials
							resp.Username = proxyUser
							resp.Password = proxyPass
						} else {
							// Fail fast if the proxy requires auth but user didn't provide credentials.
							resp.Response = fetch.AuthChallengeResponseResponseCancelAuth
						}
					}
					_ = fetch.ContinueWithAuth(reqID, resp).Do(ctx)
				}()
			}
		})
	}

	// 监听回调 URL
	chromedp.ListenTarget(ctx, func(ev interface{}) {
		if ev, ok := ev.(*network.EventRequestWillBeSent); ok {
			url := ev.Request.URL
			if strings.Contains(url, "localhost") && strings.Contains(url, "code=") {
				callbackURL = url
			}
		}
	})

	// 构建运行任务
	tasks := []chromedp.Action{
		network.Enable(),
		chromedp.Navigate(authURL),
		chromedp.WaitReady("body"),
	}

	// 只在代理需要认证时才启用 Fetch 域
	if proxyServer != "" && proxyUser != "" {
		tasks = append([]chromedp.Action{fetch.Enable().WithHandleAuthRequests(true)}, tasks...)
	}

	err := chromedp.Run(ctx, tasks...)
	if err != nil {
		return "", fmt.Errorf("访问失败: %v", err)
	}

	time.Sleep(2 * time.Second)

	if callbackURL != "" {
		return ExtractCodeFromCallbackURL(callbackURL), nil
	}

	var currentURL string
	_ = chromedp.Run(ctx, chromedp.Location(&currentURL))

	if strings.Contains(currentURL, "code=") {
		return ExtractCodeFromCallbackURL(currentURL), nil
	}

	time.Sleep(1 * time.Second)

	emailSelectors := []string{
		`input[name="email"]`,
		`input[type="email"]`,
		`input[name="username"]`,
	}

	var emailFilled bool
	for _, sel := range emailSelectors {
		err = chromedp.Run(ctx, chromedp.WaitVisible(sel, chromedp.ByQuery))
		if err == nil {
			err = chromedp.Run(ctx,
				chromedp.Clear(sel, chromedp.ByQuery),
				chromedp.SendKeys(sel, email, chromedp.ByQuery),
			)
			if err == nil {
				emailFilled = true
				break
			}
		}
	}

	if !emailFilled {
		return "", fmt.Errorf("未找到邮箱输入框")
	}

	time.Sleep(300 * time.Millisecond)

	buttonSelectors := []string{
		`button[type="submit"]`,
		`button[data-testid="login-button"]`,
		`button.continue-btn`,
		`input[type="submit"]`,
	}

	for _, sel := range buttonSelectors {
		err = chromedp.Run(ctx, chromedp.Click(sel, chromedp.ByQuery))
		if err == nil {
			break
		}
	}

	time.Sleep(1500 * time.Millisecond)

	if callbackURL != "" {
		return ExtractCodeFromCallbackURL(callbackURL), nil
	}

	_ = chromedp.Run(ctx, chromedp.Location(&currentURL))
	if strings.Contains(currentURL, "code=") {
		return ExtractCodeFromCallbackURL(currentURL), nil
	}

	passwordSelectors := []string{
		`input[name="current-password"]`,
		`input[name="password"]`,
		`input[type="password"]`,
	}

	var passwordFilled bool
	for _, sel := range passwordSelectors {
		err = chromedp.Run(ctx, chromedp.WaitVisible(sel, chromedp.ByQuery))
		if err == nil {
			err = chromedp.Run(ctx,
				chromedp.Clear(sel, chromedp.ByQuery),
				chromedp.SendKeys(sel, password, chromedp.ByQuery),
			)
			if err == nil {
				passwordFilled = true
				break
			}
		}
	}

	if !passwordFilled {
		return "", fmt.Errorf("未找到密码输入框")
	}

	time.Sleep(300 * time.Millisecond)

	for _, sel := range buttonSelectors {
		err = chromedp.Run(ctx, chromedp.Click(sel, chromedp.ByQuery))
		if err == nil {
			break
		}
	}

	for i := 0; i < 30; i++ {
		time.Sleep(500 * time.Millisecond)

		if callbackURL != "" {
			return ExtractCodeFromCallbackURL(callbackURL), nil
		}

		var url string
		if err := chromedp.Run(ctx, chromedp.Location(&url)); err == nil {
			if strings.Contains(url, "code=") {
				return ExtractCodeFromCallbackURL(url), nil
			}

			if strings.Contains(url, "consent") {
				for _, sel := range buttonSelectors {
					err = chromedp.Run(ctx, chromedp.Click(sel, chromedp.ByQuery))
					if err == nil {
						break
					}
				}
				time.Sleep(1 * time.Second)
			}

			if strings.Contains(url, "authorize") && teamID != "" {
				err = chromedp.Run(ctx,
					chromedp.Click(fmt.Sprintf(`[data-workspace-id="%s"], [data-account-id="%s"]`, teamID, teamID), chromedp.ByQuery),
				)
			}
		}
	}

	if callbackURL != "" {
		return ExtractCodeFromCallbackURL(callbackURL), nil
	}

	return "", fmt.Errorf("授权超时")
}