rate

go get golang.org/x/time/rate

type TokenBucket struct {
    capacity       int64     // 桶最大容量
    tokens         float64   // 当前令牌数量
    rate           float64   // 每秒生成令牌速度
    lastRefillTime int64     // 上次补充时间
}

type Limiter struct {
	mu     sync.Mutex
	limit  Limit
	burst  int
	tokens float64
	// last is the last time the limiter's tokens field was updated
	last time.Time
	// lastEvent is the latest time of a rate-limited event (past or future)
	lastEvent time.Time
}

limiter := rate.NewLimiter(2, 5)

package main

import (
	"context"
	"fmt"
	"golang.org/x/time/rate"
	"time"
)

func main() {

	//  0.5 个 token每秒
	// 桶容量最大 5
	limiter := rate.NewLimiter(0.5, 5)

	ctx := context.Background()

	fmt.Println("====== 初始 burst 测试 ======")

	// 初始桶是满的
	for i := 1; i <= 7; i++ {

		ok := limiter.Allow()

		fmt.Printf(
			"时间=%s 第%d次 Allow() => %v 当前Token≈%.6f\n",
			time.Now().Format("15:04:05.000"),
			i,
			ok,
			limiter.Tokens(),
		)
	}

	fmt.Println()
	fmt.Println("====== 等待 token 自动恢复 ======")

	// 等待 2 秒
	time.Sleep(2 * time.Second)

	fmt.Printf(
		"2秒后 Token≈%.6f\n",
		limiter.Tokens(),
	)

	fmt.Println()

	// 再次消费
	for i := 1; i <= 3; i++ {

		ok := limiter.Allow()

		fmt.Printf(
			"恢复后 第%d次 Allow() => %v Token≈%.6f\n",
			i,
			ok,
			limiter.Tokens(),
		)
	}

	fmt.Println()
	fmt.Println("====== Wait() 测试 ======")

	// 清空 token
	for limiter.Allow() {
	}

	fmt.Printf(
		"清空后 Token≈%.6f\n",
		limiter.Tokens(),
	)

	fmt.Println()

	// Wait 会阻塞直到 token 恢复
	for i := 1; i <= 3; i++ {

		start := time.Now()

		err := limiter.Wait(ctx)
		if err != nil {
			panic(err)
		}

		cost := time.Since(start)

		fmt.Printf(
			"第%d次 Wait() 等待时间=%v Token≈%.6f 时间=%s\n",
			i,
			cost,
			limiter.Tokens(),
			time.Now().Format("15:04:05.000"),
		)
	}
}

package main

import (
	"context"
	"fmt"
	"net"
	"strconv"
	"strings"
	"time"

	"golang.org/x/time/rate"
)

const Timeout = 2 * time.Second

// ---------------- 解析端口 ----------------
func parsePorts(portInput string) ([]int, error) {
	var ports []int

	// 判断是否是范围
	if strings.Contains(portInput, "-") {
		parts := strings.Split(portInput, "-")
		if len(parts) != 2 {
			return nil, fmt.Errorf("端口范围格式错误")
		}

		start, err1 := strconv.Atoi(parts[0])
		end, err2 := strconv.Atoi(parts[1])

		if err1 != nil || err2 != nil || start > end {
			return nil, fmt.Errorf("端口范围非法")
		}

		for i := start; i <= end; i++ {
			ports = append(ports, i)
		}
	} else {
		// 单个端口
		p, err := strconv.Atoi(portInput)
		if err != nil {
			return nil, fmt.Errorf("端口格式错误")
		}
		ports = append(ports, p)
	}

	return ports, nil
}

// ---------------- 扫描函数 ----------------
func scanPorts(ip string, ports []string) {
	// 创建速率限制器：每秒最多 10 个请求
	limiter := rate.NewLimiter(rate.Limit(10), 10)
	ctx := context.Background()

	for _, port := range ports {
		port_final, err := parsePorts(port)
		if err != nil {
			fmt.Println("解析错误:", err)
			return
		}
		for _, port_scan := range port_final {
			// ==================== QPS 限制 ====================
			// 每次扫描端口前等待 token
			if err := limiter.Wait(ctx); err != nil {
				fmt.Printf("速率限制错误: %v\n", err)
				continue
			}

			address := fmt.Sprintf("%s:%d", ip, port_scan)

			conn, err := net.DialTimeout("tcp", address, Timeout)
			if err == nil {
				conn.Close()
				fmt.Printf("✔ %d open\n", port_scan)
			}
		}
	}
}

