package services

import (
	"context"
	"fmt"
	"io/ioutil"
	"net"
	"regexp"
	"strings"
	"time"

	"golang.org/x/crypto/ssh"
	"github.com/shadow1ng/fscan/common"
	"github.com/shadow1ng/fscan/common/i18n"
	"github.com/shadow1ng/fscan/plugins"
)

// SSHPlugin SSH扫描和利用插件 - 单文件实现，包含真正的利用功能
type SSHPlugin struct {
	plugins.BasePlugin
}

// NewSSHPlugin 创建SSH插件
func NewSSHPlugin() *SSHPlugin {
	return &SSHPlugin{
		BasePlugin: plugins.NewBasePlugin("ssh"),
	}
}



// Scan 执行SSH扫描 - 支持密码和密钥认证
func (p *SSHPlugin) Scan(ctx context.Context, info *common.HostInfo) *ScanResult {
	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)

	// 如果指定了SSH密钥，优先使用密钥认证
	if common.SshKeyPath != "" {
		if result := p.scanWithKey(ctx, info); result != nil && result.Success {
			common.LogSuccess(i18n.GetText("ssh_key_auth_success", target, result.Username))
			return result
		}
	}

	// 如果禁用暴力破解，只做服务识别
	if common.DisableBrute {
		return p.identifyService(info)
	}

	// 生成测试凭据
	credentials := GenerateCredentials("ssh")
	if len(credentials) == 0 {
		// SSH默认凭据
		credentials = []Credential{
			{Username: "root", Password: ""},
			{Username: "root", Password: "root"},
			{Username: "root", Password: "toor"},
			{Username: "admin", Password: "admin"},
			{Username: "admin", Password: ""},
		}
	}

	// 逐个测试凭据
	for _, cred := range credentials {
		// 检查Context是否被取消
		select {
		case <-ctx.Done():
			return &ScanResult{
				Success: false,
				Service: "ssh",
				Error:   ctx.Err(),
			}
		default:
		}

		// 测试凭据
		if client := p.testCredential(ctx, info, cred); client != nil {
			// SSH认证成功，关闭连接
			client.Close()

			// 记录成功发现弱密码
			common.LogSuccess(i18n.GetText("ssh_pwd_auth_success", target, cred.Username, cred.Password))

			return &ScanResult{
				Success:  true,
				Service:  "ssh",
				Username: cred.Username,
				Password: cred.Password,
			}
		}
	}

	// 所有凭据都失败
	return &ScanResult{
		Success: false,
		Service: "ssh",
		Error:   fmt.Errorf("未发现弱密码"),
	}
}


// scanWithKey 使用SSH私钥扫描
func (p *SSHPlugin) scanWithKey(ctx context.Context, info *common.HostInfo) *ScanResult {
	// 读取私钥文件
	keyData, err := ioutil.ReadFile(common.SshKeyPath)
	if err != nil {
		common.LogError(i18n.GetText("ssh_key_read_failed", err))
		return nil
	}

	// 常见的SSH用户名
	usernames := common.Userdict["ssh"]
	if len(usernames) == 0 {
		usernames = []string{"root", "admin", "ubuntu", "centos", "user", "git", "www-data"}
	}

	// 逐个测试用户名
	for _, username := range usernames {
		cred := Credential{
			Username: username,
			KeyData:  keyData,
		}

		if client := p.testCredential(ctx, info, cred); client != nil {
			client.Close()
			return &ScanResult{
				Success:  true,
				Service:  "ssh",
				Username: username,
			}
		}
	}

	return nil
}

