spider/internal/handler/proxy.go

package handler

import (
	"context"
	"fmt"
	"net"
	"net/http"
	"net/url"
	"strconv"
	"sync"
	"time"

	"spider/internal/model"
	"spider/internal/store"
	"spider/internal/task"

	"github.com/gin-gonic/gin"
	"golang.org/x/net/proxy"
)

// ProxyHandler handles proxy CRUD and testing.
type ProxyHandler struct {
	store   *store.Store
	taskMgr *task.Manager
}

// List handles GET /proxies
func (h *ProxyHandler) List(c *gin.Context) {
	page, pageSize, offset := parsePage(c)

	query := h.store.DB.Model(&model.Proxy{})
	if status := c.Query("status"); status != "" {
		query = query.Where("status = ?", status)
	}
	if enabled := c.Query("enabled"); enabled != "" {
		query = query.Where("enabled = ?", enabled == "true")
	}
	if search := c.Query("search"); search != "" {
		like := "%" + search + "%"
		query = query.Where("name LIKE ? OR host LIKE ? OR region LIKE ?", like, like, like)
	}

	var total int64
	query.Count(&total)

	var proxies []model.Proxy
	if err := query.Order("id DESC").Limit(pageSize).Offset(offset).Find(&proxies).Error; err != nil {
		Fail(c, 500, err.Error())
		return
	}
	PageOK(c, proxies, total, page, pageSize)
}

// ListEnabled handles GET /proxies/enabled — returns only enabled proxies (for task dropdown)
func (h *ProxyHandler) ListEnabled(c *gin.Context) {
	var proxies []model.Proxy
	h.store.DB.Where("enabled = ?", true).Order("name ASC").Find(&proxies)
	OK(c, proxies)
}

// Create handles POST /proxies
func (h *ProxyHandler) Create(c *gin.Context) {
	var body struct {
		Name     string `json:"name" binding:"required"`
		Protocol string `json:"protocol" binding:"required"`
		Host     string `json:"host" binding:"required"`
		Port     int    `json:"port" binding:"required"`
		Username string `json:"username"`
		Password string `json:"password"`
		Region   string `json:"region"`
		Remark   string `json:"remark"`
	}
	if err := c.ShouldBindJSON(&body); err != nil {
		Fail(c, 400, err.Error())
		return
	}

	allowed := map[string]bool{"http": true, "https": true, "socks5": true}
	if !allowed[body.Protocol] {
		Fail(c, 400, "协议必须是 http/https/socks5")
		return
	}

	p := model.Proxy{
		Name:     body.Name,
		Protocol: body.Protocol,
		Host:     body.Host,
		Port:     body.Port,
		Username: body.Username,
		Password: body.Password,
		Region:   body.Region,
		Remark:   body.Remark,
		Enabled:  true,
		Status:   "unknown",
	}
	if err := h.store.DB.Create(&p).Error; err != nil {
		Fail(c, 500, err.Error())
		return
	}
	LogAudit(h.store, c, "create", "proxy", fmt.Sprintf("%d", p.ID), gin.H{"name": p.Name})
	OK(c, p)
}

// Update handles PUT /proxies/:id
func (h *ProxyHandler) Update(c *gin.Context) {
	id, err := strconv.ParseUint(c.Param("id"), 10, 64)
	if err != nil {
		Fail(c, 400, "invalid id")
		return
	}

	var p model.Proxy
	if err := h.store.DB.First(&p, id).Error; err != nil {
		Fail(c, 404, "代理不存在")
		return
	}

	var body struct {
		Name     *string `json:"name"`
		Protocol *string `json:"protocol"`
		Host     *string `json:"host"`
		Port     *int    `json:"port"`
		Username *string `json:"username"`
		Password *string `json:"password"`
		Region   *string `json:"region"`
		Remark   *string `json:"remark"`
		Enabled  *bool   `json:"enabled"`
	}
	if err := c.ShouldBindJSON(&body); err != nil {
		Fail(c, 400, err.Error())
		return
	}

	updates := map[string]any{}
	if body.Name != nil {
		updates["name"] = *body.Name
	}
	if body.Protocol != nil {
		updates["protocol"] = *body.Protocol
	}
	if body.Host != nil {
		updates["host"] = *body.Host
	}
	if body.Port != nil {
		updates["port"] = *body.Port
	}
	if body.Username != nil {
		updates["username"] = *body.Username
	}
	if body.Password != nil {
		updates["password"] = *body.Password
	}
	if body.Region != nil {
		updates["region"] = *body.Region
	}
	if body.Remark != nil {
		updates["remark"] = *body.Remark
	}
	if body.Enabled != nil {
		updates["enabled"] = *body.Enabled
	}

	h.store.DB.Model(&p).Updates(updates)
	h.store.DB.First(&p, id)
	LogAudit(h.store, c, "update", "proxy", fmt.Sprintf("%d", id), updates)
	OK(c, p)
}

// Delete handles DELETE /proxies/:id
func (h *ProxyHandler) Delete(c *gin.Context) {
	id, err := strconv.ParseUint(c.Param("id"), 10, 64)
	if err != nil {
		Fail(c, 400, "invalid id")
		return
	}
	if err := h.store.DB.Delete(&model.Proxy{}, id).Error; err != nil {
		Fail(c, 500, err.Error())
		return
	}
	LogAudit(h.store, c, "delete", "proxy", fmt.Sprintf("%d", id), nil)
	OK(c, gin.H{"message": "已删除"})
}

