proxy-platform/internal/warp/warp.go

package warp

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"

	"go.uber.org/zap"
)

var (
	ErrWARPNotInstalled   = errors.New("WARP 未安装")
	ErrWARPNotConnected    = errors.New("WARP 未连接")
	ErrIPRefreshFailed     = errors.New("IP 刷新失败")
	ErrIPPoolExhausted     = errors.New("IP 池耗尽")
)

// WARPStatus WARP 状态
type WARPStatus struct {
	Connected   bool
	AccountID   string
	DeviceID    string
	ClientID    string
	CurrentIP   string
	Country     string
	ConnectTime time.Time
}

// Client WARP 客户端
type Client struct {
	socksPort           int
	refreshCooldown     time.Duration
	maxRefreshRetries   int
	retryDelayMin       time.Duration
	retryDelayMax       time.Duration
	logger              *zap.Logger
	mu                  sync.Mutex
	lastRefresh         time.Time
	currentIP           string
	ipHistory           []string
}

// NewClient 创建 WARP 客户端
func NewClient(socksPort int, refreshCooldown int, maxRefreshRetries int, retryDelayMin int, retryDelayMax int, logger *zap.Logger) *Client {
	return &Client{
		socksPort:         socksPort,
		refreshCooldown:   time.Duration(refreshCooldown) * time.Second,
		maxRefreshRetries: maxRefreshRetries,
		retryDelayMin:     time.Duration(retryDelayMin) * time.Second,
		retryDelayMax:     time.Duration(retryDelayMax) * time.Second,
		logger:            logger,
		ipHistory:         make([]string, 0, 100),
	}
}

// Connect 连接 WARP
func (c *Client) Connect(ctx context.Context) error {
	// 检查是否安装
	if !c.isInstalled() {
		return ErrWARPNotInstalled
	}

	// 连接 WARP
	cmd := exec.CommandContext(ctx, "warp-cli", "connect")
	output, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("WARP 连接失败: %w, output: %s", err, string(output))
	}

	// 等待连接建立
	for i := 0; i < 10; i++ {
		time.Sleep(1 * time.Second)
		status, err := c.Status(ctx)
		if err == nil && status.Connected {
			c.currentIP = status.CurrentIP
			c.lastRefresh = time.Now()
			c.logger.Info("WARP 连接成功", zap.String("ip", status.CurrentIP))
			return nil
		}
	}

	return ErrWARPNotConnected
}

// Disconnect 断开 WARP
func (c *Client) Disconnect(ctx context.Context) error {
	cmd := exec.CommandContext(ctx, "warp-cli", "disconnect")
	output, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("WARP 断开失败: %w, output: %s", err, string(output))
	}
	c.currentIP = ""
	return nil
}

// Status 获取 WARP 状态
func (c *Client) Status(ctx context.Context) (*WARPStatus, error) {
	cmd := exec.CommandContext(ctx, "warp-cli", "status")
	output, err := cmd.Output()
	if err != nil {
		return nil, err
	}

	status := &WARPStatus{}
	outputStr := string(output)

	// 解析状态
	if strings.Contains(outputStr, "Status: Connected") {
		status.Connected = true
	} else if strings.Contains(outputStr, "Status: Disconnected") {
		status.Connected = false
		return status, nil
	}

	// 获取当前 IP
	ip, err := c.GetCurrentIP(ctx)
	if err == nil {
		status.CurrentIP = ip
		c.currentIP = ip
	}

	// 获取国家信息
	country, _ := c.GetCountry(ctx)
	status.Country = country

	return status, nil
}

// RefreshIP 刷新 IP
func (c *Client) RefreshIP(ctx context.Context) (string, error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	// 检查冷却时间
	if time.Since(c.lastRefresh) < c.refreshCooldown {
		waitTime := c.refreshCooldown - time.Since(c.lastRefresh)
		c.logger.Warn("IP 刷新冷却中", zap.Duration("wait_time", waitTime))
		return "", fmt.Errorf("冷却中，请等待 %v", waitTime)
	}

	// 记录旧 IP
	oldIP := c.currentIP

	// 重试刷新
	for attempt := 0; attempt < c.maxRefreshRetries; attempt++ {
		c.logger.Info("尝试刷新 IP",
			zap.Int("attempt", attempt+1),
			zap.Int("max_retries", c.maxRefreshRetries),
			zap.String("old_ip", oldIP),
		)

		// 断开连接
		if err := c.Disconnect(ctx); err != nil {
			c.logger.Error("断开 WARP 失败", zap.Error(err))
			continue
		}

		// 随机等待
		delay := c.retryDelayMin + time.Duration(randInt(0, int(c.retryDelayMax-c.retryDelayMin)))
		time.Sleep(delay)

		// 重新连接
		if err := c.Connect(ctx); err != nil {
			c.logger.Error("连接 WARP 失败", zap.Error(err))
			continue
		}

		// 获取新 IP
		newIP, err := c.GetCurrentIP(ctx)
		if err != nil {
			c.logger.Error("获取新 IP 失败", zap.Error(err))
			continue
		}

		// 检查是否真的换到了新 IP
		if newIP != oldIP {
			c.currentIP = newIP
			c.lastRefresh = time.Now()
			c.ipHistory = append(c.ipHistory, newIP)
			c.logger.Info("IP 刷新成功",
				zap.String("old_ip", oldIP),
				zap.String("new_ip", newIP),
				zap.Int("attempt", attempt+1),
			)
			return newIP, nil
		}

		c.logger.Warn("获取到相同 IP，重试中",
			zap.String("ip", newIP),
		)
	}

	return "", ErrIPPoolExhausted
}

