package ssh

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

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

// SSHConnector SSH连接器实现
type SSHConnector struct {
	host    string
	port    string
}

// NewSSHConnector 创建SSH连接器
func NewSSHConnector() *SSHConnector {
	return &SSHConnector{
		// 移除timeout字段，统一使用Context超时
	}
}

// Connect 连接到SSH服务
func (c *SSHConnector) Connect(ctx context.Context, info *common.HostInfo) (interface{}, error) {
	// 保存主机和端口信息
	c.host = info.Host
	c.port = info.Ports
	
	// SSH连接在认证时才真正建立，这里返回配置信息
	// 移除Timeout设置，统一使用Context超时控制
	config := &ssh.ClientConfig{
		Timeout:         0, // 禁用SSH库内部超时，使用Context控制
		HostKeyCallback: ssh.InsecureIgnoreHostKey(), // 忽略主机密钥验证
	}
	
	return config, nil
}

// Authenticate 认证
func (c *SSHConnector) Authenticate(ctx context.Context, conn interface{}, cred *base.Credential) error {
	config, ok := conn.(*ssh.ClientConfig)
	if !ok {
		return fmt.Errorf("无效的连接类型")
	}
	
	// 直接使用传入的Context，它已经包含了正确的超时设置
	// 创建配置副本并设置认证方法
	authConfig := *config
	
	if cred.KeyData != nil && len(cred.KeyData) > 0 {
		// 密钥认证
		signer, err := ssh.ParsePrivateKey(cred.KeyData)
		if err != nil {
			return fmt.Errorf("解析私钥失败: %v", err)
		}
		authConfig.User = cred.Username
		authConfig.Auth = []ssh.AuthMethod{ssh.PublicKeys(signer)}
	} else {
		// 密码认证
		authConfig.User = cred.Username
		authConfig.Auth = []ssh.AuthMethod{ssh.Password(cred.Password)}
	}
	
	// 构建目标地址
	target := fmt.Sprintf("%s:%s", c.host, c.port)
	
	// 使用Context控制超时的SSH连接
	type sshResult struct {
		client *ssh.Client
		err    error
	}
	
	resultChan := make(chan sshResult, 1)
	
	go func() {
		client, err := ssh.Dial("tcp", target, &authConfig)
		resultChan <- sshResult{client: client, err: err}
	}()
	
	// 等待结果或超时
	select {
	case result := <-resultChan:
		if result.client != nil {
			defer result.client.Close()
		}
		if result.err != nil {
			return fmt.Errorf("SSH认证失败: %v", result.err)
		}
		return nil
	case <-ctx.Done():
		return fmt.Errorf("SSH连接超时: %v", ctx.Err())
	}
}

// Close 关闭连接
func (c *SSHConnector) Close(conn interface{}) error {
	// SSH配置无需关闭
	return nil
}

// SSHPlugin SSH插件实现
type SSHPlugin struct {
	*base.ServicePlugin
	exploiter *SSHExploiter
}

// NewSSHPlugin 创建SSH插件
func NewSSHPlugin() *SSHPlugin {
	// 插件元数据
	metadata := &base.PluginMetadata{
		Name:        "ssh",
		Version:     "2.0.0",
		Author:      "fscan-team",
		Description: "SSH服务扫描和利用插件",
		Category:    "service",
		Ports:       []int{22, 2222, 2200, 22222}, // 添加常见的SSH替代端口
		Protocols:   []string{"tcp"},
		Tags:        []string{"ssh", "bruteforce", "remote_access"},
	}
	
	// 创建连接器和服务插件
	connector := NewSSHConnector()
	servicePlugin := base.NewServicePlugin(metadata, connector)
	
	// 创建SSH插件
	plugin := &SSHPlugin{
		ServicePlugin: servicePlugin,
		exploiter:     NewSSHExploiter(),
	}
	
	// 设置能力
	plugin.SetCapabilities([]base.Capability{
		base.CapWeakPassword,
		base.CapCommandExecution,
		base.CapDataExtraction,
	})
	
	return plugin
}

// Scan 重写扫描方法以支持密钥认证和自动利用
func (p *SSHPlugin) Scan(ctx context.Context, info *common.HostInfo) (*base.ScanResult, error) {
	// 如果指定了SSH密钥，优先使用密钥认证
	if common.SshKeyPath != "" {
		result := p.scanWithKey(ctx, info)
		if result != nil && result.Success {
			target := fmt.Sprintf("%s:%s", info.Host, info.Ports)
			common.LogSuccess(i18n.GetText("ssh_key_auth_success", target, result.Credentials[0].Username))
			
			return result, nil
		}
	}
	
	// 执行基础的密码扫描
	if common.DisableBrute {
		return p.performServiceIdentification(ctx, info)
	}
	
	result, err := p.ServicePlugin.Scan(ctx, info)
	if err != nil || !result.Success {
		return result, err
	}
	
	// 记录成功的弱密码发现
	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)
	cred := result.Credentials[0]
	common.LogSuccess(i18n.GetText("ssh_pwd_auth_success", target, cred.Username, cred.Password))
	
	return result, nil
}