// Test handles POST /proxies/:id/test — tests proxy connectivity
func (h *ProxyHandler) Test(c *gin.Context) {
	id, err := strconv.ParseUint(c.Param("id"), 10, 64)
	if err != nil {
		Fail(c, 400, "invalid id")
		return
	}

	var p model.Proxy
	if err := h.store.DB.First(&p, id).Error; err != nil {
		Fail(c, 404, "代理不存在")
		return
	}

	proxyURL := p.ProxyURL()
	status := "ok"
	errMsg := ""

	if p.Protocol == "socks5" {
		// Test SOCKS5 by dialing through it
		auth := &proxy.Auth{}
		if p.Username != "" {
			auth.User = p.Username
			auth.Password = p.Password
		} else {
			auth = nil
		}
		dialer, err := proxy.SOCKS5("tcp", fmt.Sprintf("%s:%d", p.Host, p.Port), auth, proxy.Direct)
		if err != nil {
			status = "fail"
			errMsg = err.Error()
		} else {
			ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
			defer cancel()
			conn, err := dialer.(proxy.ContextDialer).DialContext(ctx, "tcp", "www.google.com:80")
			if err != nil {
				status = "fail"
				errMsg = err.Error()
			} else {
				conn.Close()
			}
		}
	} else {
		// Test HTTP/HTTPS proxy
		pURL, _ := url.Parse(proxyURL)
		client := &http.Client{
			Timeout: 10 * time.Second,
			Transport: &http.Transport{
				Proxy: http.ProxyURL(pURL),
				DialContext: (&net.Dialer{Timeout: 5 * time.Second}).DialContext,
			},
		}
		resp, err := client.Get("https://httpbin.org/ip")
		if err != nil {
			status = "fail"
			errMsg = err.Error()
		} else {
			resp.Body.Close()
			if resp.StatusCode != 200 {
				status = "fail"
				errMsg = fmt.Sprintf("HTTP %d", resp.StatusCode)
			}
		}
	}

	now := time.Now()
	h.store.DB.Model(&p).Updates(map[string]any{
		"status":          status,
		"last_checked_at": now,
	})

	result := gin.H{"status": status, "proxy_url": proxyURL}
	if errMsg != "" {
		result["error"] = errMsg
	}
	OK(c, result)
}

// TestAll handles POST /proxies/test-all — tests all enabled proxies in parallel.
func (h *ProxyHandler) TestAll(c *gin.Context) {
	var proxies []model.Proxy
	h.store.DB.Where("enabled = ?", true).Find(&proxies)
	if len(proxies) == 0 {
		Fail(c, 404, "没有已启用的代理")
		return
	}

	type testResult struct {
		ID     uint   `json:"id"`
		Name   string `json:"name"`
		Status string `json:"status"`
		Error  string `json:"error,omitempty"`
	}

	results := make([]testResult, len(proxies))
	var wg sync.WaitGroup

	for i, p := range proxies {
		wg.Add(1)
		go func(idx int, px model.Proxy) {
			defer wg.Done()
			tr := testResult{ID: px.ID, Name: px.Name}

			if px.Protocol == "socks5" {
				auth := &proxy.Auth{}
				if px.Username != "" {
					auth.User = px.Username
					auth.Password = px.Password
				} else {
					auth = nil
				}
				dialer, err := proxy.SOCKS5("tcp", fmt.Sprintf("%s:%d", px.Host, px.Port), auth, proxy.Direct)
				if err != nil {
					tr.Status = "fail"
					tr.Error = err.Error()
				} else {
					ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
					defer cancel()
					conn, err := dialer.(proxy.ContextDialer).DialContext(ctx, "tcp", "www.google.com:80")
					if err != nil {
						tr.Status = "fail"
						tr.Error = err.Error()
					} else {
						conn.Close()
						tr.Status = "ok"
					}
				}
			} else {
				pURL, _ := url.Parse(px.ProxyURL())
				client := &http.Client{
					Timeout: 10 * time.Second,
					Transport: &http.Transport{
						Proxy:       http.ProxyURL(pURL),
						DialContext: (&net.Dialer{Timeout: 5 * time.Second}).DialContext,
					},
				}
				resp, err := client.Get("https://httpbin.org/ip")
				if err != nil {
					tr.Status = "fail"
					tr.Error = err.Error()
				} else {
					resp.Body.Close()
					if resp.StatusCode != 200 {
						tr.Status = "fail"
						tr.Error = fmt.Sprintf("HTTP %d", resp.StatusCode)
					} else {
						tr.Status = "ok"
					}
				}
			}

			// Update DB
			now := time.Now()
			h.store.DB.Model(&model.Proxy{}).Where("id = ?", px.ID).Updates(map[string]any{
				"status":          tr.Status,
				"last_checked_at": now,
			})

			results[idx] = tr
		}(i, p)
	}

	wg.Wait()

	okCount := 0
	failCount := 0
	for _, r := range results {
		if r.Status == "ok" {
			okCount++
		} else {
			failCount++
		}
	}

	OK(c, gin.H{
		"total":   len(results),
		"ok":      okCount,
		"fail":    failCount,
		"results": results,
	})
}

// PoolStatus handles GET /proxies/pool-status — returns live proxy pool health.
func (h *ProxyHandler) PoolStatus(c *gin.Context) {
	if h.taskMgr == nil {
		OK(c, gin.H{"active": false, "message": "任务管理器未初始化"})
		return
	}
	pool := h.taskMgr.GetProxyPool()
	if pool == nil {
		OK(c, gin.H{"active": false, "message": "当前没有使用代理池"})
		return
	}

	entries := pool.AllEntries()
	OK(c, gin.H{
		"active":       true,
		"total":        pool.Size(),
		"active_count": pool.ActiveCount(),
		"proxies":      entries,
	})
}