AntiCaptcha/internal/captcha/handler.go

package captcha

import (
	"fmt"
	"net/http"
	"os"
	"path/filepath"
	"sync"

	"anticaptcha/pkg/opencv"
)

type Handler struct {
	modelPath string
	mu        sync.RWMutex
}

// 模型配置
var modelConfigs = map[string]string{
	"math.onnx":    "math.onnx",
	"ocr.onnx":     "ocr.onnx",
	"rotate.onnx":  "rotate.onnx",
	"siamese.onnx": "siamese.onnx",
	"charsets.txt": "charsets.txt",
}

// 从 Gitea 仓库下载（无需认证，公开仓库）
const modelBaseURL = "https://git.viaeon.com/admin/anticaptcha/raw/branch/main/models"

func NewHandler(modelPath string) *Handler {
	h := &Handler{
		modelPath: modelPath,
	}
	// 确保模型目录存在并下载缺失的模型
	h.ensureModels()
	return h
}

// ensureModels 检查并下载缺失的模型
func (h *Handler) ensureModels() {
	// 确保目录存在
	if err := os.MkdirAll(h.modelPath, 0755); err != nil {
		fmt.Printf("警告: 创建模型目录失败: %v\n", err)
		return
	}

	for localName, remoteName := range modelConfigs {
		localPath := filepath.Join(h.modelPath, localName)
		if _, err := os.Stat(localPath); os.IsNotExist(err) {
			fmt.Printf("下载模型: %s -> %s\n", remoteName, localName)
			if err := h.downloadModel(remoteName, localPath); err != nil {
				fmt.Printf("警告: 下载模型 %s 失败: %v\n", remoteName, err)
			} else {
				fmt.Printf("模型下载完成: %s\n", localName)
			}
		}
	}
}

// downloadModel 下载模型文件
func (h *Handler) downloadModel(remoteName, localPath string) error {
	url := fmt.Sprintf("%s/%s", modelBaseURL, remoteName)

	resp, err := http.Get(url)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("HTTP %d", resp.StatusCode)
	}

	out, err := os.Create(localPath)
	if err != nil {
		return err
	}
	defer out.Close()

	_, err = out.ReadFrom(resp.Body)
	return err
}

// OCR 文字识别（需要 ONNX 模型）
func (h *Handler) OCR(imageBase64 string) (string, error) {
	// 暂时返回模拟结果
	// 实际实现需要加载 OCR 模型
	return "OCR result", nil
}

// Math 数学计算识别
func (h *Handler) Math(imageBase64 string) (string, error) {
	// 暂时返回模拟结果
	return "0", nil
}

// SliderMatch 滑块缺口匹配
func (h *Handler) SliderMatch(targetBase64, backgroundBase64 string) (int, error) {
	target, err := opencv.DecodeFromBase64(targetBase64)
	if err != nil {
		return 0, err
	}
	defer target.Free()

	background, err := opencv.DecodeFromBase64(backgroundBase64)
	if err != nil {
		return 0, err
	}
	defer background.Free()

	return opencv.SliderMatch(target, background)
}

// SliderComparison 阴影滑块匹配
func (h *Handler) SliderComparison(targetBase64, backgroundBase64 string) (int, error) {
	target, err := opencv.DecodeFromBase64(targetBase64)
	if err != nil {
		return 0, err
	}
	defer target.Free()

	background, err := opencv.DecodeFromBase64(backgroundBase64)
	if err != nil {
		return 0, err
	}
	defer background.Free()

	return opencv.SliderComparison(target, background)
}

// CompareSimilarity 图片相似度对比
func (h *Handler) CompareSimilarity(img1Base64, img2Base64 string) (float32, error) {
	img1, err := opencv.DecodeFromBase64(img1Base64)
	if err != nil {
		return 0, err
	}
	defer img1.Free()

	img2, err := opencv.DecodeFromBase64(img2Base64)
	if err != nil {
		return 0, err
	}
	defer img2.Free()

	return opencv.CompareSimilarity(img1, img2)
}

// SingleRotate 单图旋转验证码
func (h *Handler) SingleRotate(imageBase64 string) (float32, error) {
	img, err := opencv.DecodeFromBase64(imageBase64)
	if err != nil {
		return 0, err
	}
	defer img.Free()

	return opencv.DetectRotation(img)
}

// DoubleRotate 双图旋转验证码
func (h *Handler) DoubleRotate(insideBase64, outsideBase64 string) (float32, error) {
	// 简化处理
	inside, err := opencv.DecodeFromBase64(insideBase64)
	if err != nil {
		return 0, err
	}
	defer inside.Free()

	outside, err := opencv.DecodeFromBase64(outsideBase64)
	if err != nil {
		return 0, err
	}
	defer outside.Free()

	angleInside, err := opencv.DetectRotation(inside)
	if err != nil {
		return 0, err
	}

	angleOutside, err := opencv.DetectRotation(outside)
	if err != nil {
		return 0, err
	}

	return angleInside - angleOutside, nil
}

// DetectionIcon 图标检测
func (h *Handler) DetectionIcon(imageBase64 string) ([]map[string]int, error) {
	// 暂时返回空结果
	// 实际需要目标检测模型
	return []map[string]int{}, nil
}

// DetectionText 文字检测
func (h *Handler) DetectionText(imageBase64 string) ([]map[string]int, error) {
	// 暂时返回空结果
	return []map[string]int{}, nil
}

// DetectionIconOrder 按序检测图标
func (h *Handler) DetectionIconOrder(orderImgBase64, targetImgBase64 string) ([]map[string]int, error) {
	// 暂时返回空结果
	return []map[string]int{}, nil
}

// DetectionTextOrder 按序检测文字
func (h *Handler) DetectionTextOrder(orderImgBase64, targetImgBase64 string) ([]map[string]int, error) {
	// 暂时返回空结果
	return []map[string]int{}, nil
}