// ---------------- main ----------------
func main() {
	ip := "www.baidu.com"
	port_field := []string{"80", "443"}

	scanPorts(ip, port_field)
}

package main

import (
	"bytes"
	"context"
	"encoding/xml"
	"fmt"
	"os/exec"
	"strings"
	"time"
)

type ServiceResult struct {
	IP      string
	Port    int
	State   string
	Service string
	Product string
	Version string
}

// 只保留最关键的 XML 结构
type nmapResult struct {
	Hosts []struct {
		Ports []struct {
			PortID   int    `xml:"portid,attr"`
			Protocol string `xml:"protocol,attr"`

			State struct {
				State string `xml:"state,attr"`
			} `xml:"state"`

			Service struct {
				Name    string `xml:"name,attr"`
				Product string `xml:"product,attr"`
				Version string `xml:"version,attr"` //将product+version其存入数据库用于前端显示，但这里并不参与后续的逻辑处理
			} `xml:"service"`
		} `xml:"ports>port"`
	} `xml:"host"`
}

func IdentifyService(host string, ports []string) ([]ServiceResult, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) ////给足够长的时间用于程序运行，覆盖dns解析等消耗的时长
	defer cancel()

	portScan := strings.Join(ports, ",")
	if portScan == "" {
		return nil, fmt.Errorf("port不能为空")
	}

	cmd := exec.CommandContext(ctx,
		"nmap",
		"-sV",               //探测开放端口以获取服务/版本信息，对于服务探测非常有必要
		"-Pn",               //将所有主机视为已上线状态，跳过主机发现步骤（因为前面的主机存活以及端口扫描模块已经完成了对应的内容）
		"-n",                //禁止反向域名解析
		"--max-rate", "100", //每秒最多发送多少请求包
		"-p", portScan,
		"-oX", "-", //结果输出为xml格式到stdout中
		host,
	)

	var out bytes.Buffer
	cmd.Stdout = &out

	if err := cmd.Run(); err != nil {
		return nil, fmt.Errorf("nmap 执行失败: %v", err)
	}

	var result nmapResult
	if err := xml.Unmarshal(out.Bytes(), &result); err != nil {
		return nil, fmt.Errorf("解析失败: %v", err)
	}

	if len(result.Hosts) == 0 {
		return nil, fmt.Errorf("无结果")
	}

	h := result.Hosts[0] //从流程来看只需要一个即可，但是看处理方式要不要将多个host的结果合并在一起，就可以使用数组，后面综合起来的时候再看看

	if len(h.Ports) == 0 {
		return nil, fmt.Errorf("无端口信息")
	}

	var results []ServiceResult

	for _, p := range h.Ports {
		results = append(results, ServiceResult{
			IP:      host,
			Port:    p.PortID,
			State:   p.State.State,
			Service: p.Service.Name,
			Product: p.Service.Product,
			Version: p.Service.Version,
		})
	}

	return results, nil
}

func main() {
	ports := []string{"3306", "8888"}

	start := time.Now()

	results, err := IdentifyService("127.0.0.1", ports)
	if err != nil {
		fmt.Println(err)
		return
	}

	for _, r := range results {
		fmt.Printf("IP=%s Port=%d Service=%s State=%s Version=%s\n",
			r.IP, r.Port, r.Service, r.State, r.Product+" "+r.Version)
	}

	end := time.Since(start)
	fmt.Println("整体运行消耗时间为：", end)
}

nuclei.WithGlobalRateLimitCtx(ctx, 100, time.Second) // 先定总上限
nuclei.WithConcurrency(nuclei.Concurrency{
    HostConcurrency:             50,
    TemplateConcurrency:         10,
    TemplatePayloadConcurrency:  10,
    ProbeConcurrency:            50,
})

package main

import (
	"context"
	"fmt"
	"time"

	nuclei "github.com/projectdiscovery/nuclei/v3/lib"
	"github.com/projectdiscovery/nuclei/v3/pkg/catalog/disk"
	"github.com/projectdiscovery/nuclei/v3/pkg/output"
	"os"
	"path/filepath"
)

