fscan/Common/Core/Plugin.go

package core

import (
	"fmt"
	"sync"
)

/*
Plugin.go - 插件系统管理

整合Types.go中的插件系统，提供统一的插件注册和管理机制。
*/

// =============================================================================
// 核心数据结构 (从Types.go迁移)
// =============================================================================

// HostInfo 主机信息结构
type HostInfo struct {
	Host    string   // 主机地址
	Ports   string   // 端口范围
	Url     string   // URL地址
	Infostr []string // 附加信息
}

// =============================================================================
// 插件类型常量
// =============================================================================

const (
	PluginTypeService = "service" // 服务类型插件
	PluginTypeWeb     = "web"     // Web类型插件
	PluginTypeLocal   = "local"   // 本地类型插件
	PluginTypeBrute   = "brute"   // 暴力破解插件
	PluginTypePoc     = "poc"     // POC验证插件
	PluginTypeScan    = "scan"    // 扫描探测插件
)

// =============================================================================
// 插件定义和管理
// =============================================================================

// ScanPlugin 定义扫描插件的结构
type ScanPlugin struct {
	Name        string                // 插件名称
	Version     string                // 插件版本
	Description string                // 插件描述
	Author      string                // 插件作者
	Ports       []int                 // 适用端口
	Types       []string              // 插件类型标签，一个插件可以有多个类型
	Priority    int                   // 插件优先级(数字越小优先级越高)
	Enabled     bool                  // 是否启用
	ScanFunc    func(*HostInfo) error // 扫描函数
}

// PluginManager 插件管理器
type PluginManager struct {
	mu      sync.RWMutex
	plugins map[string]*ScanPlugin
	types   map[string][]*ScanPlugin // 按类型索引的插件
	ports   map[int][]*ScanPlugin    // 按端口索引的插件
}

// 全局插件管理器实例
var globalPluginManager = NewPluginManager()

// NewPluginManager 创建新的插件管理器
func NewPluginManager() *PluginManager {
	return &PluginManager{
		plugins: make(map[string]*ScanPlugin),
		types:   make(map[string][]*ScanPlugin),
		ports:   make(map[int][]*ScanPlugin),
	}
}

// =============================================================================
// 插件基础方法
// =============================================================================

// HasType 检查插件是否具有指定类型
func (p *ScanPlugin) HasType(typeName string) bool {
	for _, t := range p.Types {
		if t == typeName {
			return true
		}
	}
	return false
}

// HasPort 检查插件是否支持指定端口
func (p *ScanPlugin) HasPort(port int) bool {
	// 如果没有指定端口列表，表示支持所有端口
	if len(p.Ports) == 0 {
		return true
	}

	// 检查端口是否在支持列表中
	for _, supportedPort := range p.Ports {
		if port == supportedPort {
			return true
		}
	}
	return false
}

// IsEnabled 检查插件是否启用
func (p *ScanPlugin) IsEnabled() bool {
	return p.Enabled
}

// GetInfo 获取插件基本信息
func (p *ScanPlugin) GetInfo() map[string]interface{} {
	return map[string]interface{}{
		"name":        p.Name,
		"version":     p.Version,
		"description": p.Description,
		"author":      p.Author,
		"types":       p.Types,
		"ports":       p.Ports,
		"priority":    p.Priority,
		"enabled":     p.Enabled,
	}
}

// =============================================================================
// 插件管理器方法
// =============================================================================

// RegisterPlugin 注册插件
func (pm *PluginManager) RegisterPlugin(plugin *ScanPlugin) error {
	if plugin == nil {
		return fmt.Errorf("plugin cannot be nil")
	}
	if plugin.Name == "" {
		return fmt.Errorf("plugin name cannot be empty")
	}
	if plugin.ScanFunc == nil {
		return fmt.Errorf("plugin scan function cannot be nil")
	}

	pm.mu.Lock()
	defer pm.mu.Unlock()

	// 检查插件是否已存在
	if _, exists := pm.plugins[plugin.Name]; exists {
		return fmt.Errorf("plugin %s already registered", plugin.Name)
	}

	// 注册插件
	pm.plugins[plugin.Name] = plugin

	// 按类型索引
	for _, pluginType := range plugin.Types {
		pm.types[pluginType] = append(pm.types[pluginType], plugin)
	}

	// 按端口索引
	for _, port := range plugin.Ports {
		pm.ports[port] = append(pm.ports[port], plugin)
	}

	return nil
}

// GetPlugin 获取指定名称的插件
func (pm *PluginManager) GetPlugin(name string) (*ScanPlugin, bool) {
	pm.mu.RLock()
	defer pm.mu.RUnlock()
	plugin, exists := pm.plugins[name]
	return plugin, exists
}

