perf: 日常优化

2026-02-09 02:09:17 +08:00 · 2025-05-05 04:00:35 +08:00
parent 2b4a4024b8
commit 0dc4a6c360
5 changed files with 809 additions and 786 deletions
--- a/WebScan/WebScan.go
+++ b/WebScan/WebScan.go
@@ -1,90 +1,171 @@
 package WebScan

 import (
+	"context"
 	"embed"
+	"errors"
 	"fmt"
-	"github.com/shadow1ng/fscan/Common"
-	"github.com/shadow1ng/fscan/WebScan/lib"
 	"net/http"
 	"net/url"
 	"os"
 	"path/filepath"
 	"strings"
 	"sync"
+	"time"
+
+	"github.com/shadow1ng/fscan/Common"
+	"github.com/shadow1ng/fscan/WebScan/lib"
+)
+
+// 常量定义
+const (
+	protocolHTTP     = "http://"
+	protocolHTTPS    = "https://"
+	yamlExt          = ".yaml"
+	ymlExt           = ".yml"
+	defaultTimeout   = 30 * time.Second
+	concurrencyLimit = 10 // 并发加载POC的限制
+)
+
+// 错误定义
+var (
+	ErrInvalidURL    = errors.New("无效的URL格式")
+	ErrEmptyTarget   = errors.New("目标URL为空")
+	ErrPocNotFound   = errors.New("未找到匹配的POC")
+	ErrPocLoadFailed = errors.New("POC加载失败")
 )

 //go:embed pocs