func main() {
	// 扫描目标数组（支持单目标或多目标）
	targets := []string{
		"127.0.0.1:8888",
	}

	// 扫描模板 ID 数组；为空则表示全量扫描
	templateIDs := []string{"local-flask", "multi-request-check"}

	//定义模版目录
	home, err := os.UserHomeDir()
	if err != nil {
		fmt.Println(err)
	}
	templateDir := filepath.Join(home, "nuclei-templates")

	start := time.Now()

	if err := runNucleiScan(context.Background(), targets, templateIDs, templateDir); err != nil {
		fmt.Fprintln(os.Stderr, "scan failed:", err)
		os.Exit(1)
	}

	end := time.Since(start)
	fmt.Println("运行消耗的时间为：", end)

}

func runNucleiScan(ctx context.Context, targets []string, templateIDs []string, templateDir string) error {
	if len(targets) == 0 {
		return fmt.Errorf("targets is empty")
	}

	opts := []nuclei.NucleiSDKOptions{
		nuclei.WithCatalog(disk.NewCatalog(templateDir)), //模版目录
		nuclei.DisableUpdateCheck(),

		// 全局速率限制：200 QPS
		nuclei.WithGlobalRateLimitCtx(ctx, 200, time.Second),
	}

	// 指定模板 ID 时，相当于 CLI 的 -id id1,id2
	if len(templateIDs) > 0 {
		opts = append(opts, nuclei.WithTemplateFilters(nuclei.TemplateFilters{
			IDs: templateIDs,
		}))
	}

	engine, err := nuclei.NewNucleiEngineCtx(ctx, opts...)
	if err != nil {
		return fmt.Errorf("create nuclei engine failed: %w", err)
	}
	defer engine.Close()

	// 全量扫描 / 指定 ID 扫描，注意这里只是加载了全量模版，但是如果指定了filter，就只会扫描指定id的模版文件
	if err := engine.LoadAllTemplates(); err != nil {
		return fmt.Errorf("load templates failed: %w", err)
	}

	// 这里直接加载对应的目标数组即可，另类实现-l
	engine.LoadTargets(targets, true) //如果是false的话就输入目标为ip就必须加上http，只走https了，故设置为true用于自行探测是什么协议

	// 回调里直接输出命中的结果
	return engine.ExecuteCallbackWithCtx(ctx, func(ev *output.ResultEvent) {
		if ev == nil || !ev.MatcherStatus {
			return
		}
		sev := ev.Info.SeverityHolder.Severity.String()
		if sev == "" {
			sev = "unknown"
		}
		fmt.Printf("[VULN] severity=%s template=%s name=%s Matched=%s\n",
			sev,
			ev.TemplateID,
			ev.Info.Name,
			ev.Matched,
		)
	})
}

package main

import (
	"context"
	"fmt"
	"sync"
	"time"

	"Go_ENV/checker"
	"Go_ENV/dictionary"
	"Go_ENV/model"
	"golang.org/x/time/rate"
)

// ServiceInfo 服务信息
type ServiceInfo struct {
	Service string
	Port    int
}

// HostTarget 主机目标（双层结构）
type HostTarget struct {
	Host     string
	Services []ServiceInfo
}

func main() {
	// ==================== 双层数据结构 ====================
	// 第一层：Host
	// 第二层：该 Host 上识别出的服务
	targets := []HostTarget{
		{
			Host: "127.0.0.1",
			Services: []ServiceInfo{
				{"ssh", 22},
				{"mysql", 3306},
				{"redis", 6379},
			},
		},
		// 可以添加更多 Host
		// {
		// 	Host: "192.168.1.100",
		// 	Services: []ServiceInfo{
		// 		{"ssh", 22},
		// 		{"mysql", 3306},
		// 	},
		// },
	}

	// QPS 配置
	const qpsLimit = 1
	const burstSize = 10

	// 使用 WaitGroup 支持多 Host 并发
	var wg sync.WaitGroup

	// 遍历每个 Host（可并发）
	for _, hostTarget := range targets {
		wg.Add(1)
		go func(ht HostTarget) {
			defer wg.Done()
			scanHost(ht, qpsLimit, burstSize)
		}(hostTarget)
	}

	wg.Wait()
	fmt.Println("\n所有主机扫描完成！")
}