// scanWithKey 使用密钥扫描
func (p *SSHPlugin) scanWithKey(ctx context.Context, info *common.HostInfo) *base.ScanResult {
	// 读取私钥
	keyData, err := ioutil.ReadFile(common.SshKeyPath)
	if err != nil {
		common.LogError(i18n.GetText("ssh_key_read_failed", err))
		return nil
	}
	
	// 尝试不同的用户名
	usernames := common.Userdict["ssh"]
	if len(usernames) == 0 {
		usernames = []string{"root", "admin", "ubuntu", "centos", "user"}
	}
	
	for _, username := range usernames {
		cred := &base.Credential{
			Username: username,
			KeyData:  keyData,
		}
		
		result, err := p.ScanCredential(ctx, info, cred)
		if err == nil && result.Success {
			return result
		}
	}
	
	return nil
}

// generateCredentials 重写凭据生成方法
func (p *SSHPlugin) generateCredentials() []*base.Credential {
	// 获取SSH专用的用户名字典
	usernames := common.Userdict["ssh"]
	if len(usernames) == 0 {
		// 默认SSH用户名
		usernames = []string{"root", "admin", "ubuntu", "centos", "user", "test"}
	}
	
	return base.GenerateCredentials(usernames, common.Passwords)
}


// Exploit 使用exploiter执行利用
func (p *SSHPlugin) Exploit(ctx context.Context, info *common.HostInfo, creds *base.Credential) (*base.ExploitResult, error) {
	return p.exploiter.Exploit(ctx, info, creds)
}

// GetExploitMethods 获取利用方法
func (p *SSHPlugin) GetExploitMethods() []base.ExploitMethod {
	return p.exploiter.GetExploitMethods()
}

// IsExploitSupported 检查利用支持
func (p *SSHPlugin) IsExploitSupported(method base.ExploitType) bool {
	return p.exploiter.IsExploitSupported(method)
}

// performServiceIdentification 执行SSH服务识别（-nobr模式）
func (p *SSHPlugin) performServiceIdentification(ctx context.Context, info *common.HostInfo) (*base.ScanResult, error) {
	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)
	
	// 尝试连接到SSH服务获取Banner
	conn, err := common.WrapperTcpWithTimeout("tcp", target, time.Duration(common.Timeout)*time.Second)
	if err != nil {
		return &base.ScanResult{
			Success: false,
			Error:   err,
		}, nil
	}
	defer conn.Close()
	
	// 读取SSH Banner
	sshInfo, isSSH := p.identifySSHService(conn)
	if isSSH {
		// 记录服务识别成功
		common.LogSuccess(i18n.GetText("ssh_service_identified", target, sshInfo))
		
		return &base.ScanResult{
			Success: true,
			Service: "SSH",
			Banner:  sshInfo,
			Extra: map[string]interface{}{
				"service": "SSH",
				"port":    info.Ports,
				"info":    sshInfo,
			},
		}, nil
	}
	
	// 如果无法识别为SSH，返回失败
	return &base.ScanResult{
		Success: false,
		Error:   fmt.Errorf("无法识别为SSH服务"),
	}, nil
}

// identifySSHService 通过Banner识别SSH服务
func (p *SSHPlugin) identifySSHService(conn net.Conn) (string, bool) {
	// 设置读取超时
	conn.SetReadDeadline(time.Now().Add(time.Duration(common.Timeout) * time.Second))
	
	// SSH服务器在连接后会发送Banner
	banner := make([]byte, 512)
	n, err := conn.Read(banner)
	if err != nil || n < 4 {
		return "", false
	}
	
	bannerStr := strings.TrimSpace(string(banner[:n]))
	
	// 检查SSH协议标识
	if strings.HasPrefix(bannerStr, "SSH-") {
		// 提取SSH版本信息
		parts := strings.Fields(bannerStr)
		if len(parts) > 0 {
			// 提取协议版本和服务器标识
			versionPart := parts[0]
			serverInfo := ""
			if len(parts) > 1 {
				serverInfo = strings.Join(parts[1:], " ")
			}
			
			// 使用正则表达式提取更详细信息
			if matched := regexp.MustCompile(`SSH-([0-9.]+)-(.+)`).FindStringSubmatch(versionPart); len(matched) >= 3 {
				protocolVersion := matched[1]
				serverVersion := matched[2]
				if serverInfo != "" {
					return fmt.Sprintf("SSH %s (%s) %s", protocolVersion, serverVersion, serverInfo), true
				}
				return fmt.Sprintf("SSH %s (%s)", protocolVersion, serverVersion), true
			}
			
			return fmt.Sprintf("SSH服务: %s", bannerStr), true
		}
		return "SSH服务", true
	}
	
	return "", false
}

// =============================================================================
// 插件注册
// =============================================================================

// RegisterSSHPlugin 注册SSH插件
func RegisterSSHPlugin() {
	factory := base.NewSimplePluginFactory(
		&base.PluginMetadata{
			Name:        "ssh",
			Version:     "2.0.0",
			Author:      "fscan-team",
			Description: "SSH服务扫描和利用插件",
			Category:    "service",
			Ports:       []int{22, 2222, 2200, 22222}, // 添加常见的SSH替代端口
			Protocols:   []string{"tcp"},
			Tags:        []string{"ssh", "bruteforce", "remote_access"},
		},
		func() base.Plugin {
			return NewSSHPlugin()
		},
	)
	
	base.GlobalPluginRegistry.Register("ssh", factory)
}

// 自动注册
func init() {
	RegisterSSHPlugin()
}