perf: 优化代码结构

Core/ICMP.go

@@ -45,6 +45,16 @@ func CheckLive(hostslist []string, Ping bool) []string {
 	return AliveHosts
 }
 
+// IsContain 检查切片中是否包含指定元素
+func IsContain(items []string, item string) bool {
+	for _, eachItem := range items {
+		if eachItem == item {
+			return true
+		}
+	}
+	return false
+}
+
 func handleAliveHosts(chanHosts chan string, hostslist []string, isPing bool) {
 	for ip := range chanHosts {
 		if _, ok := ExistHosts[ip]; !ok && IsContain(hostslist, ip) {

Core/Scanner.go

@@ -29,103 +29,121 @@ func Scan(info Common.HostInfo) {
 	ch := make(chan struct{}, Common.ThreadNum)
 	wg := sync.WaitGroup{}
 
-	// 根据不同模式执行扫描
+	// 执行扫描逻辑
 	switch {
 	case Common.LocalMode:
-		// 本地信息收集模式
-		LocalScan = true
-
-		// 定义本地模式允许的插件
-		validLocalPlugins := make(map[string]bool)
-		for _, plugin := range Common.PluginGroups[Common.ModeLocal] {
-			validLocalPlugins[plugin] = true
-		}
-
-		// 如果没有指定扫描模式或为默认的All，设置为 ModeLocal
-		if Common.ScanMode == "" || Common.ScanMode == "All" {
-			Common.ScanMode = Common.ModeLocal
-		} else if Common.ScanMode != Common.ModeLocal {
-			// 不是完整模式时，检查是否是合法的单插件
-			if !validLocalPlugins[Common.ScanMode] {
-				Common.LogError(fmt.Sprintf("无效的本地模式插件: %s, 仅支持 localinfo", Common.ScanMode))
-				return
-			}
-		}
-
-		if Common.ScanMode == Common.ModeLocal {
-			Common.LogInfo("执行本地信息收集 - 使用全部本地插件")
-		} else {
-			Common.LogInfo(fmt.Sprintf("执行本地信息收集 - 使用插件: %s", Common.ScanMode))
-		}
-
-		executeScans([]Common.HostInfo{info}, &ch, &wg)
-
+		executeLocalScan(info, &ch, &wg)
 	case len(Common.URLs) > 0:
-		// Web模式
-		WebScan = true
-
-		// 从 pluginGroups 获取Web模式允许的插件
-		validWebPlugins := make(map[string]bool)
-		for _, plugin := range Common.PluginGroups[Common.ModeWeb] {
-			validWebPlugins[plugin] = true
-		}
-
-		// 如果没有指定扫描模式，默认设置为 ModeWeb
-		if Common.ScanMode == "" || Common.ScanMode == "All" {
-			Common.ScanMode = Common.ModeWeb
-		}
-
-		// 如果不是 ModeWeb，检查是否是合法的单插件
-		if Common.ScanMode != Common.ModeWeb {
-			if !validWebPlugins[Common.ScanMode] {
-				Common.LogError(fmt.Sprintf("无效的Web插件: %s, 仅支持 webtitle 和 webpoc", Common.ScanMode))
-				return
-			}
-			// ScanMode 保持为单插件名
-		}
-
-		var targetInfos []Common.HostInfo
-		for _, url := range Common.URLs {
-			urlInfo := info
-			if !strings.HasPrefix(url, "http://") && !strings.HasPrefix(url, "https://") {
-				url = "http://" + url
-			}
-			urlInfo.Url = url
-			targetInfos = append(targetInfos, urlInfo)
-		}
-
-		if Common.ScanMode == Common.ModeWeb {
-			Common.LogInfo("开始Web扫描 - 使用全部Web插件")
-		} else {
-			Common.LogInfo(fmt.Sprintf("开始Web扫描 - 使用插件: %s", Common.ScanMode))
-		}
-		executeScans(targetInfos, &ch, &wg)
-
+		executeWebScan(info, &ch, &wg)
 	default:
-		// 主机扫描模式
-		if info.Host == "" {
-			Common.LogError("未指定扫描目标")
-			return
-		}
-
-		hosts, err := Common.ParseIP(info.Host, Common.HostsFile, Common.ExcludeHosts)
-		if err != nil {
-			Common.LogError(fmt.Sprintf("解析主机错误: %v", err))
-			return
-		}
-		Common.LogInfo("开始主机扫描")
-		executeScan(hosts, info, &ch, &wg)
+		executeHostScan(info, &ch, &wg)
 	}
 
+	// 等待扫描完成
 	finishScan(&wg)
 }
 
+// 执行本地扫描
+func executeLocalScan(info Common.HostInfo, ch *chan struct{}, wg *sync.WaitGroup) {
+	Common.LogInfo("执行本地信息收集")
+
+	// 定义本地模式允许的插件
+	validLocalPlugins := getValidPlugins(Common.ModeLocal)
+
+	// 校验扫描模式
+	if err := validateScanMode(validLocalPlugins, Common.ModeLocal); err != nil {
+		Common.LogError(err.Error())
+		return
+	}
+
+	if Common.ScanMode == Common.ModeLocal {
+		Common.LogInfo("使用全部本地插件")
+	} else {
+		Common.LogInfo(fmt.Sprintf("使用插件: %s", Common.ScanMode))
+	}
+
+	// 执行扫描
+	executeScans([]Common.HostInfo{info}, ch, wg)
+}
+
+// 执行Web扫描
+func executeWebScan(info Common.HostInfo, ch *chan struct{}, wg *sync.WaitGroup) {
+	Common.LogInfo("开始Web扫描")
+
+	// 从 pluginGroups 获取Web模式允许的插件
+	validWebPlugins := getValidPlugins(Common.ModeWeb)
+
+	// 校验扫描模式
+	if err := validateScanMode(validWebPlugins, Common.ModeWeb); err != nil {
+		Common.LogError(err.Error())
+		return
+	}
+
+	// 创建目标URL信息
+	var targetInfos []Common.HostInfo
+	for _, url := range Common.URLs {
+		urlInfo := info
+		if !strings.HasPrefix(url, "http://") && !strings.HasPrefix(url, "https://") {
+			url = "http://" + url
+		}
+		urlInfo.Url = url
+		targetInfos = append(targetInfos, urlInfo)
+	}
+
+	if Common.ScanMode == Common.ModeWeb {
+		Common.LogInfo("使用全部Web插件")
+	} else {
+		Common.LogInfo(fmt.Sprintf("使用插件: %s", Common.ScanMode))
+	}
+
+	// 执行扫描
+	executeScans(targetInfos, ch, wg)
+}
+
+// 执行主机扫描
+func executeHostScan(info Common.HostInfo, ch *chan struct{}, wg *sync.WaitGroup) {
+	if info.Host == "" {
+		Common.LogError("未指定扫描目标")
+		return
+	}
+
+	hosts, err := Common.ParseIP(info.Host, Common.HostsFile, Common.ExcludeHosts)
+	if err != nil {
+		Common.LogError(fmt.Sprintf("解析主机错误: %v", err))
+		return
+	}
+
+	Common.LogInfo("开始主机扫描")
+	executeScan(hosts, info, ch, wg)
+}
+
+// 获取合法的插件列表
+func getValidPlugins(mode string) map[string]bool {
+	validPlugins := make(map[string]bool)
+	for _, plugin := range Common.PluginGroups[mode] {
+		validPlugins[plugin] = true
+	}
+	return validPlugins
+}
+
+// 校验扫描模式是否有效
+func validateScanMode(validPlugins map[string]bool, mode string) error {
+	if Common.ScanMode == "" || Common.ScanMode == "All" {
+		Common.ScanMode = mode
+	} else if _, exists := validPlugins[Common.ScanMode]; !exists {
+		return fmt.Errorf("无效的%s插件: %s", mode, Common.ScanMode)
+	}
+	return nil
+}
+
 // executeScan 执行主扫描流程
 func executeScan(hosts []string, info Common.HostInfo, ch *chan struct{}, wg *sync.WaitGroup) {
 	var targetInfos []Common.HostInfo
 
+	// 扫描主机和端口
 	if len(hosts) > 0 || len(Common.HostPort) > 0 {
-		if (Common.DisablePing == false && len(hosts) > 1) || Common.IsICMPScan() {
+		// 处理活跃主机
+		if shouldPingScan(hosts) {
 			hosts = CheckLive(hosts, Common.UsePing)
 			Common.LogInfo(fmt.Sprintf("存活主机数量: %d", len(hosts)))
 			if Common.IsICMPScan() {
@@ -133,39 +151,62 @@ func executeScan(hosts []string, info Common.HostInfo, ch *chan struct{}, wg *sy
 			}
 		}
 
-		var alivePorts []string
-		if Common.IsWebScan() {
-			alivePorts = NoPortScan(hosts, Common.Ports)
-		} else if len(hosts) > 0 {
-			alivePorts = PortScan(hosts, Common.Ports, Common.Timeout)
-			Common.LogInfo(fmt.Sprintf("存活端口数量: %d", len(alivePorts)))
-			if Common.IsPortScan() {
-				return
-			}
+		// 处理活跃端口
+		alivePorts := getAlivePorts(hosts)
+		if len(alivePorts) > 0 {
+			targetInfos = prepareTargetInfos(alivePorts, info)
 		}
-
-		if len(Common.HostPort) > 0 {
-			alivePorts = append(alivePorts, Common.HostPort...)
-			alivePorts = Common.RemoveDuplicate(alivePorts)
-			Common.HostPort = nil
-			Common.LogInfo(fmt.Sprintf("存活端口数量: %d", len(alivePorts)))
-		}
-
-		targetInfos = prepareTargetInfos(alivePorts, info)
 	}
 
-	for _, url := range Common.URLs {
-		urlInfo := info
-		urlInfo.Url = url
-		targetInfos = append(targetInfos, urlInfo)
-	}
+	// 添加 URL 扫描目标
+	targetInfos = appendURLTargets(targetInfos, info)
 
+	// 如果有扫描目标，执行漏洞扫描
 	if len(targetInfos) > 0 {
 		Common.LogInfo("开始漏洞扫描")
 		executeScans(targetInfos, ch, wg)
 	}
 }
 
+// shouldPingScan 判断是否需要进行 ping 扫描
+func shouldPingScan(hosts []string) bool {
+	return (Common.DisablePing == false && len(hosts) > 1) || Common.IsICMPScan()
+}
+
+// getAlivePorts 获取存活端口
+func getAlivePorts(hosts []string) []string {
+	var alivePorts []string
+	if Common.IsWebScan() {
+		alivePorts = NoPortScan(hosts, Common.Ports)
+	} else if len(hosts) > 0 {
+		alivePorts = PortScan(hosts, Common.Ports, Common.Timeout)
+		Common.LogInfo(fmt.Sprintf("存活端口数量: %d", len(alivePorts)))
+		if Common.IsPortScan() {
+			return nil // 结束扫描
+		}
+	}
+
+	// 合并传入的端口信息
+	if len(Common.HostPort) > 0 {
+		alivePorts = append(alivePorts, Common.HostPort...)
+		alivePorts = Common.RemoveDuplicate(alivePorts)
+		Common.HostPort = nil
+		Common.LogInfo(fmt.Sprintf("存活端口数量: %d", len(alivePorts)))
+	}
+
+	return alivePorts
+}
+
+// appendURLTargets 添加 URL 扫描目标
+func appendURLTargets(targetInfos []Common.HostInfo, baseInfo Common.HostInfo) []Common.HostInfo {
+	for _, url := range Common.URLs {
+		urlInfo := baseInfo
+		urlInfo.Url = url
+		targetInfos = append(targetInfos, urlInfo)
+	}
+	return targetInfos
+}
+
 // prepareTargetInfos 准备扫描目标信息
 func prepareTargetInfos(alivePorts []string, baseInfo Common.HostInfo) []Common.HostInfo {
 	var infos []Common.HostInfo
@@ -183,8 +224,58 @@ func prepareTargetInfos(alivePorts []string, baseInfo Common.HostInfo) []Common.
 	return infos
 }
 
+// 扫描任务结构体定义
+type ScanTask struct {
+	pluginName string
+	target     Common.HostInfo
+}
+
+// executeScans 执行扫描任务
 func executeScans(targets []Common.HostInfo, ch *chan struct{}, wg *sync.WaitGroup) {
 	mode := Common.GetScanMode()
+
+	// 获取待执行的插件列表
+	pluginsToRun, isSinglePlugin := getPluginsToRun(mode)
+
+	var tasks []ScanTask
+	actualTasks := 0
+	loadedPlugins := make([]string, 0)
+
+	// 遍历目标，收集任务
+	for _, target := range targets {
+		targetPort, _ := strconv.Atoi(target.Ports)
+		for _, pluginName := range pluginsToRun {
+			plugin, exists := Common.PluginManager[pluginName]
+			if !exists {
+				continue
+			}
+
+			taskAdded, newTasks := collectScanTasks(plugin, target, targetPort, pluginName, isSinglePlugin)
+			if taskAdded {
+				actualTasks += len(newTasks)
+				loadedPlugins = append(loadedPlugins, pluginName)
+				tasks = append(tasks, newTasks...)
+			}
+		}
+	}
+
+	// 去重并排序插件
+	finalPlugins := getUniquePlugins(loadedPlugins)
+
+	// 输出加载的插件信息
+	Common.LogInfo(fmt.Sprintf("加载的插件: %s", strings.Join(finalPlugins, ", ")))
+
+	// 初始化进度条
+	initializeProgressBar(actualTasks)
+
+	// 执行收集的任务
+	for _, task := range tasks {
+		AddScan(task.pluginName, task.target, ch, wg)
+	}
+}
+
+// 获取待执行插件列表
+func getPluginsToRun(mode string) ([]string, bool) {
 	var pluginsToRun []string
 	isSinglePlugin := false
 
@@ -195,82 +286,28 @@ func executeScans(targets []Common.HostInfo, ch *chan struct{}, wg *sync.WaitGro
 		isSinglePlugin = true
 	}
 
-	loadedPlugins := make([]string, 0)
-	actualTasks := 0
+	return pluginsToRun, isSinglePlugin
+}
 
-	type ScanTask struct {
-		pluginName string
-		target     Common.HostInfo
-	}
-	tasks := make([]ScanTask, 0)
+// 收集扫描任务
+func collectScanTasks(plugin Common.ScanPlugin, target Common.HostInfo, targetPort int, pluginName string, isSinglePlugin bool) (bool, []ScanTask) {
+	var tasks []ScanTask
+	taskAdded := false
 
-	// 第一次遍历：计算任务数和收集要执行的插件
-	for _, target := range targets {
-		targetPort, _ := strconv.Atoi(target.Ports)
-
-		for _, pluginName := range pluginsToRun {
-			plugin, exists := Common.PluginManager[pluginName]
-			if !exists {
-				continue
-			}
-
-			// Web模式特殊处理
-			if WebScan {
-				actualTasks++
-				loadedPlugins = append(loadedPlugins, pluginName)
-				tasks = append(tasks, ScanTask{
-					pluginName: pluginName,
-					target:     target,
-				})
-				continue
-			}
-
-			// 本地扫描模式
-			if LocalScan {
-				if len(plugin.Ports) == 0 {
-					actualTasks++
-					loadedPlugins = append(loadedPlugins, pluginName)
-					tasks = append(tasks, ScanTask{
-						pluginName: pluginName,
-						target:     target,
-					})
-				}
-				continue
-			}
-
-			// 单插件模式
-			if isSinglePlugin {
-				actualTasks++
-				loadedPlugins = append(loadedPlugins, pluginName)
-				tasks = append(tasks, ScanTask{
-					pluginName: pluginName,
-					target:     target,
-				})
-				continue
-			}
-
-			// 常规模式
-			if len(plugin.Ports) > 0 {
-				if plugin.HasPort(targetPort) {
-					actualTasks++
-					loadedPlugins = append(loadedPlugins, pluginName)
-					tasks = append(tasks, ScanTask{
-						pluginName: pluginName,
-						target:     target,
-					})
-				}
-			} else {
-				actualTasks++
-				loadedPlugins = append(loadedPlugins, pluginName)
-				tasks = append(tasks, ScanTask{
-					pluginName: pluginName,
-					target:     target,
-				})
-			}
-		}
+	// Web模式特殊处理
+	if WebScan || LocalScan || isSinglePlugin || len(plugin.Ports) == 0 || plugin.HasPort(targetPort) {
+		taskAdded = true
+		tasks = append(tasks, ScanTask{
+			pluginName: pluginName,
+			target:     target,
+		})
 	}
 
-	// 去重并输出实际加载的插件
+	return taskAdded, tasks
+}
+
+// 获取去重后的插件列表
+func getUniquePlugins(loadedPlugins []string) []string {
 	uniquePlugins := make(map[string]struct{})
 	for _, p := range loadedPlugins {
 		uniquePlugins[p] = struct{}{}
@@ -280,11 +317,13 @@ func executeScans(targets []Common.HostInfo, ch *chan struct{}, wg *sync.WaitGro
 	for p := range uniquePlugins {
 		finalPlugins = append(finalPlugins, p)
 	}
+
 	sort.Strings(finalPlugins)
+	return finalPlugins
+}
 
-	Common.LogInfo(fmt.Sprintf("加载的插件: %s", strings.Join(finalPlugins, ", ")))
-
-	// 初始化进度条
+// 初始化进度条
+func initializeProgressBar(actualTasks int) {
 	if Common.ShowProgress {
 		Common.ProgressBar = progressbar.NewOptions(actualTasks,
 			progressbar.OptionEnableColorCodes(true),
@@ -303,11 +342,6 @@ func executeScans(targets []Common.HostInfo, ch *chan struct{}, wg *sync.WaitGro
 			progressbar.OptionSetRenderBlankState(true),
 		)
 	}
-
-	// 执行收集的任务
-	for _, task := range tasks {
-		AddScan(task.pluginName, task.target, ch, wg)
-	}
 }
 
 // finishScan 完成扫描任务
@@ -324,7 +358,7 @@ func finishScan(wg *sync.WaitGroup) {
 // Mutex用于保护共享资源的并发访问
 var Mutex = &sync.Mutex{}
 
-// AddScan
+// AddScan 添加扫描任务并启动扫描
 func AddScan(plugin string, info Common.HostInfo, ch *chan struct{}, wg *sync.WaitGroup) {
 	*ch <- struct{}{}
 	wg.Add(1)
@@ -335,20 +369,14 @@ func AddScan(plugin string, info Common.HostInfo, ch *chan struct{}, wg *sync.Wa
 			<-*ch
 		}()
 
-		Mutex.Lock()
+		// 使用原子操作更新扫描计数
 		atomic.AddInt64(&Common.Num, 1)
-		Mutex.Unlock()
 
+		// 执行扫描插件
 		ScanFunc(&plugin, &info)
 
-		Common.OutputMutex.Lock()
-		atomic.AddInt64(&Common.End, 1)
-		if Common.ProgressBar != nil {
-			// 清除当前行
-			fmt.Print("\033[2K\r")
-			Common.ProgressBar.Add(1)
-		}
-		Common.OutputMutex.Unlock()
+		// 更新扫描结束后的状态
+		updateScanProgress(&info)
 	}()
 }
 
@@ -371,12 +399,15 @@ func ScanFunc(name *string, info *Common.HostInfo) {
 	}
 }
 
-// IsContain 检查切片中是否包含指定元素
-func IsContain(items []string, item string) bool {
-	for _, eachItem := range items {
-		if eachItem == item {
-			return true
-		}
+// updateScanProgress 更新扫描进度
+func updateScanProgress(info *Common.HostInfo) {
+	// 输出互斥锁更新进度条
+	Common.OutputMutex.Lock()
+	atomic.AddInt64(&Common.End, 1)
+	if Common.ProgressBar != nil {
+		// 清除当前行并更新进度条
+		fmt.Print("\033[2K\r")
+		Common.ProgressBar.Add(1)
 	}
-	return false
+	Common.OutputMutex.Unlock()
 }