main.go

package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"sync"

	"ApiKeyScan/config"
	"ApiKeyScan/output"
	"ApiKeyScan/scanner"
	"ApiKeyScan/validator"
)

func main() {
	log.Println("程序开始执行...")

	// 命令行参数解析
	validatorType := flag.String("validator", "all", "Specify the validator to use: openai, gemini, claude, or all. Default is all.")
	flag.Parse()

	if flag.NArg() < 2 {
		fmt.Println("Usage: ApiKeyScan [-validator <type>] <search_keyword> <output_file>")
		flag.PrintDefaults()
		os.Exit(1)
	}

	searchKeyword := flag.Arg(0)
	outputFile := flag.Arg(1)

	log.Printf("解析参数: searchKeyword=%s, outputFile=%s, validator=%s\n", searchKeyword, outputFile, *validatorType)
	fmt.Printf("开始搜索关键词: %s\n", searchKeyword)
	fmt.Printf("结果将写入文件: %s\n", outputFile)

	// 初始化正则表达式
	regexes := config.GetDefaultRegexes()

	// 扫描器
	githubScanner := scanner.NewGitHubScanner()
	potentialKeys := make(chan string)
	doneScanning := make(chan struct{})

	go func() {
		defer close(doneScanning)
		err := githubScanner.Scan(searchKeyword, regexes, potentialKeys)
		if err != nil {
			log.Printf("扫描过程中发生错误: %v", err)
		}
	}()

	// 验证器
	var validators []validator.Validator
	switch *validatorType {
	case "openai":
		validators = append(validators, validator.NewOpenAIValidator("https://api.openai.com/v1/chat/completions", "gpt-3.5-turbo"))
		log.Println("使用 OpenAI 验证器")
	case "gemini":
		validators = append(validators, validator.NewGeminiValidator())
		log.Println("使用 Gemini 验证器")
	case "claude":
		validators = append(validators, validator.NewClaudeValidator())
		log.Println("使用 Claude 验证器")
	case "all":
		validators = []validator.Validator{
			validator.NewOpenAIValidator("https://api.openai.com/v1/chat/completions", "gpt-3.5-turbo"),
			validator.NewGeminiValidator(),
			validator.NewClaudeValidator(),
		}
		log.Println("使用所有验证器")
	default:
		log.Fatalf("无效的验证器类型: %s. 请使用 'openai', 'gemini', 'claude', 或 'all'", *validatorType)
	}
	validKeys := make(chan string)
	doneValidating := make(chan struct{})

	go func() {
		defer close(doneValidating)
		var wg sync.WaitGroup
		processedKeys := make(map[string]bool)
		var mu sync.Mutex

		for key := range potentialKeys {
			mu.Lock()
			if processedKeys[key] {
				mu.Unlock()
				continue
			}
			processedKeys[key] = true
			mu.Unlock()

			wg.Add(1)
			go func(k string) {
				defer wg.Done()
				for _, v := range validators {
					isValid, err := v.Validate(k)
					if err != nil {
						// 不记录错误，因为一个 key 可能对一个验证器无效，但对另一个有效
						continue
					}
					if isValid {
						validKeys <- k
						break // 一旦验证通过，就不再用其他验证器尝试
					}
				}
			}(key)
		}
		wg.Wait()
	}()

	// 输出器
	writer := output.NewFileWriter(outputFile)
	defer func() {
		if err := writer.Close(); err != nil {
			log.Printf("关闭输出文件失败: %v", err)
		}
	}()

	var wgOutput sync.WaitGroup
	wgOutput.Add(1)
	go func() {
		defer wgOutput.Done()
		for key := range validKeys {
			err := writer.WriteKey(key)
			if err != nil {
				log.Printf("写入密钥 %s 到文件时发生错误: %v", key, err)
			} else {
				fmt.Printf("发现并验证通过密钥: %s\n", key)
			}
		}
	}()

	// 等待扫描和验证完成
	<-doneScanning
	close(potentialKeys) // 关闭 potentialKeys channel，通知验证器没有更多密钥
	<-doneValidating
	close(validKeys) // 关闭 validKeys channel，通知输出器没有更多密钥

	wgOutput.Wait()

	fmt.Println("所有操作完成。")
}