-var Pocs embed.FS
-var once sync.Once
-var AllPocs []*lib.Poc
+var pocsFS embed.FS
+var (
+	once    sync.Once
+	allPocs []*lib.Poc
+)

 // WebScan 执行Web漏洞扫描
 func WebScan(info *Common.HostInfo) {
-	once.Do(initpoc)
+	// 初始化POC
+	once.Do(initPocs)

-	var pocinfo = Common.Pocinfo
-
-	// 自动构建URL
-	if info.Url == "" {
-		info.Url = fmt.Sprintf("http://%s:%s", info.Host, info.Ports)
+	// 验证输入
+	if info == nil {
+		Common.LogError("无效的扫描目标")
+		return
 	}

-	urlParts := strings.Split(info.Url, "/")
-
-	// 检查切片长度并构建目标URL
-	if len(urlParts) >= 3 {
-		pocinfo.Target = strings.Join(urlParts[:3], "/")
-	} else {
-		pocinfo.Target = info.Url
+	if len(allPocs) == 0 {
+		Common.LogError("POC加载失败，无法执行扫描")
+		return
 	}

-	Common.LogDebug(fmt.Sprintf("扫描目标: %s", pocinfo.Target))
+	// 构建目标URL
+	target, err := buildTargetURL(info)
+	if err != nil {
+		Common.LogError(fmt.Sprintf("构建目标URL失败: %v", err))
+		return
+	}

-	// 如果是直接调用WebPoc（没有指定pocName），执行所有POC
-	if pocinfo.PocName == "" && len(info.Infostr) == 0 {
-		Common.LogDebug("直接调用WebPoc，执行所有POC")
-		Execute(pocinfo)
-	} else {
-		// 根据指纹信息选择性执行POC
-		if len(info.Infostr) > 0 {
-			for _, infostr := range info.Infostr {
-				pocinfo.PocName = lib.CheckInfoPoc(infostr)
-				if pocinfo.PocName != "" {
-					Common.LogDebug(fmt.Sprintf("根据指纹 %s 执行对应POC", infostr))
-					Execute(pocinfo)
-				}
-			}
-		} else if pocinfo.PocName != "" {
-			// 指定了特定的POC
-			Common.LogDebug(fmt.Sprintf("执行指定POC: %s", pocinfo.PocName))
-			Execute(pocinfo)
-		}
+	// 使用带超时的上下文
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTimeout)
+	defer cancel()
+
+	// 根据扫描策略执行POC
+	if Common.Pocinfo.PocName == "" && len(info.Infostr) == 0 {
+		// 执行所有POC
+		executePOCs(ctx, Common.PocInfo{Target: target})
+	} else if len(info.Infostr) > 0 {
+		// 基于指纹信息执行POC
+		scanByFingerprints(ctx, target, info.Infostr)
+	} else if Common.Pocinfo.PocName != "" {
+		// 基于指定POC名称执行
+		executePOCs(ctx, Common.PocInfo{Target: target, PocName: Common.Pocinfo.PocName})
 	}
 }

-// Execute 执行具体的POC检测
-func Execute(PocInfo Common.PocInfo) {
-	Common.LogDebug(fmt.Sprintf("开始执行POC检测，目标: %s", PocInfo.Target))
+// buildTargetURL 构建规范的目标URL
+func buildTargetURL(info *Common.HostInfo) (string, error) {
+	// 自动构建URL
+	if info.Url == "" {
+		info.Url = fmt.Sprintf("%s%s:%s", protocolHTTP, info.Host, info.Ports)
+	} else if !hasProtocolPrefix(info.Url) {
+		info.Url = protocolHTTP + info.Url
+	}
+
+	// 解析URL以提取基础部分
+	parsedURL, err := url.Parse(info.Url)
+	if err != nil {
+		return "", fmt.Errorf("%w: %v", ErrInvalidURL, err)
+	}
+
+	return fmt.Sprintf("%s://%s", parsedURL.Scheme, parsedURL.Host), nil
+}
+
+// hasProtocolPrefix 检查URL是否包含协议前缀
+func hasProtocolPrefix(urlStr string) bool {
+	return strings.HasPrefix(urlStr, protocolHTTP) || strings.HasPrefix(urlStr, protocolHTTPS)
+}
+
+// scanByFingerprints 根据指纹执行POC
+func scanByFingerprints(ctx context.Context, target string, fingerprints []string) {
+	for _, fingerprint := range fingerprints {
+		if fingerprint == "" {
+			continue
+		}
+
+		pocName := lib.CheckInfoPoc(fingerprint)
+		if pocName == "" {
+			continue
+		}
+
+		executePOCs(ctx, Common.PocInfo{Target: target, PocName: pocName})
+	}
+}
+
+// executePOCs 执行POC检测
+func executePOCs(ctx context.Context, pocInfo Common.PocInfo) {
+	// 验证目标
+	if pocInfo.Target == "" {
+		Common.LogError(ErrEmptyTarget.Error())
+		return
+	}

 	// 确保URL格式正确
-	if !strings.HasPrefix(PocInfo.Target, "http://") && !strings.HasPrefix(PocInfo.Target, "https://") {
-		PocInfo.Target = "http://" + PocInfo.Target
+	if !hasProtocolPrefix(pocInfo.Target) {
+		pocInfo.Target = protocolHTTP + pocInfo.Target
 	}

-	// 验证URL格式
-	_, err := url.Parse(PocInfo.Target)
+	// 验证URL
+	_, err := url.Parse(pocInfo.Target)
 	if err != nil {
-		Common.LogError(fmt.Sprintf("无效的URL格式 %v: %v", PocInfo.Target, err))
+		Common.LogError(fmt.Sprintf("%v %s: %v", ErrInvalidURL, pocInfo.Target, err))
 		return
 	}

-	// 创建基础HTTP请求
-	req, err := http.NewRequest("GET", PocInfo.Target, nil)
+	// 创建基础请求
+	req, err := createBaseRequest(ctx, pocInfo.Target)
 	if err != nil {
-		Common.LogError(fmt.Sprintf("初始化请求失败 %v: %v", PocInfo.Target, err))
+		Common.LogError(fmt.Sprintf("创建HTTP请求失败: %v", err))
 		return
 	}

+	// 筛选POC
+	matchedPocs := filterPocs(pocInfo.PocName)
+	if len(matchedPocs) == 0 {
+		Common.LogDebug(fmt.Sprintf("%v: %s", ErrPocNotFound, pocInfo.PocName))
+		return
+	}
+
+	// 执行POC检测
+	lib.CheckMultiPoc(req, matchedPocs, Common.PocNum)
+}
+
+// createBaseRequest 创建带上下文的HTTP请求
+func createBaseRequest(ctx context.Context, target string) (*http.Request, error) {
+	req, err := http.NewRequestWithContext(ctx, "GET", target, nil)
+	if err != nil {
+		return nil, err
+	}
+
 	// 设置请求头
 	req.Header.Set("User-agent", Common.UserAgent)
 	req.Header.Set("Accept", Common.Accept)
@@ -93,75 +174,150 @@ func Execute(PocInfo Common.PocInfo) {
 		req.Header.Set("Cookie", Common.Cookie)
 	}

-	// 根据名称筛选POC并执行
-	pocs := filterPoc(PocInfo.PocName)
-	Common.LogDebug(fmt.Sprintf("筛选到的POC数量: %d", len(pocs)))
-	lib.CheckMultiPoc(req, pocs, Common.PocNum)
+	return req, nil
 }

-// initpoc 初始化POC加载
-func initpoc() {
-	Common.LogDebug("开始初始化POC")
+// initPocs 初始化并加载POC
+func initPocs() {
+	allPocs = make([]*lib.Poc, 0)

 	if Common.PocPath == "" {
-		Common.LogDebug("从内置目录加载POC")
-		// 从嵌入的POC目录加载
-		entries, err := Pocs.ReadDir("pocs")
-		if err != nil {
-			Common.LogError(fmt.Sprintf("加载内置POC失败: %v", err))
-			return
-		}
-
-		// 加载YAML格式的POC文件
-		for _, entry := range entries {
-			filename := entry.Name()
-			if strings.HasSuffix(filename, ".yaml") || strings.HasSuffix(filename, ".yml") {
-				if poc, err := lib.LoadPoc(filename, Pocs); err == nil && poc != nil {
-					AllPocs = append(AllPocs, poc)
-				} else if err != nil {
-				}
-			}
-		}
-		Common.LogDebug(fmt.Sprintf("内置POC加载完成，共加载 %d 个", len(AllPocs)))
+		loadEmbeddedPocs()
 	} else {
-		// 从指定目录加载POC
-		Common.LogSuccess(fmt.Sprintf("从目录加载POC: %s", Common.PocPath))
-		err := filepath.Walk(Common.PocPath, func(path string, info os.FileInfo, err error) error {
-			if err != nil || info == nil {
-				return err
-			}
-
-			if !info.IsDir() && (strings.HasSuffix(path, ".yaml") || strings.HasSuffix(path, ".yml")) {
-				if poc, err := lib.LoadPocbyPath(path); err == nil && poc != nil {
-					AllPocs = append(AllPocs, poc)
-				} else if err != nil {
-				}
-			}
-			return nil
-		})
-
-		if err != nil {
-			Common.LogError(fmt.Sprintf("加载外部POC失败: %v", err))
-		}
-		Common.LogDebug(fmt.Sprintf("外部POC加载完成，共加载 %d 个", len(AllPocs)))
+		loadExternalPocs(Common.PocPath)
 	}
 }

-// filterPoc 根据POC名称筛选
-func filterPoc(pocname string) []*lib.Poc {
-	Common.LogDebug(fmt.Sprintf("开始筛选POC，筛选条件: %s", pocname))
-
-	if pocname == "" {
-		Common.LogDebug(fmt.Sprintf("未指定POC名称，返回所有POC: %d 个", len(AllPocs)))
-		return AllPocs
+// loadEmbeddedPocs 加载内置POC
+func loadEmbeddedPocs() {
+	entries, err := pocsFS.ReadDir("pocs")
+	if err != nil {
+		Common.LogError(fmt.Sprintf("加载内置POC目录失败: %v", err))
+		return
 	}

+	// 收集所有POC文件
+	var pocFiles []string
+	for _, entry := range entries {
+		if isPocFile(entry.Name()) {
+			pocFiles = append(pocFiles, entry.Name())
+		}
+	}
+
+	// 并发加载POC文件
+	loadPocsConcurrently(pocFiles, true, "")
+}
+
+// loadExternalPocs 从外部路径加载POC
+func loadExternalPocs(pocPath string) {
+	if !directoryExists(pocPath) {
+		Common.LogError(fmt.Sprintf("POC目录不存在: %s", pocPath))
+		return
+	}
+
+	// 收集所有POC文件路径
+	var pocFiles []string
+	err := filepath.Walk(pocPath, func(path string, info os.FileInfo, err error) error {
+		if err != nil || info == nil || info.IsDir() {
+			return nil
+		}
+
+		if isPocFile(info.Name()) {
+			pocFiles = append(pocFiles, path)
+		}
+		return nil
+	})
+
+	if err != nil {
+		Common.LogError(fmt.Sprintf("遍历POC目录失败: %v", err))
+		return
+	}
+
+	// 并发加载POC文件
+	loadPocsConcurrently(pocFiles, false, pocPath)
+}
+
+// loadPocsConcurrently 并发加载POC文件
+func loadPocsConcurrently(pocFiles []string, isEmbedded bool, pocPath string) {
+	pocCount := len(pocFiles)
+	if pocCount == 0 {
+		return
+	}
+
+	var wg sync.WaitGroup
+	var mu sync.Mutex
+	var successCount, failCount int
+
+	// 使用信号量控制并发数
+	semaphore := make(chan struct{}, concurrencyLimit)
+
+	for _, file := range pocFiles {
+		wg.Add(1)
+		semaphore <- struct{}{} // 获取信号量
+
+		go func(filename string) {
+			defer func() {
+				<-semaphore // 释放信号量
+				wg.Done()
+			}()
+
+			var poc *lib.Poc
+			var err error
+
+			// 根据不同的来源加载POC
+			if isEmbedded {
+				poc, err = lib.LoadPoc(filename, pocsFS)
+			} else {
+				poc, err = lib.LoadPocbyPath(filename)
+			}
+
+			mu.Lock()
+			defer mu.Unlock()
+
+			if err != nil {
+				failCount++
+				return
+			}
+
+			if poc != nil {
+				allPocs = append(allPocs, poc)
+				successCount++
+			}
+		}(file)
+	}
+
+	wg.Wait()
+	Common.LogSuccess(fmt.Sprintf("POC加载完成: 总共%d个，成功%d个，失败%d个",
+		pocCount, successCount, failCount))
+}
+
+// directoryExists 检查目录是否存在
+func directoryExists(path string) bool {
+	info, err := os.Stat(path)
+	return err == nil && info.IsDir()
+}
+
+// isPocFile 检查文件是否为POC文件
+func isPocFile(filename string) bool {
+	lowerName := strings.ToLower(filename)
+	return strings.HasSuffix(lowerName, yamlExt) || strings.HasSuffix(lowerName, ymlExt)
+}
+
+// filterPocs 根据POC名称筛选
+func filterPocs(pocName string) []*lib.Poc {
+	if pocName == "" {
+		return allPocs
+	}
+
+	// 转换为小写以进行不区分大小写的匹配
+	searchName := strings.ToLower(pocName)
+
 	var matchedPocs []*lib.Poc
-	for _, poc := range AllPocs {
-		if strings.Contains(poc.Name, pocname) {
+	for _, poc := range allPocs {
+		if poc != nil && strings.Contains(strings.ToLower(poc.Name), searchName) {
 			matchedPocs = append(matchedPocs, poc)
 		}
 	}
-	Common.LogDebug(fmt.Sprintf("POC筛选完成，匹配到 %d 个", len(matchedPocs)))
+
 	return matchedPocs
 }