// scanHost 扫描单个主机的所有服务（带 QPS 限制）
func scanHost(hostTarget HostTarget, qpsLimit, burstSize int) {
	fmt.Printf("\n========================================\n")
	fmt.Printf("开始扫描主机: %s\n", hostTarget.Host)
	fmt.Printf("服务数量: %d\n", len(hostTarget.Services))
	fmt.Printf("QPS 限制: %d 请求/秒\n", qpsLimit)
	fmt.Printf("========================================\n")

	// ==================== 为该 Host 创建独立的速率限制器 ====================
	limiter := rate.NewLimiter(rate.Limit(qpsLimit), burstSize)
	ctx := context.Background()

	// 遍历该 Host 的所有服务
	for _, svc := range hostTarget.Services {
		fmt.Printf("\n[%s] 开始爆破 %s://%s:%d\n", hostTarget.Host, svc.Service, hostTarget.Host, svc.Port)

		// 获取对应服务的弱口令字典
		credentials := dictionary.GetCredentials(svc.Service)
		fmt.Printf("[%s] 字典大小: %d 组凭证\n", hostTarget.Host, len(credentials))

		// 根据字典生成任务列表
		tasks := make([]model.Task, 0, len(credentials))
		for _, cred := range credentials {
			task := model.Task{
				Service:  svc.Service,
				Host:     hostTarget.Host,
				Port:     svc.Port,
				Username: cred.Username,
				Password: cred.Password,
				Timeout:  3 * time.Second,
			}

			// 针对不同服务设置特定参数
			switch svc.Service {
			case "mysql":
				task.MySQLDB = ""
			case "redis":
				task.RedisDB = 0
			}

			tasks = append(tasks, task)
		}

		// ==================== 应用 QPS 限制 ====================
		dispatcher := checker.NewDispatcher()
		startTime := time.Now()

		for _, task := range tasks {
			if err := limiter.Wait(ctx); err != nil {
				continue
			}

			result := dispatcher.Run(ctx, task)
			if result.OK {
				fmt.Printf("[%s] ✓ service=%s user='%s' pass='%s'\n",
					hostTarget.Host, result.Service, result.Username, result.Password)
			}
		}

		elapsed := time.Since(startTime)
		actualQPS := float64(len(tasks)) / elapsed.Seconds()
		fmt.Printf("[%s] %s 完成: 请求数=%d 耗时=%.2fs 实际QPS=%.2f\n",
			hostTarget.Host, svc.Service, len(tasks), elapsed.Seconds(), actualQPS)
	}

	fmt.Printf("\n[%s] 主机扫描完成！\n", hostTarget.Host)
}

package checker

import (
	"context"
	"fmt"
	"strings"

	"Go_ENV/model"
	"Go_ENV/services"
)

type Dispatcher struct {
	registry map[string]Checker
}

func NewDispatcher() *Dispatcher {
	return &Dispatcher{
		registry: map[string]Checker{
			"ssh":   services.NewSSHChecker(),
			"mysql": services.NewMySQLChecker(),
			"redis": services.NewRedisChecker(),
		},
	}
}

func (d *Dispatcher) Run(ctx context.Context, task model.Task) model.Result {
	task.Service = strings.ToLower(strings.TrimSpace(task.Service))
	checker, ok := d.registry[task.Service]
	if !ok {
		return model.Result{
			Service:  task.Service,
			Target:   fmt.Sprintf("%s:%d", task.Host, task.Port),
			Username: task.Username,
			Password: task.Password,
			OK:       false,
			Err:      fmt.Errorf("unsupported service: %s", task.Service),
		}
	}

	return checker.Check(ctx, task)
}

opts := []nuclei.NucleiSDKOptions{
		nuclei.WithCatalog(disk.NewCatalog(templateDir)),
		nuclei.DisableUpdateCheck(),

		nuclei.WithGlobalRateLimitCtx(ctx, qpsPerTarget, time.Second),

		nuclei.WithConcurrency(nuclei.Concurrency{
			TemplateConcurrency:           25,
			HostConcurrency:               qpsPerTarget,
			HeadlessHostConcurrency:       qpsPerTarget,
			HeadlessTemplateConcurrency:   10,
			JavascriptTemplateConcurrency: 10,
			TemplatePayloadConcurrency:    25,
			ProbeConcurrency:              50,
		}),
	}