// GetCurrentIP 获取当前 IP
func (c *Client) GetCurrentIP(ctx context.Context) (string, error) {
	// 通过 WARP 出口获取 IP
	client := &http.Client{
		Transport: &http.Transport{
			Proxy: http.ProxyURL(nil), // 使用 WARP 接口
			DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
				// 使用 WARP 的 SOCKS5 代理
				return net.DialTimeout("tcp", "127.0.0.1:"+strconv.Itoa(c.socksPort), 5*time.Second)
			},
		},
		Timeout: 10 * time.Second,
	}

	resp, err := client.Get("https://cloudflare.com/cdn-cgi/trace")
	if err != nil {
		// 回退：使用系统默认出口
		return c.getIPFromSystem()
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}

	// 解析 IP
	lines := strings.Split(string(body), "\n")
	for _, line := range lines {
		if strings.HasPrefix(line, "ip=") {
			return strings.TrimPrefix(line, "ip="), nil
		}
	}

	return "", errors.New("无法解析 IP")
}

// getIPFromSystem 从系统获取 IP
func (c *Client) getIPFromSystem() (string, error) {
	resp, err := http.Get("https://cloudflare.com/cdn-cgi/trace")
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}

	lines := strings.Split(string(body), "\n")
	for _, line := range lines {
		if strings.HasPrefix(line, "ip=") {
			return strings.TrimPrefix(line, "ip="), nil
		}
	}

	return "", errors.New("无法解析 IP")
}

// GetCountry 获取国家
func (c *Client) GetCountry(ctx context.Context) (string, error) {
	// 使用 WARP 出口
	resp, err := http.Get("https://cloudflare.com/cdn-cgi/trace")
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}

	lines := strings.Split(string(body), "\n")
	for _, line := range lines {
		if strings.HasPrefix(line, "loc=") {
			return strings.TrimPrefix(line, "loc="), nil
		}
	}

	return "??", nil
}

// isInstalled 检查是否安装
func (c *Client) isInstalled() bool {
	_, err := exec.LookPath("warp-cli")
	return err == nil
}

// GetIPHistory 获取 IP 历史
func (c *Client) GetIPHistory() []string {
	c.mu.Lock()
	defer c.mu.Unlock()

	result := make([]string, len(c.ipHistory))
	copy(result, c.ipHistory)
	return result
}

// RegisterNewAccount 注册新账号（用于获取新 IP 池）
func (c *Client) RegisterNewAccount(ctx context.Context) error {
	// 删除当前账号
	cmd := exec.CommandContext(ctx, "warp-cli", "delete")
	_, _ = cmd.CombinedOutput()

	// 注册新账号
	cmd = exec.CommandContext(ctx, "warp-cli", "register")
	output, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("注册新账号失败: %w, output: %s", err, string(output))
	}

	c.logger.Info("注册新 WARP 账号成功")
	return nil
}

// randInt 生成随机整数
func randInt(min, max int) int {
	return min + int(randFloat()*float64(max-min+1))
}

// randFloat 生成随机浮点数
func randFloat() float64 {
	return float64(time.Now().UnixNano()%1000) / 1000.0
}

// ParseWARPAccount 从输出解析账号信息
func ParseWARPAccount(output string) (accountID, deviceID, clientID string) {
	accountRegex := regexp.MustCompile(`Account ID:\s*([a-f0-9-]+)`)
	deviceRegex := regexp.MustCompile(`Device ID:\s*([a-f0-9-]+)`)
	clientRegex := regexp.MustCompile(`Client ID:\s*([a-f0-9-]+)`)

	if match := accountRegex.FindStringSubmatch(output); len(match) > 1 {
		accountID = match[1]
	}
	if match := deviceRegex.FindStringSubmatch(output); len(match) > 1 {
		deviceID = match[1]
	}
	if match := clientRegex.FindStringSubmatch(output); len(match) > 1 {
		clientID = match[1]
	}

	return
}