Initial sanitized code sync
This commit is contained in:
zqq61
97
internal/task/manager.go
Normal file
97
internal/task/manager.go
Normal file
@@ -0,0 +1,97 @@
|
||||
package task
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"log"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"gpt-plus/internal/db"
|
||||
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// TaskManager controls the lifecycle of task execution.
|
||||
// Only one task may run at a time (single-task serial constraint).
|
||||
type TaskManager struct {
|
||||
mu sync.Mutex
|
||||
current *TaskRunner
|
||||
gormDB *gorm.DB
|
||||
}
|
||||
|
||||
func NewTaskManager(d *gorm.DB) *TaskManager {
|
||||
return &TaskManager{gormDB: d}
|
||||
}
|
||||
|
||||
// Init marks any leftover running/stopping tasks as interrupted on startup.
|
||||
func (m *TaskManager) Init() {
|
||||
m.gormDB.Model(&db.Task{}).
|
||||
Where("status IN ?", []string{StatusRunning, StatusStopping}).
|
||||
Updates(map[string]interface{}{
|
||||
"status": StatusInterrupted,
|
||||
"stopped_at": time.Now(),
|
||||
})
|
||||
log.Println("[task-manager] init: marked leftover tasks as interrupted")
|
||||
}
|
||||
|
||||
func (m *TaskManager) Start(taskID string) error {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
if m.current != nil && m.current.IsRunning() {
|
||||
return fmt.Errorf("已有任务正在运行 (ID: %s)", m.current.taskID)
|
||||
}
|
||||
|
||||
var t db.Task
|
||||
if err := m.gormDB.First(&t, "id = ?", taskID).Error; err != nil {
|
||||
return fmt.Errorf("任务不存在: %w", err)
|
||||
}
|
||||
if t.Status != StatusPending && t.Status != StatusStopped && t.Status != StatusInterrupted {
|
||||
return fmt.Errorf("任务状态不允许启动: %s", t.Status)
|
||||
}
|
||||
|
||||
runner, err := NewTaskRunner(taskID, m.gormDB)
|
||||
if err != nil {
|
||||
return fmt.Errorf("创建任务运行器失败: %w", err)
|
||||
}
|
||||
m.current = runner
|
||||
|
||||
go func() {
|
||||
runner.Run()
|
||||
m.mu.Lock()
|
||||
if m.current == runner {
|
||||
m.current = nil
|
||||
}
|
||||
m.mu.Unlock()
|
||||
}()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *TaskManager) Stop(taskID string) error {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
if m.current == nil || m.current.taskID != taskID {
|
||||
return fmt.Errorf("该任务未在运行")
|
||||
}
|
||||
m.current.GracefulStop()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *TaskManager) ForceStop(taskID string) error {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
if m.current == nil || m.current.taskID != taskID {
|
||||
return fmt.Errorf("该任务未在运行")
|
||||
}
|
||||
m.current.ForceStop()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *TaskManager) IsRunning() bool {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
return m.current != nil && m.current.IsRunning()
|
||||
}
|
||||
144
internal/task/manager_test.go
Normal file
144
internal/task/manager_test.go
Normal file
@@ -0,0 +1,144 @@
|
||||
package task
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gpt-plus/internal/db"
|
||||
|
||||
"github.com/glebarez/sqlite"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/logger"
|
||||
)
|
||||
|
||||
func setupTaskTestDB(t *testing.T) *gorm.DB {
|
||||
t.Helper()
|
||||
d, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{
|
||||
Logger: logger.Default.LogMode(logger.Silent),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("open db: %v", err)
|
||||
}
|
||||
d.AutoMigrate(&db.SystemConfig{}, &db.EmailRecord{}, &db.Task{}, &db.TaskLog{},
|
||||
&db.CardCode{}, &db.Card{}, &db.Account{})
|
||||
db.DB = d
|
||||
return d
|
||||
}
|
||||
|
||||
func TestTaskManagerInit(t *testing.T) {
|
||||
d := setupTaskTestDB(t)
|
||||
|
||||
// Create a "leftover" running task
|
||||
d.Create(&db.Task{ID: "left-1", Type: "plus", Status: StatusRunning, TotalCount: 10})
|
||||
d.Create(&db.Task{ID: "left-2", Type: "team", Status: StatusStopping, TotalCount: 5})
|
||||
|
||||
tm := NewTaskManager(d)
|
||||
tm.Init()
|
||||
|
||||
var t1, t2 db.Task
|
||||
d.First(&t1, "id = ?", "left-1")
|
||||
d.First(&t2, "id = ?", "left-2")
|
||||
|
||||
if t1.Status != StatusInterrupted {
|
||||
t.Fatalf("left-1 status = %q, want interrupted", t1.Status)
|
||||
}
|
||||
if t2.Status != StatusInterrupted {
|
||||
t.Fatalf("left-2 status = %q, want interrupted", t2.Status)
|
||||
}
|
||||
if t1.StoppedAt == nil {
|
||||
t.Fatal("left-1 stopped_at should be set")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskManagerStartNonexistent(t *testing.T) {
|
||||
d := setupTaskTestDB(t)
|
||||
tm := NewTaskManager(d)
|
||||
|
||||
err := tm.Start("nonexistent-id")
|
||||
if err == nil {
|
||||
t.Fatal("expected error for nonexistent task")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskManagerStartWrongStatus(t *testing.T) {
|
||||
d := setupTaskTestDB(t)
|
||||
d.Create(&db.Task{ID: "completed-1", Type: "plus", Status: StatusCompleted, TotalCount: 10})
|
||||
|
||||
tm := NewTaskManager(d)
|
||||
err := tm.Start("completed-1")
|
||||
if err == nil {
|
||||
t.Fatal("expected error for completed task")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskManagerStopNoRunning(t *testing.T) {
|
||||
d := setupTaskTestDB(t)
|
||||
tm := NewTaskManager(d)
|
||||
|
||||
err := tm.Stop("any-id")
|
||||
if err == nil {
|
||||
t.Fatal("expected error when no task running")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskManagerForceStopNoRunning(t *testing.T) {
|
||||
d := setupTaskTestDB(t)
|
||||
tm := NewTaskManager(d)
|
||||
|
||||
err := tm.ForceStop("any-id")
|
||||
if err == nil {
|
||||
t.Fatal("expected error when no task running")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskManagerIsRunning(t *testing.T) {
|
||||
d := setupTaskTestDB(t)
|
||||
tm := NewTaskManager(d)
|
||||
|
||||
if tm.IsRunning() {
|
||||
t.Fatal("should not be running initially")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskManagerStopWrongID(t *testing.T) {
|
||||
d := setupTaskTestDB(t)
|
||||
tm := NewTaskManager(d)
|
||||
|
||||
// Simulate a running task by setting current directly
|
||||
tm.current = &TaskRunner{taskID: "real-id"}
|
||||
tm.current.running.Store(true)
|
||||
|
||||
err := tm.Stop("wrong-id")
|
||||
if err == nil {
|
||||
t.Fatal("expected error for wrong task ID")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskManagerInitDoesNotAffectPending(t *testing.T) {
|
||||
d := setupTaskTestDB(t)
|
||||
d.Create(&db.Task{ID: "pending-1", Type: "plus", Status: StatusPending, TotalCount: 5})
|
||||
|
||||
tm := NewTaskManager(d)
|
||||
tm.Init()
|
||||
|
||||
var task db.Task
|
||||
d.First(&task, "id = ?", "pending-1")
|
||||
if task.Status != StatusPending {
|
||||
t.Fatalf("pending task should stay pending, got %q", task.Status)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskManagerInitDoesNotAffectCompleted(t *testing.T) {
|
||||
d := setupTaskTestDB(t)
|
||||
stopped := time.Now()
|
||||
d.Create(&db.Task{ID: "done-1", Type: "plus", Status: StatusCompleted, TotalCount: 10, StoppedAt: &stopped})
|
||||
|
||||
tm := NewTaskManager(d)
|
||||
tm.Init()
|
||||
|
||||
var task db.Task
|
||||
d.First(&task, "id = ?", "done-1")
|
||||
if task.Status != StatusCompleted {
|
||||
t.Fatalf("completed task should stay completed, got %q", task.Status)
|
||||
}
|
||||
}
|
||||
229
internal/task/membership.go
Normal file
229
internal/task/membership.go
Normal file
@@ -0,0 +1,229 @@
|
||||
package task
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"gpt-plus/pkg/auth"
|
||||
"gpt-plus/pkg/chatgpt"
|
||||
)
|
||||
|
||||
const (
|
||||
finalMembershipPolls = 5
|
||||
finalMembershipPollDelay = 2 * time.Second
|
||||
)
|
||||
|
||||
type finalMembershipState struct {
|
||||
Personal *chatgpt.AccountInfo
|
||||
Workspace *chatgpt.AccountInfo
|
||||
}
|
||||
|
||||
func (s *finalMembershipState) plusActive() bool {
|
||||
return plusMembershipActive(s.Personal)
|
||||
}
|
||||
|
||||
func (s *finalMembershipState) teamActive() bool {
|
||||
return teamMembershipActive(s.Workspace)
|
||||
}
|
||||
|
||||
func (s *finalMembershipState) satisfied(taskType string) bool {
|
||||
switch taskType {
|
||||
case TaskTypePlus:
|
||||
return s.plusActive()
|
||||
case TaskTypeTeam:
|
||||
return s.teamActive()
|
||||
case TaskTypeBoth:
|
||||
return s.plusActive() && s.teamActive()
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func (s *finalMembershipState) resultPlanForTask(taskType string) string {
|
||||
if s != nil && s.satisfied(taskType) {
|
||||
return taskType
|
||||
}
|
||||
return s.actualPlan()
|
||||
}
|
||||
|
||||
func (s *finalMembershipState) actualPlan() string {
|
||||
switch {
|
||||
case s == nil:
|
||||
return "unknown"
|
||||
case s.plusActive() && s.teamActive():
|
||||
return TaskTypeBoth
|
||||
case s.teamActive():
|
||||
return "team"
|
||||
case s.plusActive():
|
||||
return "plus"
|
||||
case s.Personal != nil && s.Personal.PlanType != "":
|
||||
return s.Personal.PlanType
|
||||
case s.Workspace != nil && s.Workspace.PlanType != "":
|
||||
return s.Workspace.PlanType
|
||||
default:
|
||||
return "unknown"
|
||||
}
|
||||
}
|
||||
|
||||
func (s *finalMembershipState) describe() string {
|
||||
return fmt.Sprintf("personal=%s workspace=%s", describeMembership(s.Personal), describeMembership(s.Workspace))
|
||||
}
|
||||
|
||||
func plusMembershipActive(info *chatgpt.AccountInfo) bool {
|
||||
return info != nil &&
|
||||
info.Structure == "personal" &&
|
||||
info.PlanType == "plus" &&
|
||||
info.HasActiveSubscription &&
|
||||
info.SubscriptionID != ""
|
||||
}
|
||||
|
||||
func teamMembershipActive(info *chatgpt.AccountInfo) bool {
|
||||
return info != nil &&
|
||||
info.Structure == "workspace" &&
|
||||
info.PlanType == "team"
|
||||
}
|
||||
|
||||
func describeMembership(info *chatgpt.AccountInfo) string {
|
||||
if info == nil {
|
||||
return "none"
|
||||
}
|
||||
return fmt.Sprintf("%s/%s(active=%v sub=%t id=%s)",
|
||||
info.Structure, info.PlanType, info.HasActiveSubscription, info.SubscriptionID != "", info.AccountID)
|
||||
}
|
||||
|
||||
func selectPersonalMembership(accounts []*chatgpt.AccountInfo) *chatgpt.AccountInfo {
|
||||
for _, acct := range accounts {
|
||||
if acct.Structure == "personal" {
|
||||
return acct
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func selectWorkspaceMembership(accounts []*chatgpt.AccountInfo, preferredID string) *chatgpt.AccountInfo {
|
||||
if preferredID != "" {
|
||||
for _, acct := range accounts {
|
||||
if acct.Structure == "workspace" && acct.AccountID == preferredID {
|
||||
return acct
|
||||
}
|
||||
}
|
||||
}
|
||||
for _, acct := range accounts {
|
||||
if acct.Structure == "workspace" && acct.PlanType == "team" {
|
||||
return acct
|
||||
}
|
||||
}
|
||||
for _, acct := range accounts {
|
||||
if acct.Structure == "workspace" {
|
||||
return acct
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *TaskRunner) verifyTaskMembership(
|
||||
ctx context.Context,
|
||||
taskType string,
|
||||
session *chatgpt.Session,
|
||||
teamAccountID string,
|
||||
statusFn chatgpt.StatusFunc,
|
||||
) (*finalMembershipState, error) {
|
||||
var lastState *finalMembershipState
|
||||
var lastErr error
|
||||
|
||||
for attempt := 1; attempt <= finalMembershipPolls; attempt++ {
|
||||
if err := refreshVerificationTokens(ctx, session); err != nil && statusFn != nil {
|
||||
statusFn(" -> Final token refresh %d/%d failed: %v", attempt, finalMembershipPolls, err)
|
||||
}
|
||||
|
||||
accounts, err := chatgpt.CheckAccountFull(session.Client, session.AccessToken, session.DeviceID)
|
||||
if err != nil {
|
||||
lastErr = fmt.Errorf("accounts/check failed: %w", err)
|
||||
} else {
|
||||
lastState = &finalMembershipState{
|
||||
Personal: selectPersonalMembership(accounts),
|
||||
Workspace: selectWorkspaceMembership(accounts, teamAccountID),
|
||||
}
|
||||
}
|
||||
|
||||
if (taskType == TaskTypeTeam || taskType == TaskTypeBoth) &&
|
||||
(lastState == nil || !lastState.teamActive()) &&
|
||||
teamAccountID != "" {
|
||||
workspaceToken, wsErr := chatgpt.GetWorkspaceAccessToken(session.Client, teamAccountID)
|
||||
if wsErr == nil {
|
||||
if wsAccounts, wsCheckErr := chatgpt.CheckAccountFull(session.Client, workspaceToken, session.DeviceID); wsCheckErr == nil {
|
||||
if lastState == nil {
|
||||
lastState = &finalMembershipState{}
|
||||
}
|
||||
lastState.Workspace = selectWorkspaceMembership(wsAccounts, teamAccountID)
|
||||
} else if lastErr == nil {
|
||||
lastErr = fmt.Errorf("workspace accounts/check failed: %w", wsCheckErr)
|
||||
}
|
||||
} else if lastErr == nil {
|
||||
lastErr = fmt.Errorf("workspace token refresh failed: %w", wsErr)
|
||||
}
|
||||
}
|
||||
|
||||
if lastState != nil && statusFn != nil {
|
||||
statusFn(" -> Final membership %d/%d: %s", attempt, finalMembershipPolls, lastState.describe())
|
||||
}
|
||||
if lastState != nil && lastState.satisfied(taskType) {
|
||||
return lastState, nil
|
||||
}
|
||||
if lastState != nil {
|
||||
lastErr = fmt.Errorf("membership mismatch for %s: %s", taskType, lastState.describe())
|
||||
}
|
||||
|
||||
if attempt < finalMembershipPolls {
|
||||
if err := sleepWithContext(ctx, finalMembershipPollDelay); err != nil {
|
||||
return lastState, err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if lastState != nil {
|
||||
return lastState, fmt.Errorf("final membership mismatch for %s: %s", taskType, lastState.describe())
|
||||
}
|
||||
if lastErr != nil {
|
||||
return nil, fmt.Errorf("final membership check failed for %s: %w", taskType, lastErr)
|
||||
}
|
||||
return nil, fmt.Errorf("final membership check failed for %s", taskType)
|
||||
}
|
||||
|
||||
func refreshVerificationTokens(ctx context.Context, session *chatgpt.Session) error {
|
||||
if err := session.RefreshSession(); err == nil {
|
||||
return nil
|
||||
} else {
|
||||
tokens, tokenErr := auth.ObtainCodexTokens(ctx, session.Client, session.DeviceID, "")
|
||||
if tokenErr != nil {
|
||||
return fmt.Errorf("session refresh failed: %v; codex refresh failed: %w", err, tokenErr)
|
||||
}
|
||||
if tokens.AccessToken == "" {
|
||||
return fmt.Errorf("codex refresh returned empty access token")
|
||||
}
|
||||
session.AccessToken = tokens.AccessToken
|
||||
if tokens.RefreshToken != "" {
|
||||
session.RefreshToken = tokens.RefreshToken
|
||||
}
|
||||
if tokens.IDToken != "" {
|
||||
session.IDToken = tokens.IDToken
|
||||
}
|
||||
if tokens.ChatGPTAccountID != "" {
|
||||
session.AccountID = tokens.ChatGPTAccountID
|
||||
}
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func sleepWithContext(ctx context.Context, delay time.Duration) error {
|
||||
timer := time.NewTimer(delay)
|
||||
defer timer.Stop()
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case <-timer.C:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
125
internal/task/membership_test.go
Normal file
125
internal/task/membership_test.go
Normal file
@@ -0,0 +1,125 @@
|
||||
package task
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"gpt-plus/internal/db"
|
||||
"gpt-plus/pkg/chatgpt"
|
||||
)
|
||||
|
||||
func TestFinalMembershipStateSatisfiedByTaskType(t *testing.T) {
|
||||
state := &finalMembershipState{
|
||||
Personal: &chatgpt.AccountInfo{
|
||||
AccountID: "personal-1",
|
||||
Structure: "personal",
|
||||
PlanType: "plus",
|
||||
HasActiveSubscription: true,
|
||||
SubscriptionID: "sub_123",
|
||||
},
|
||||
Workspace: &chatgpt.AccountInfo{
|
||||
AccountID: "workspace-1",
|
||||
Structure: "workspace",
|
||||
PlanType: "team",
|
||||
},
|
||||
}
|
||||
|
||||
if !state.satisfied(TaskTypePlus) {
|
||||
t.Fatal("expected plus task to be satisfied")
|
||||
}
|
||||
if !state.satisfied(TaskTypeTeam) {
|
||||
t.Fatal("expected team task to be satisfied")
|
||||
}
|
||||
if !state.satisfied(TaskTypeBoth) {
|
||||
t.Fatal("expected both task to be satisfied")
|
||||
}
|
||||
|
||||
if got := state.resultPlanForTask(TaskTypePlus); got != TaskTypePlus {
|
||||
t.Fatalf("plus resultPlanForTask = %q, want %q", got, TaskTypePlus)
|
||||
}
|
||||
if got := state.resultPlanForTask(TaskTypeTeam); got != TaskTypeTeam {
|
||||
t.Fatalf("team resultPlanForTask = %q, want %q", got, TaskTypeTeam)
|
||||
}
|
||||
if got := state.resultPlanForTask(TaskTypeBoth); got != TaskTypeBoth {
|
||||
t.Fatalf("both resultPlanForTask = %q, want %q", got, TaskTypeBoth)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFinalMembershipStateFallbackToActualPlan(t *testing.T) {
|
||||
state := &finalMembershipState{
|
||||
Personal: &chatgpt.AccountInfo{
|
||||
AccountID: "personal-1",
|
||||
Structure: "personal",
|
||||
PlanType: "free",
|
||||
},
|
||||
}
|
||||
|
||||
if state.satisfied(TaskTypePlus) {
|
||||
t.Fatal("free personal account should not satisfy plus task")
|
||||
}
|
||||
if got := state.resultPlanForTask(TaskTypePlus); got != "free" {
|
||||
t.Fatalf("resultPlanForTask = %q, want free", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAccountStatusForResult(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
result *chatgpt.AccountResult
|
||||
want string
|
||||
}{
|
||||
{name: "nil", result: nil, want: "active"},
|
||||
{name: "free", result: &chatgpt.AccountResult{PlanType: "free"}, want: "free"},
|
||||
{name: "plus", result: &chatgpt.AccountResult{PlanType: "plus"}, want: "plus"},
|
||||
{name: "team", result: &chatgpt.AccountResult{PlanType: "team"}, want: "team"},
|
||||
{name: "unknown", result: &chatgpt.AccountResult{PlanType: "mystery"}, want: "active"},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
if got := accountStatusForResult(tc.result); got != tc.want {
|
||||
t.Fatalf("accountStatusForResult() = %q, want %q", got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSaveAccountToDBUsesRealMembershipStatus(t *testing.T) {
|
||||
d := setupTaskTestDB(t)
|
||||
runner := &TaskRunner{taskID: "task-1", gormDB: d}
|
||||
|
||||
runner.saveAccountToDB(&chatgpt.AccountResult{
|
||||
Email: "free@test.com",
|
||||
Password: "pw",
|
||||
PlanType: "free",
|
||||
}, "plus", nil, "task-1")
|
||||
|
||||
var acct db.Account
|
||||
if err := d.First(&acct, "email = ?", "free@test.com").Error; err != nil {
|
||||
t.Fatalf("load account: %v", err)
|
||||
}
|
||||
if acct.Plan != "plus" {
|
||||
t.Fatalf("plan = %q, want plus", acct.Plan)
|
||||
}
|
||||
if acct.Status != "free" {
|
||||
t.Fatalf("status = %q, want free", acct.Status)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLogResultStoresUsefulMessage(t *testing.T) {
|
||||
d := setupTaskTestDB(t)
|
||||
taskRow := &db.Task{ID: "task-2", Type: "plus", Status: StatusRunning, TotalCount: 1}
|
||||
if err := d.Create(taskRow).Error; err != nil {
|
||||
t.Fatalf("create task: %v", err)
|
||||
}
|
||||
|
||||
runner := &TaskRunner{taskID: "task-2", gormDB: d}
|
||||
runner.logResult(taskRow, 1, "a@test.com", LogStatusFailed, "free", "final membership mismatch", 3)
|
||||
|
||||
var logRow db.TaskLog
|
||||
if err := d.First(&logRow, "task_id = ?", "task-2").Error; err != nil {
|
||||
t.Fatalf("load task log: %v", err)
|
||||
}
|
||||
if logRow.Message != "final membership mismatch" {
|
||||
t.Fatalf("message = %q, want final membership mismatch", logRow.Message)
|
||||
}
|
||||
}
|
||||
793
internal/task/runner.go
Normal file
793
internal/task/runner.go
Normal file
@@ -0,0 +1,793 @@
|
||||
package task
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log"
|
||||
"math/big"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"gpt-plus/config"
|
||||
"gpt-plus/internal/db"
|
||||
"gpt-plus/pkg/auth"
|
||||
"gpt-plus/pkg/captcha"
|
||||
"gpt-plus/pkg/chatgpt"
|
||||
"gpt-plus/pkg/httpclient"
|
||||
"gpt-plus/pkg/provider/card"
|
||||
"gpt-plus/pkg/provider/email"
|
||||
"gpt-plus/pkg/proxy"
|
||||
"gpt-plus/pkg/storage"
|
||||
"gpt-plus/pkg/stripe"
|
||||
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type TaskRunner struct {
|
||||
taskID string
|
||||
gormDB *gorm.DB
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
stopping atomic.Bool
|
||||
running atomic.Bool
|
||||
}
|
||||
|
||||
func NewTaskRunner(taskID string, d *gorm.DB) (*TaskRunner, error) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
return &TaskRunner{
|
||||
taskID: taskID,
|
||||
gormDB: d,
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (r *TaskRunner) IsRunning() bool {
|
||||
return r.running.Load()
|
||||
}
|
||||
|
||||
func (r *TaskRunner) GracefulStop() {
|
||||
r.stopping.Store(true)
|
||||
r.gormDB.Model(&db.Task{}).Where("id = ?", r.taskID).Update("status", StatusStopping)
|
||||
log.Printf("[task-runner] %s: graceful stop requested", r.taskID)
|
||||
}
|
||||
|
||||
func (r *TaskRunner) ForceStop() {
|
||||
r.stopping.Store(true)
|
||||
r.cancel()
|
||||
log.Printf("[task-runner] %s: force stop (context cancelled)", r.taskID)
|
||||
}
|
||||
|
||||
func (r *TaskRunner) Run() {
|
||||
r.running.Store(true)
|
||||
defer r.running.Store(false)
|
||||
|
||||
var t db.Task
|
||||
if err := r.gormDB.First(&t, "id = ?", r.taskID).Error; err != nil {
|
||||
log.Printf("[task-runner] %s: load task failed: %v", r.taskID, err)
|
||||
return
|
||||
}
|
||||
|
||||
// Load config from DB
|
||||
cfg, err := config.LoadFromDB(r.gormDB)
|
||||
if err != nil {
|
||||
log.Printf("[task-runner] %s: load config failed: %v", r.taskID, err)
|
||||
r.failTask(&t, "配置加载失败: "+err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
// Override task-specific settings based on task type
|
||||
switch t.Type {
|
||||
case TaskTypePlus:
|
||||
cfg.Team.Enabled = false
|
||||
case TaskTypeTeam:
|
||||
cfg.Team.Enabled = true
|
||||
case TaskTypeBoth:
|
||||
cfg.Team.Enabled = true
|
||||
}
|
||||
|
||||
// Mark task as running
|
||||
now := time.Now()
|
||||
t.Status = StatusRunning
|
||||
t.StartedAt = &now
|
||||
r.gormDB.Save(&t)
|
||||
|
||||
// Wire locale
|
||||
if cfg.Account.Locale != "" {
|
||||
lang := cfg.Account.Locale
|
||||
acceptLang := lang + ",en;q=0.9"
|
||||
if lang == "en-US" {
|
||||
acceptLang = "en-US,en;q=0.9"
|
||||
}
|
||||
auth.SetLocale(lang, acceptLang)
|
||||
chatgpt.SetDefaultLanguage(lang)
|
||||
}
|
||||
|
||||
// Initialize components
|
||||
proxyHasSession := strings.Contains(cfg.Proxy.URL, "{SESSION}")
|
||||
b2proxyEnabled := cfg.Proxy.B2Proxy.Enabled
|
||||
var cardCountry string
|
||||
if b2proxyEnabled {
|
||||
cardCountry = cfg.Card.ExpectedCountry()
|
||||
}
|
||||
|
||||
var client *httpclient.Client
|
||||
if !b2proxyEnabled && !proxyHasSession {
|
||||
client, err = httpclient.NewClient(cfg.Proxy.URL)
|
||||
if err != nil {
|
||||
r.failTask(&t, "创建 HTTP 客户端失败: "+err.Error())
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// Email provider
|
||||
emailProv := email.NewMailGateway(
|
||||
cfg.Email.MailGateway.BaseURL,
|
||||
cfg.Email.MailGateway.APIKey,
|
||||
cfg.Email.MailGateway.Provider,
|
||||
)
|
||||
|
||||
// Card provider — always use DB mode
|
||||
cardProv := card.NewDBCardProvider(card.DBCardProviderConfig{
|
||||
DB: r.gormDB,
|
||||
DefaultName: cfg.Card.API.DefaultName,
|
||||
DefaultCountry: cfg.Card.API.DefaultCountry,
|
||||
DefaultCurrency: cfg.Card.API.DefaultCurrency,
|
||||
DefaultAddress: cfg.Card.API.DefaultAddress,
|
||||
DefaultCity: cfg.Card.API.DefaultCity,
|
||||
DefaultState: cfg.Card.API.DefaultState,
|
||||
DefaultPostalCode: cfg.Card.API.DefaultPostalCode,
|
||||
APIBaseURL: cfg.Card.API.BaseURL,
|
||||
})
|
||||
|
||||
// Captcha solver
|
||||
var solver *captcha.Solver
|
||||
if cfg.Captcha.APIKey != "" {
|
||||
captchaProxy := cfg.Captcha.Proxy
|
||||
if captchaProxy == "" {
|
||||
captchaProxy = cfg.Proxy.URL
|
||||
}
|
||||
solver = captcha.NewSolver(cfg.Captcha.APIKey, captchaProxy)
|
||||
}
|
||||
|
||||
// Browser fingerprint pool
|
||||
var fpPool *stripe.BrowserFingerprintPool
|
||||
if cfg.Stripe.FingerprintDir != "" {
|
||||
fpPool, _ = stripe.NewBrowserFingerprintPool(cfg.Stripe.FingerprintDir, "en")
|
||||
}
|
||||
|
||||
// Stripe constants
|
||||
stripe.SetFallbackConstants(cfg.Stripe.BuildHash, cfg.Stripe.TagVersion)
|
||||
stripe.FetchStripeConstants()
|
||||
|
||||
log.Printf("[task-runner] %s: starting batch of %d", r.taskID, t.TotalCount)
|
||||
|
||||
// Main loop
|
||||
for i := 1; i <= t.TotalCount; i++ {
|
||||
if r.stopping.Load() {
|
||||
break
|
||||
}
|
||||
if r.ctx.Err() != nil {
|
||||
break
|
||||
}
|
||||
|
||||
log.Printf("[task-runner] %s: iteration %d/%d", r.taskID, i, t.TotalCount)
|
||||
|
||||
// Per-iteration client (B2Proxy / session template)
|
||||
iterClient := client
|
||||
if b2proxyEnabled {
|
||||
r.logStep(i, "", "正在连接 B2Proxy...")
|
||||
iterClient = r.getB2ProxyClient(cfg, cardCountry, i, solver)
|
||||
if iterClient == nil {
|
||||
r.logResult(&t, i, "", LogStatusFailed, "", "B2Proxy 连接失败", 0)
|
||||
continue
|
||||
}
|
||||
r.logStep(i, "", "B2Proxy 连接成功")
|
||||
} else if proxyHasSession {
|
||||
r.logStep(i, "", "正在连接代理...")
|
||||
iterClient = r.getSessionClient(cfg, i, solver)
|
||||
if iterClient == nil {
|
||||
r.logResult(&t, i, "", LogStatusFailed, "", "代理连接失败", 0)
|
||||
continue
|
||||
}
|
||||
r.logStep(i, "", "代理连接成功")
|
||||
}
|
||||
|
||||
start := time.Now()
|
||||
result := r.runOnce(r.ctx, i, t.Type, iterClient, emailProv, cardProv, solver, cfg, fpPool)
|
||||
dur := int(time.Since(start).Seconds())
|
||||
|
||||
status := LogStatusSuccess
|
||||
errMsg := ""
|
||||
if result.Error != nil {
|
||||
status = LogStatusFailed
|
||||
errMsg = result.Error.Error()
|
||||
}
|
||||
|
||||
r.logResult(&t, i, result.Email, status, result.Plan, errMsg, dur)
|
||||
|
||||
if i < t.TotalCount {
|
||||
time.Sleep(3 * time.Second)
|
||||
}
|
||||
}
|
||||
|
||||
// Finalize task status
|
||||
stopped := time.Now()
|
||||
t.StoppedAt = &stopped
|
||||
if r.stopping.Load() || r.ctx.Err() != nil {
|
||||
t.Status = StatusStopped
|
||||
} else {
|
||||
t.Status = StatusCompleted
|
||||
}
|
||||
r.gormDB.Save(&t)
|
||||
log.Printf("[task-runner] %s: finished (status=%s, success=%d, fail=%d)", r.taskID, t.Status, t.SuccessCount, t.FailCount)
|
||||
}
|
||||
|
||||
type iterResult struct {
|
||||
Email string
|
||||
Plan string
|
||||
Error error
|
||||
}
|
||||
|
||||
func (r *TaskRunner) runOnce(ctx context.Context, index int, taskType string, client *httpclient.Client,
|
||||
emailProv email.EmailProvider, cardProv card.CardProvider, solver *captcha.Solver,
|
||||
cfg *config.Config, fpPool *stripe.BrowserFingerprintPool) iterResult {
|
||||
|
||||
result := iterResult{Plan: "failed"}
|
||||
|
||||
password := generatePassword(cfg.Account.PasswordLength)
|
||||
|
||||
// Step: Register
|
||||
r.logStep(index, "", "正在注册账号...")
|
||||
regResult, err := auth.Register(ctx, client, emailProv, password)
|
||||
if err != nil {
|
||||
result.Error = fmt.Errorf("注册失败: %w", err)
|
||||
return result
|
||||
}
|
||||
result.Email = regResult.Email
|
||||
r.logStep(index, regResult.Email, "注册成功: "+regResult.Email)
|
||||
|
||||
r.saveEmailRecord(regResult.Email, "in_use", "owner")
|
||||
|
||||
// Step: Login
|
||||
sentinel := auth.NewSentinelGeneratorWithDeviceID(client, regResult.DeviceID)
|
||||
var loginResult *auth.LoginResult
|
||||
|
||||
if regResult.Tokens != nil {
|
||||
loginResult = regResult.Tokens
|
||||
r.logStep(index, regResult.Email, "登录成功 (注册时已获取令牌)")
|
||||
} else {
|
||||
r.logStep(index, regResult.Email, "正在登录...")
|
||||
loginResult, err = auth.Login(ctx, client, regResult.Email, password,
|
||||
regResult.DeviceID, sentinel, regResult.MailboxID, emailProv)
|
||||
if err != nil {
|
||||
result.Error = fmt.Errorf("登录失败: %w", err)
|
||||
r.updateEmailRecord(regResult.Email, "used_failed")
|
||||
return result
|
||||
}
|
||||
r.logStep(index, regResult.Email, "登录成功")
|
||||
}
|
||||
|
||||
session := chatgpt.NewSession(client, loginResult.AccessToken, loginResult.RefreshToken,
|
||||
loginResult.IDToken, regResult.DeviceID, loginResult.ChatGPTAccountID, loginResult.ChatGPTUserID)
|
||||
|
||||
// Step: Activate Plus
|
||||
r.logStep(index, regResult.Email, "正在开通 Plus...")
|
||||
var browserFP *stripe.BrowserFingerprint
|
||||
if fpPool != nil {
|
||||
browserFP = fpPool.Get()
|
||||
}
|
||||
statusFn := func(f string, a ...interface{}) {
|
||||
msg := fmt.Sprintf(f, a...)
|
||||
log.Printf("[task-runner] %s #%d: %s", r.taskID, index, msg)
|
||||
r.logStep(index, regResult.Email, msg)
|
||||
}
|
||||
plusResult, plusErr := chatgpt.ActivatePlus(ctx, session, sentinel, cardProv, cfg.Stripe, regResult.Email, solver, statusFn, browserFP)
|
||||
|
||||
plusSkipped := false
|
||||
if plusErr != nil {
|
||||
plusSkipped = true
|
||||
if errors.Is(plusErr, chatgpt.ErrPlusNotEligible) || errors.Is(plusErr, chatgpt.ErrCaptchaRequired) ||
|
||||
errors.Is(plusErr, stripe.ErrCardDeclined) || strings.Contains(plusErr.Error(), "challenge") {
|
||||
r.logStep(index, regResult.Email, "Plus 跳过: "+plusErr.Error())
|
||||
log.Printf("[task-runner] %s #%d: Plus skipped: %v", r.taskID, index, plusErr)
|
||||
} else {
|
||||
r.logStep(index, regResult.Email, "Plus 失败: "+plusErr.Error())
|
||||
log.Printf("[task-runner] %s #%d: Plus failed: %v", r.taskID, index, plusErr)
|
||||
}
|
||||
}
|
||||
|
||||
buildAccountResult := func(fp *stripe.Fingerprint) *chatgpt.AccountResult {
|
||||
return &chatgpt.AccountResult{
|
||||
Email: regResult.Email, Password: password,
|
||||
AccessToken: session.AccessToken, RefreshToken: session.RefreshToken,
|
||||
IDToken: session.IDToken, ChatGPTAccountID: session.AccountID,
|
||||
ChatGPTUserID: session.UserID, GUID: fp.GUID, MUID: fp.MUID, SID: fp.SID,
|
||||
Proxy: cfg.Proxy.URL, CreatedAt: time.Now().UTC().Format(time.RFC3339),
|
||||
}
|
||||
}
|
||||
|
||||
var fingerprint *stripe.Fingerprint
|
||||
if !plusSkipped {
|
||||
r.logStep(index, regResult.Email, "Plus 开通成功: "+plusResult.PlanType)
|
||||
result.Plan = plusResult.PlanType
|
||||
fingerprint = &stripe.Fingerprint{GUID: plusResult.GUID, MUID: plusResult.MUID, SID: plusResult.SID}
|
||||
|
||||
session.RefreshSession()
|
||||
if session.AccountID == "" {
|
||||
if acctInfo, acctErr := chatgpt.CheckAccount(session.Client, session.AccessToken, session.DeviceID); acctErr == nil {
|
||||
session.AccountID = acctInfo.AccountID
|
||||
}
|
||||
}
|
||||
|
||||
r.logStep(index, regResult.Email, "正在获取 Plus OAuth 授权...")
|
||||
plusCodexTokens, plusCodexErr := auth.ObtainCodexTokens(ctx, client, regResult.DeviceID, "")
|
||||
if plusCodexErr == nil {
|
||||
session.AccessToken = plusCodexTokens.AccessToken
|
||||
session.RefreshToken = plusCodexTokens.RefreshToken
|
||||
if plusCodexTokens.IDToken != "" {
|
||||
session.IDToken = plusCodexTokens.IDToken
|
||||
}
|
||||
if plusCodexTokens.ChatGPTAccountID != "" {
|
||||
session.AccountID = plusCodexTokens.ChatGPTAccountID
|
||||
}
|
||||
r.logStep(index, regResult.Email, "Plus OAuth 授权成功")
|
||||
} else {
|
||||
r.logStep(index, regResult.Email, "Plus OAuth 授权失败 (非致命)")
|
||||
}
|
||||
|
||||
plusAccount := buildAccountResult(fingerprint)
|
||||
plusAccount.PlanType = plusResult.PlanType
|
||||
plusAccount.StripeSessionID = plusResult.StripeSessionID
|
||||
storage.SavePlusAuthFile(cfg.Output.Dir, plusAccount)
|
||||
|
||||
r.saveAccountToDB(plusAccount, "plus", nil, r.taskID)
|
||||
r.updateEmailRecord(regResult.Email, "used")
|
||||
r.logStep(index, regResult.Email, "Plus 账号已保存")
|
||||
} else {
|
||||
sc := stripe.FetchStripeConstants()
|
||||
if browserFP != nil {
|
||||
fingerprint, _ = stripe.GetFingerprint(session.Client, chatgpt.GetDefaultUA(), "chatgpt.com", sc.TagVersion, browserFP)
|
||||
} else {
|
||||
fingerprint, _ = stripe.GetFingerprintAuto(ctx, session.Client, chatgpt.GetDefaultUA(), "chatgpt.com", sc.TagVersion)
|
||||
}
|
||||
if fingerprint == nil {
|
||||
fingerprint = &stripe.Fingerprint{}
|
||||
}
|
||||
}
|
||||
|
||||
// Team activation
|
||||
var teamResult *chatgpt.TeamResult
|
||||
var teamErr error
|
||||
if cfg.Team.Enabled {
|
||||
r.logStep(index, regResult.Email, "正在开通 Team...")
|
||||
teamStatusFn := func(f string, a ...interface{}) {
|
||||
msg := fmt.Sprintf(f, a...)
|
||||
log.Printf("[task-runner] %s #%d: %s", r.taskID, index, msg)
|
||||
r.logStep(index, regResult.Email, msg)
|
||||
}
|
||||
teamResult, teamErr = chatgpt.ActivateTeam(ctx, session, sentinel, cardProv,
|
||||
cfg.Stripe, cfg.Team, fingerprint, regResult.Email, solver,
|
||||
emailProv, regResult.MailboxID, teamStatusFn)
|
||||
if teamErr != nil {
|
||||
r.logStep(index, regResult.Email, "Team 失败: "+teamErr.Error())
|
||||
log.Printf("[task-runner] %s #%d: Team failed: %v", r.taskID, index, teamErr)
|
||||
}
|
||||
if teamResult != nil {
|
||||
r.logStep(index, regResult.Email, "Team 开通成功: "+teamResult.TeamAccountID)
|
||||
if result.Plan == "failed" {
|
||||
result.Plan = "team"
|
||||
}
|
||||
savedAccessToken := session.AccessToken
|
||||
r.logStep(index, regResult.Email, "正在获取 Team OAuth 授权...")
|
||||
codexTokens, codexErr := auth.ObtainCodexTokens(ctx, client, regResult.DeviceID, teamResult.TeamAccountID)
|
||||
|
||||
var wsAccessToken string
|
||||
if codexErr != nil {
|
||||
wsAccessToken, _ = chatgpt.GetWorkspaceAccessToken(client, teamResult.TeamAccountID)
|
||||
if wsAccessToken == "" {
|
||||
wsAccessToken = savedAccessToken
|
||||
}
|
||||
} else {
|
||||
wsAccessToken = codexTokens.AccessToken
|
||||
}
|
||||
|
||||
teamAccount := buildAccountResult(fingerprint)
|
||||
teamAccount.AccessToken = wsAccessToken
|
||||
teamAccount.ChatGPTAccountID = teamResult.TeamAccountID
|
||||
teamAccount.TeamAccountID = teamResult.TeamAccountID
|
||||
teamAccount.WorkspaceToken = wsAccessToken
|
||||
if codexTokens != nil {
|
||||
teamAccount.RefreshToken = codexTokens.RefreshToken
|
||||
if codexTokens.IDToken != "" {
|
||||
teamAccount.IDToken = codexTokens.IDToken
|
||||
}
|
||||
}
|
||||
teamAccount.PlanType = "team"
|
||||
storage.SaveTeamAuthFile(cfg.Output.Dir, teamAccount)
|
||||
|
||||
parentID := r.saveAccountToDB(teamAccount, "team_owner", nil, r.taskID)
|
||||
r.updateEmailRecord(regResult.Email, "used")
|
||||
r.logStep(index, regResult.Email, "Team 账号已保存")
|
||||
|
||||
if cfg.Team.InviteCount > 0 && teamResult.TeamAccountID != "" {
|
||||
r.logStep(index, regResult.Email, fmt.Sprintf("正在邀请 %d 个成员...", cfg.Team.InviteCount))
|
||||
r.inviteMembers(ctx, cfg, client, emailProv, session, regResult, teamResult, fingerprint, parentID)
|
||||
r.logStep(index, regResult.Email, "成员邀请完成")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if plusResult == nil && teamResult == nil && taskType == TaskTypePlus {
|
||||
result.Plan = "free"
|
||||
freeAccount := buildAccountResult(fingerprint)
|
||||
freeAccount.PlanType = "free"
|
||||
storage.SaveFreeAuthFile(cfg.Output.Dir, freeAccount)
|
||||
r.saveAccountToDB(freeAccount, "plus", nil, r.taskID)
|
||||
}
|
||||
|
||||
membershipState, verifyErr := r.verifyTaskMembership(ctx, taskType, session, teamAccountIDFromResult(teamResult), statusFn)
|
||||
if membershipState != nil {
|
||||
result.Plan = membershipState.resultPlanForTask(taskType)
|
||||
}
|
||||
if verifyErr != nil {
|
||||
reasons := []string{verifyErr.Error()}
|
||||
if plusErr != nil {
|
||||
reasons = append(reasons, "plus="+plusErr.Error())
|
||||
}
|
||||
if teamErr != nil {
|
||||
reasons = append(reasons, "team="+teamErr.Error())
|
||||
}
|
||||
result.Error = fmt.Errorf("最终会员校验失败: %s", strings.Join(reasons, "; "))
|
||||
r.updateEmailRecord(regResult.Email, "used_failed")
|
||||
return result
|
||||
}
|
||||
|
||||
r.updateEmailRecord(regResult.Email, "used")
|
||||
return result
|
||||
}
|
||||
|
||||
func (r *TaskRunner) inviteMembers(ctx context.Context, cfg *config.Config,
|
||||
client *httpclient.Client, emailProv email.EmailProvider,
|
||||
session *chatgpt.Session, regResult *auth.RegisterResult,
|
||||
teamResult *chatgpt.TeamResult, fingerprint *stripe.Fingerprint, parentID *uint) {
|
||||
|
||||
invWsToken, invWsErr := chatgpt.GetWorkspaceAccessToken(client, teamResult.TeamAccountID)
|
||||
if invWsErr != nil {
|
||||
invWsToken = session.AccessToken
|
||||
}
|
||||
if invWsToken == "" {
|
||||
return
|
||||
}
|
||||
|
||||
for invIdx := 1; invIdx <= cfg.Team.InviteCount; invIdx++ {
|
||||
if r.stopping.Load() || r.ctx.Err() != nil {
|
||||
break
|
||||
}
|
||||
|
||||
r.logStep(0, regResult.Email, fmt.Sprintf("成员 %d/%d: 创建代理客户端...", invIdx, cfg.Team.InviteCount))
|
||||
|
||||
// Use B2Proxy for member too, not just cfg.Proxy.URL (which may be empty in B2Proxy mode)
|
||||
var memberClient *httpclient.Client
|
||||
var memberClientErr error
|
||||
if cfg.Proxy.B2Proxy.Enabled {
|
||||
cardCountry := cfg.Card.ExpectedCountry()
|
||||
for try := 1; try <= 3; try++ {
|
||||
proxyURL, fetchErr := proxy.FetchB2Proxy(cfg.Proxy.B2Proxy, cardCountry)
|
||||
if fetchErr != nil {
|
||||
continue
|
||||
}
|
||||
memberClient, memberClientErr = httpclient.NewClient(proxyURL)
|
||||
if memberClientErr == nil {
|
||||
break
|
||||
}
|
||||
}
|
||||
} else if strings.Contains(cfg.Proxy.URL, "{SESSION}") {
|
||||
memberProxyURL := strings.ReplaceAll(cfg.Proxy.URL, "{SESSION}", fmt.Sprintf("inv%s%d", randomSessionID(), invIdx))
|
||||
memberClient, memberClientErr = httpclient.NewClient(memberProxyURL)
|
||||
} else {
|
||||
memberClient, memberClientErr = httpclient.NewClient(cfg.Proxy.URL)
|
||||
}
|
||||
_ = memberClientErr
|
||||
if memberClient == nil {
|
||||
r.logStep(0, regResult.Email, fmt.Sprintf("成员 %d: 代理连接失败，跳过", invIdx))
|
||||
continue
|
||||
}
|
||||
|
||||
r.logStep(0, regResult.Email, fmt.Sprintf("成员 %d: 注册中...", invIdx))
|
||||
memberPassword := auth.GenerateRandomPassword(cfg.Account.PasswordLength)
|
||||
memberReg, err := auth.Register(ctx, memberClient, emailProv, memberPassword)
|
||||
if err != nil {
|
||||
r.logStep(0, regResult.Email, fmt.Sprintf("成员 %d: 注册失败: %v", invIdx, err))
|
||||
continue
|
||||
}
|
||||
r.logStep(0, memberReg.Email, fmt.Sprintf("成员 %d: 注册成功: %s", invIdx, memberReg.Email))
|
||||
|
||||
r.saveEmailRecord(memberReg.Email, "in_use", "member")
|
||||
|
||||
r.logStep(0, memberReg.Email, fmt.Sprintf("成员 %d: 发送邀请...", invIdx))
|
||||
if err := chatgpt.InviteToTeam(client, invWsToken, teamResult.TeamAccountID, session.DeviceID, memberReg.Email); err != nil {
|
||||
r.logStep(0, memberReg.Email, fmt.Sprintf("成员 %d: 邀请失败: %v", invIdx, err))
|
||||
r.updateEmailRecord(memberReg.Email, "used_failed")
|
||||
continue
|
||||
}
|
||||
r.logStep(0, memberReg.Email, fmt.Sprintf("成员 %d: 邀请成功，等待传播...", invIdx))
|
||||
|
||||
time.Sleep(3 * time.Second)
|
||||
|
||||
var memberLoginResult *auth.LoginResult
|
||||
if memberReg.Tokens != nil {
|
||||
memberLoginResult = memberReg.Tokens
|
||||
r.logStep(0, memberReg.Email, fmt.Sprintf("成员 %d: 使用注册令牌", invIdx))
|
||||
} else {
|
||||
r.logStep(0, memberReg.Email, fmt.Sprintf("成员 %d: 登录中...", invIdx))
|
||||
memberSentinel := auth.NewSentinelGeneratorWithDeviceID(memberClient, memberReg.DeviceID)
|
||||
memberLoginResult, err = auth.Login(ctx, memberClient, memberReg.Email, memberPassword,
|
||||
memberReg.DeviceID, memberSentinel, memberReg.MailboxID, emailProv)
|
||||
if err != nil {
|
||||
r.logStep(0, memberReg.Email, fmt.Sprintf("成员 %d: 登录失败: %v", invIdx, err))
|
||||
r.updateEmailRecord(memberReg.Email, "used_failed")
|
||||
continue
|
||||
}
|
||||
r.logStep(0, memberReg.Email, fmt.Sprintf("成员 %d: 登录成功", invIdx))
|
||||
}
|
||||
|
||||
var memberWsToken, memberRefreshToken, memberIDToken, memberAccountID string
|
||||
memberWsToken = memberLoginResult.AccessToken
|
||||
memberRefreshToken = memberLoginResult.RefreshToken
|
||||
memberIDToken = memberLoginResult.IDToken
|
||||
memberAccountID = memberLoginResult.ChatGPTAccountID
|
||||
|
||||
r.logStep(0, memberReg.Email, fmt.Sprintf("成员 %d: 获取 Team OAuth 授权...", invIdx))
|
||||
for codexTry := 1; codexTry <= 3; codexTry++ {
|
||||
memberCodex, err := auth.ObtainCodexTokens(ctx, memberClient, memberReg.DeviceID, teamResult.TeamAccountID)
|
||||
if err != nil {
|
||||
if codexTry < 3 {
|
||||
time.Sleep(3 * time.Second)
|
||||
}
|
||||
continue
|
||||
}
|
||||
if memberCodex.ChatGPTAccountID == teamResult.TeamAccountID {
|
||||
memberWsToken = memberCodex.AccessToken
|
||||
memberRefreshToken = memberCodex.RefreshToken
|
||||
if memberCodex.IDToken != "" {
|
||||
memberIDToken = memberCodex.IDToken
|
||||
}
|
||||
memberAccountID = memberCodex.ChatGPTAccountID
|
||||
r.logStep(0, memberReg.Email, fmt.Sprintf("成员 %d: Team OAuth 成功", invIdx))
|
||||
break
|
||||
}
|
||||
if codexTry < 3 {
|
||||
time.Sleep(3 * time.Second)
|
||||
}
|
||||
}
|
||||
|
||||
memberAuth := &chatgpt.AccountResult{
|
||||
Email: memberReg.Email, Password: memberPassword,
|
||||
AccessToken: memberWsToken, RefreshToken: memberRefreshToken,
|
||||
IDToken: memberIDToken, ChatGPTAccountID: memberAccountID,
|
||||
ChatGPTUserID: memberLoginResult.ChatGPTUserID,
|
||||
TeamAccountID: teamResult.TeamAccountID, WorkspaceToken: memberWsToken,
|
||||
PlanType: "team", Proxy: cfg.Proxy.URL,
|
||||
CreatedAt: time.Now().UTC().Format(time.RFC3339),
|
||||
}
|
||||
storage.SaveTeamAuthFile(cfg.Output.Dir, memberAuth)
|
||||
r.saveAccountToDB(memberAuth, "team_member", parentID, r.taskID)
|
||||
r.updateEmailRecord(memberReg.Email, "used_member")
|
||||
r.logStep(0, memberReg.Email, fmt.Sprintf("成员 %d: 已保存 (account=%s)", invIdx, memberAccountID))
|
||||
}
|
||||
}
|
||||
|
||||
func (r *TaskRunner) logStep(index int, email, message string) {
|
||||
r.gormDB.Create(&db.TaskLog{
|
||||
TaskID: r.taskID,
|
||||
Index: index,
|
||||
Email: email,
|
||||
Status: "step",
|
||||
Message: message,
|
||||
})
|
||||
}
|
||||
|
||||
func (r *TaskRunner) logResult(t *db.Task, index int, email, status, plan, errMsg string, duration int) {
|
||||
message := status
|
||||
if errMsg != "" {
|
||||
message = errMsg
|
||||
} else if plan != "" {
|
||||
message = "membership verified: " + plan
|
||||
}
|
||||
logEntry := &db.TaskLog{
|
||||
TaskID: r.taskID,
|
||||
Index: index,
|
||||
Email: email,
|
||||
Status: status,
|
||||
Plan: plan,
|
||||
Error: errMsg,
|
||||
Duration: duration,
|
||||
Message: message,
|
||||
}
|
||||
r.gormDB.Create(logEntry)
|
||||
|
||||
t.DoneCount++
|
||||
if status == LogStatusSuccess {
|
||||
t.SuccessCount++
|
||||
} else {
|
||||
t.FailCount++
|
||||
}
|
||||
r.gormDB.Model(t).Updates(map[string]interface{}{
|
||||
"done_count": t.DoneCount,
|
||||
"success_count": t.SuccessCount,
|
||||
"fail_count": t.FailCount,
|
||||
})
|
||||
}
|
||||
|
||||
func (r *TaskRunner) failTask(t *db.Task, errMsg string) {
|
||||
now := time.Now()
|
||||
t.Status = StatusStopped
|
||||
t.StoppedAt = &now
|
||||
r.gormDB.Save(t)
|
||||
log.Printf("[task-runner] %s: task failed: %s", r.taskID, errMsg)
|
||||
}
|
||||
|
||||
func (r *TaskRunner) saveEmailRecord(email, status, role string) {
|
||||
rec := &db.EmailRecord{
|
||||
Email: email,
|
||||
Status: status,
|
||||
UsedForRole: role,
|
||||
TaskID: r.taskID,
|
||||
}
|
||||
r.gormDB.Create(rec)
|
||||
}
|
||||
|
||||
func (r *TaskRunner) updateEmailRecord(email, status string) {
|
||||
r.gormDB.Model(&db.EmailRecord{}).Where("email = ?", email).Update("status", status)
|
||||
}
|
||||
|
||||
func teamAccountIDFromResult(result *chatgpt.TeamResult) string {
|
||||
if result == nil {
|
||||
return ""
|
||||
}
|
||||
return result.TeamAccountID
|
||||
}
|
||||
|
||||
func accountStatusForResult(result *chatgpt.AccountResult) string {
|
||||
if result == nil {
|
||||
return "active"
|
||||
}
|
||||
switch result.PlanType {
|
||||
case "free", "plus", "team":
|
||||
return result.PlanType
|
||||
default:
|
||||
return "active"
|
||||
}
|
||||
}
|
||||
|
||||
func (r *TaskRunner) saveAccountToDB(result *chatgpt.AccountResult, plan string, parentID *uint, taskID string) *uint {
|
||||
// Check if account already exists (same email)
|
||||
var existing db.Account
|
||||
if r.gormDB.Where("email = ?", result.Email).First(&existing).Error == nil {
|
||||
// Update existing record (e.g. Plus account now also has Team)
|
||||
updates := map[string]interface{}{
|
||||
"plan": plan,
|
||||
"access_token": result.AccessToken,
|
||||
"refresh_token": result.RefreshToken,
|
||||
"id_token": result.IDToken,
|
||||
"account_id": result.ChatGPTAccountID,
|
||||
"user_id": result.ChatGPTUserID,
|
||||
"team_workspace_id": result.TeamAccountID,
|
||||
"workspace_token": result.WorkspaceToken,
|
||||
"status": accountStatusForResult(result),
|
||||
}
|
||||
if parentID != nil {
|
||||
updates["parent_id"] = *parentID
|
||||
}
|
||||
r.gormDB.Model(&existing).Updates(updates)
|
||||
return &existing.ID
|
||||
}
|
||||
|
||||
acct := &db.Account{
|
||||
TaskID: taskID,
|
||||
Email: result.Email,
|
||||
Password: result.Password,
|
||||
Plan: plan,
|
||||
ParentID: parentID,
|
||||
AccessToken: result.AccessToken,
|
||||
RefreshToken: result.RefreshToken,
|
||||
IDToken: result.IDToken,
|
||||
AccountID: result.ChatGPTAccountID,
|
||||
UserID: result.ChatGPTUserID,
|
||||
TeamWorkspaceID: result.TeamAccountID,
|
||||
WorkspaceToken: result.WorkspaceToken,
|
||||
Status: accountStatusForResult(result),
|
||||
}
|
||||
r.gormDB.Create(acct)
|
||||
return &acct.ID
|
||||
}
|
||||
|
||||
func (r *TaskRunner) getB2ProxyClient(cfg *config.Config, cardCountry string, index int, solver *captcha.Solver) *httpclient.Client {
|
||||
for try := 1; try <= 5; try++ {
|
||||
proxyURL, err := proxy.FetchB2Proxy(cfg.Proxy.B2Proxy, cardCountry)
|
||||
if err != nil {
|
||||
time.Sleep(time.Second)
|
||||
continue
|
||||
}
|
||||
c, err := httpclient.NewClient(proxyURL)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
if testStripeConnectivity(c) {
|
||||
if solver != nil {
|
||||
solver.SetProxy(proxyURL)
|
||||
}
|
||||
return c
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *TaskRunner) getSessionClient(cfg *config.Config, index int, solver *captcha.Solver) *httpclient.Client {
|
||||
for try := 1; try <= 5; try++ {
|
||||
sessionID := randomSessionID()
|
||||
proxyURL := strings.ReplaceAll(cfg.Proxy.URL, "{SESSION}", sessionID)
|
||||
c, err := httpclient.NewClient(proxyURL)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
if testStripeConnectivity(c) {
|
||||
if solver != nil {
|
||||
solver.SetProxy(proxyURL)
|
||||
}
|
||||
return c
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func testStripeConnectivity(client *httpclient.Client) bool {
|
||||
req, err := http.NewRequest("GET", "https://m.stripe.com/6", nil)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
req.Header.Set("User-Agent", "Mozilla/5.0")
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
resp.Body.Close()
|
||||
|
||||
req2, _ := http.NewRequest("GET", "https://api.stripe.com", nil)
|
||||
req2.Header.Set("User-Agent", "Mozilla/5.0")
|
||||
resp2, err := client.Do(req2)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
resp2.Body.Close()
|
||||
return true
|
||||
}
|
||||
|
||||
func generatePassword(length int) string {
|
||||
const charset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789!@#$%"
|
||||
b := make([]byte, length)
|
||||
for i := range b {
|
||||
n, err := rand.Int(rand.Reader, big.NewInt(int64(len(charset))))
|
||||
if err != nil {
|
||||
b[i] = charset[i%len(charset)]
|
||||
continue
|
||||
}
|
||||
b[i] = charset[n.Int64()]
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
func randomSessionID() string {
|
||||
const chars = "abcdefghijklmnopqrstuvwxyz0123456789"
|
||||
b := make([]byte, 16)
|
||||
for i := range b {
|
||||
n, _ := rand.Int(rand.Reader, big.NewInt(int64(len(chars))))
|
||||
b[i] = chars[n.Int64()]
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
18
internal/task/types.go
Normal file
18
internal/task/types.go
Normal file
@@ -0,0 +1,18 @@
|
||||
package task
|
||||
|
||||
const (
|
||||
StatusPending = "pending"
|
||||
StatusRunning = "running"
|
||||
StatusStopping = "stopping"
|
||||
StatusStopped = "stopped"
|
||||
StatusInterrupted = "interrupted"
|
||||
StatusCompleted = "completed"
|
||||
|
||||
LogStatusSuccess = "success"
|
||||
LogStatusFailed = "failed"
|
||||
LogStatusSkipped = "skipped"
|
||||
|
||||
TaskTypePlus = "plus"
|
||||
TaskTypeTeam = "team"
|
||||
TaskTypeBoth = "both"
|
||||
)
|
||||
Reference in New Issue
Block a user