perf: 优化代码结构

This commit is contained in:
ZacharyZcR 2025-02-07 11:39:04 +08:00
parent c94ec76292
commit 102d100c25
2 changed files with 245 additions and 204 deletions

View File

@ -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) {

View File

@ -29,61 +29,57 @@ func Scan(info Common.HostInfo) {
ch := make(chan struct{}, Common.ThreadNum)
wg := sync.WaitGroup{}
// 根据不同模式执行扫描
// 执行扫描逻辑
switch {
case Common.LocalMode:
// 本地信息收集模式
LocalScan = true
executeLocalScan(info, &ch, &wg)
case len(Common.URLs) > 0:
executeWebScan(info, &ch, &wg)
default:
executeHostScan(info, &ch, &wg)
}
// 等待扫描完成
finishScan(&wg)
}
// 执行本地扫描
func executeLocalScan(info Common.HostInfo, ch *chan struct{}, wg *sync.WaitGroup) {
Common.LogInfo("执行本地信息收集")
// 定义本地模式允许的插件
validLocalPlugins := make(map[string]bool)
for _, plugin := range Common.PluginGroups[Common.ModeLocal] {
validLocalPlugins[plugin] = true
}
validLocalPlugins := getValidPlugins(Common.ModeLocal)
// 如果没有指定扫描模式或为默认的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))
// 校验扫描模式
if err := validateScanMode(validLocalPlugins, Common.ModeLocal); err != nil {
Common.LogError(err.Error())
return
}
}
if Common.ScanMode == Common.ModeLocal {
Common.LogInfo("执行本地信息收集 - 使用全部本地插件")
Common.LogInfo("使用全部本地插件")
} else {
Common.LogInfo(fmt.Sprintf("执行本地信息收集 - 使用插件: %s", Common.ScanMode))
Common.LogInfo(fmt.Sprintf("使用插件: %s", Common.ScanMode))
}
executeScans([]Common.HostInfo{info}, &ch, &wg)
// 执行扫描
executeScans([]Common.HostInfo{info}, ch, wg)
}
case len(Common.URLs) > 0:
// Web模式
WebScan = true
// 执行Web扫描
func executeWebScan(info Common.HostInfo, ch *chan struct{}, wg *sync.WaitGroup) {
Common.LogInfo("开始Web扫描")
// 从 pluginGroups 获取Web模式允许的插件
validWebPlugins := make(map[string]bool)
for _, plugin := range Common.PluginGroups[Common.ModeWeb] {
validWebPlugins[plugin] = true
}
validWebPlugins := getValidPlugins(Common.ModeWeb)
// 如果没有指定扫描模式,默认设置为 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))
// 校验扫描模式
if err := validateScanMode(validWebPlugins, Common.ModeWeb); err != nil {
Common.LogError(err.Error())
return
}
// ScanMode 保持为单插件名
}
// 创建目标URL信息
var targetInfos []Common.HostInfo
for _, url := range Common.URLs {
urlInfo := info
@ -95,14 +91,17 @@ func Scan(info Common.HostInfo) {
}
if Common.ScanMode == Common.ModeWeb {
Common.LogInfo("开始Web扫描 - 使用全部Web插件")
Common.LogInfo("使用全部Web插件")
} else {
Common.LogInfo(fmt.Sprintf("开始Web扫描 - 使用插件: %s", Common.ScanMode))
Common.LogInfo(fmt.Sprintf("使用插件: %s", Common.ScanMode))
}
executeScans(targetInfos, &ch, &wg)
default:
// 主机扫描模式
// 执行扫描
executeScans(targetInfos, ch, wg)
}
// 执行主机扫描
func executeHostScan(info Common.HostInfo, ch *chan struct{}, wg *sync.WaitGroup) {
if info.Host == "" {
Common.LogError("未指定扫描目标")
return
@ -113,19 +112,38 @@ func Scan(info Common.HostInfo) {
Common.LogError(fmt.Sprintf("解析主机错误: %v", err))
return
}
Common.LogInfo("开始主机扫描")
executeScan(hosts, info, &ch, &wg)
}
finishScan(&wg)
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,6 +151,30 @@ func executeScan(hosts []string, info Common.HostInfo, ch *chan struct{}, wg *sy
}
}
// 处理活跃端口
alivePorts := getAlivePorts(hosts)
if len(alivePorts) > 0 {
targetInfos = prepareTargetInfos(alivePorts, info)
}
}
// 添加 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)
@ -140,10 +182,11 @@ func executeScan(hosts []string, info Common.HostInfo, ch *chan struct{}, wg *sy
alivePorts = PortScan(hosts, Common.Ports, Common.Timeout)
Common.LogInfo(fmt.Sprintf("存活端口数量: %d", len(alivePorts)))
if Common.IsPortScan() {
return
return nil // 结束扫描
}
}
// 合并传入的端口信息
if len(Common.HostPort) > 0 {
alivePorts = append(alivePorts, Common.HostPort...)
alivePorts = Common.RemoveDuplicate(alivePorts)
@ -151,19 +194,17 @@ func executeScan(hosts []string, info Common.HostInfo, ch *chan struct{}, wg *sy
Common.LogInfo(fmt.Sprintf("存活端口数量: %d", len(alivePorts)))
}
targetInfos = prepareTargetInfos(alivePorts, info)
}
return alivePorts
}
// appendURLTargets 添加 URL 扫描目标
func appendURLTargets(targetInfos []Common.HostInfo, baseInfo Common.HostInfo) []Common.HostInfo {
for _, url := range Common.URLs {
urlInfo := info
urlInfo := baseInfo
urlInfo.Url = url
targetInfos = append(targetInfos, urlInfo)
}
if len(targetInfos) > 0 {
Common.LogInfo("开始漏洞扫描")
executeScans(targetInfos, ch, wg)
}
return targetInfos
}
// prepareTargetInfos 准备扫描目标信息
@ -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)
// 第一次遍历:计算任务数和收集要执行的插件
for _, target := range targets {
targetPort, _ := strconv.Atoi(target.Ports)
for _, pluginName := range pluginsToRun {
plugin, exists := Common.PluginManager[pluginName]
if !exists {
continue
}
// 收集扫描任务
func collectScanTasks(plugin Common.ScanPlugin, target Common.HostInfo, targetPort int, pluginName string, isSinglePlugin bool) (bool, []ScanTask) {
var tasks []ScanTask
taskAdded := false
// Web模式特殊处理
if WebScan {
actualTasks++
loadedPlugins = append(loadedPlugins, pluginName)
if WebScan || LocalScan || isSinglePlugin || len(plugin.Ports) == 0 || plugin.HasPort(targetPort) {
taskAdded = true
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
}
return taskAdded, tasks
}
// 单插件模式
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,
})
}
}
}
// 去重并输出实际加载的插件
// 获取去重后的插件列表
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()
}