// testCredential 测试单个凭据 - 返回SSH客户端或nil
func (p *SSHPlugin) testCredential(ctx context.Context, info *common.HostInfo, cred Credential) *ssh.Client {
	// 创建SSH配置
	config := &ssh.ClientConfig{
		User:            cred.Username,
		Timeout:         time.Duration(common.Timeout) * time.Second,
		HostKeyCallback: ssh.InsecureIgnoreHostKey(), // 忽略主机密钥验证
	}

	// 设置认证方法
	if len(cred.KeyData) > 0 {
		// 私钥认证
		signer, err := ssh.ParsePrivateKey(cred.KeyData)
		if err != nil {
			return nil
		}
		config.Auth = []ssh.AuthMethod{ssh.PublicKeys(signer)}
	} else {
		// 密码认证
		config.Auth = []ssh.AuthMethod{ssh.Password(cred.Password)}
	}

	// 建立连接
	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)
	
	// 使用Context控制超时
	type sshResult struct {
		client *ssh.Client
		err    error
	}

	// 检查发包限制
	if canSend, reason := common.CanSendPacket(); !canSend {
		common.LogError(fmt.Sprintf("SSH连接 %s 受限: %s", target, reason))
		return nil
	}

	resultChan := make(chan sshResult, 1)
	go func() {
		client, err := ssh.Dial("tcp", target, config)
		resultChan <- sshResult{client: client, err: err}
	}()

	select {
	case result := <-resultChan:
		if result.err != nil {
			// 计数TCP连接失败包
			common.IncrementTCPFailedPacketCount()
			return nil
		}
		// 计数TCP连接成功包
		common.IncrementTCPSuccessPacketCount()
		return result.client
	case <-ctx.Done():
		return nil
	}
}

// executeCommand 在SSH连接上执行命令
func (p *SSHPlugin) executeCommand(client *ssh.Client, cmd string) (string, error) {
	session, err := client.NewSession()
	if err != nil {
		return "", err
	}
	defer session.Close()

	output, err := session.CombinedOutput(cmd)
	return string(output), err
}

// identifyService 服务识别 - 不进行暴力破解时的功能
func (p *SSHPlugin) identifyService(info *common.HostInfo) *ScanResult {
	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)

	// 检查发包限制
	if canSend, reason := common.CanSendPacket(); !canSend {
		common.LogError(fmt.Sprintf("SSH服务识别 %s 受限: %s", target, reason))
		return &ScanResult{
			Success: false,
			Service: "ssh",
			Error:   fmt.Errorf("发包受限: %s", reason),
		}
	}

	// 尝试连接获取SSH Banner
	conn, err := common.WrapperTcpWithTimeout("tcp", target, time.Duration(common.Timeout)*time.Second)
	if err != nil {
		return &ScanResult{
			Success: false,
			Service: "ssh",
			Error:   err,
		}
	}
	defer conn.Close()

	// 读取SSH Banner
	if banner := p.readSSHBanner(conn); banner != "" {
		common.LogSuccess(i18n.GetText("ssh_service_identified", target, banner))
		return &ScanResult{
			Success: true,
			Service: "ssh",
			Banner:  banner,
		}
	}

	return &ScanResult{
		Success: false,
		Service: "ssh",
		Error:   fmt.Errorf("无法识别为SSH服务"),
	}
}

// readSSHBanner 读取SSH服务器Banner
func (p *SSHPlugin) readSSHBanner(conn net.Conn) string {
	// 设置读取超时
	conn.SetReadDeadline(time.Now().Add(time.Duration(common.Timeout) * time.Second))

	// 读取Banner
	banner := make([]byte, 256)
	n, err := conn.Read(banner)
	if err != nil || n < 4 {
		return ""
	}

	bannerStr := strings.TrimSpace(string(banner[:n]))

	// 检查SSH协议标识
	if strings.HasPrefix(bannerStr, "SSH-") {
		// 使用正则表达式解析版本信息
		if matched := regexp.MustCompile(`SSH-([0-9.]+)-(.+)`).FindStringSubmatch(bannerStr); len(matched) >= 3 {
			protocolVersion := matched[1]
			serverVersion := matched[2]
			return fmt.Sprintf("SSH %s (%s)", protocolVersion, serverVersion)
		}
		return fmt.Sprintf("SSH服务: %s", bannerStr)
	}

	return ""
}

// init 自动注册插件
func init() {
	// 使用高效注册方式：直接传递端口信息，避免实例创建
	RegisterPluginWithPorts("ssh", func() Plugin {
		return NewSSHPlugin()
	}, []int{22, 2222, 2200, 22222})
}