// GetPluginsByType 获取指定类型的所有插件
func (pm *PluginManager) GetPluginsByType(pluginType string) []*ScanPlugin {
	pm.mu.RLock()
	defer pm.mu.RUnlock()
	
	plugins := make([]*ScanPlugin, 0)
	if typePlugins, exists := pm.types[pluginType]; exists {
		for _, plugin := range typePlugins {
			if plugin.IsEnabled() {
				plugins = append(plugins, plugin)
			}
		}
	}
	return plugins
}

// GetPluginsByPort 获取支持指定端口的所有插件
func (pm *PluginManager) GetPluginsByPort(port int) []*ScanPlugin {
	pm.mu.RLock()
	defer pm.mu.RUnlock()
	
	plugins := make([]*ScanPlugin, 0)
	
	// 检查按端口索引的插件
	if portPlugins, exists := pm.ports[port]; exists {
		for _, plugin := range portPlugins {
			if plugin.IsEnabled() {
				plugins = append(plugins, plugin)
			}
		}
	}
	
	// 检查支持所有端口的插件
	for _, plugin := range pm.plugins {
		if len(plugin.Ports) == 0 && plugin.IsEnabled() {
			plugins = append(plugins, plugin)
		}
	}
	
	return plugins
}

// GetAllPlugins 获取所有插件
func (pm *PluginManager) GetAllPlugins() map[string]*ScanPlugin {
	pm.mu.RLock()
	defer pm.mu.RUnlock()
	
	// 创建副本避免并发修改
	result := make(map[string]*ScanPlugin)
	for name, plugin := range pm.plugins {
		result[name] = plugin
	}
	return result
}

// EnablePlugin 启用插件
func (pm *PluginManager) EnablePlugin(name string) error {
	pm.mu.Lock()
	defer pm.mu.Unlock()
	
	plugin, exists := pm.plugins[name]
	if !exists {
		return fmt.Errorf("plugin %s not found", name)
	}
	
	plugin.Enabled = true
	return nil
}

// DisablePlugin 禁用插件
func (pm *PluginManager) DisablePlugin(name string) error {
	pm.mu.Lock()
	defer pm.mu.Unlock()
	
	plugin, exists := pm.plugins[name]
	if !exists {
		return fmt.Errorf("plugin %s not found", name)
	}
	
	plugin.Enabled = false
	return nil
}

// UnregisterPlugin 注销插件
func (pm *PluginManager) UnregisterPlugin(name string) error {
	pm.mu.Lock()
	defer pm.mu.Unlock()
	
	plugin, exists := pm.plugins[name]
	if !exists {
		return fmt.Errorf("plugin %s not found", name)
	}
	
	// 从主索引删除
	delete(pm.plugins, name)
	
	// 从类型索引删除
	for _, pluginType := range plugin.Types {
		if typePlugins, exists := pm.types[pluginType]; exists {
			for i, p := range typePlugins {
				if p.Name == name {
					pm.types[pluginType] = append(typePlugins[:i], typePlugins[i+1:]...)
					break
				}
			}
		}
	}
	
	// 从端口索引删除
	for _, port := range plugin.Ports {
		if portPlugins, exists := pm.ports[port]; exists {
			for i, p := range portPlugins {
				if p.Name == name {
					pm.ports[port] = append(portPlugins[:i], portPlugins[i+1:]...)
					break
				}
			}
		}
	}
	
	return nil
}

// GetPluginCount 获取插件总数
func (pm *PluginManager) GetPluginCount() int {
	pm.mu.RLock()
	defer pm.mu.RUnlock()
	return len(pm.plugins)
}

// GetEnabledPluginCount 获取启用的插件数量
func (pm *PluginManager) GetEnabledPluginCount() int {
	pm.mu.RLock()
	defer pm.mu.RUnlock()
	
	count := 0
	for _, plugin := range pm.plugins {
		if plugin.IsEnabled() {
			count++
		}
	}
	return count
}

// =============================================================================
// 全局插件管理函数 (保持向后兼容)
// =============================================================================

// RegisterPlugin 注册插件到全局管理器
func RegisterPlugin(name string, plugin ScanPlugin) error {
	// 转换为新的插件结构
	newPlugin := &ScanPlugin{
		Name:     name,
		Ports:    plugin.Ports,
		Types:    plugin.Types,
		Enabled:  true,
		ScanFunc: plugin.ScanFunc,
	}
	
	// 注册到新的插件管理器
	err := globalPluginManager.RegisterPlugin(newPlugin)
	if err != nil {
		return err
	}
	
	// 同时更新Legacy管理器以保持向后兼容
	LegacyPluginManager[name] = plugin
	
	return nil
}

// GetGlobalPluginManager 获取全局插件管理器
func GetGlobalPluginManager() *PluginManager {
	return globalPluginManager
}

// 向后兼容的全局变量 (已废弃，建议使用PluginManager)
var LegacyPluginManager = make(map[string]ScanPlugin)