package services

import (
	"bytes"
	"context"
	"fmt"
	"strings"
	"time"

	"github.com/shadow1ng/fscan/common"
)

// SmbGhostPlugin CVE-2020-0796 SMB Ghost漏洞检测插件
type SmbGhostPlugin struct {
	name  string
	ports []int
}

// SMB Ghost 检测数据包
const smbGhostPacket = "\x00" + // session
	"\x00\x00\xc0" + // length

	"\xfeSMB@\x00" + // protocol

	//[MS-SMB2]: SMB2 NEGOTIATE Request
	//https://docs.microsoft.com/en-us/openspecs/windows_protocols/ms-smb2/e14db7ff-763a-4263-8b10-0c3944f52fc5

	"\x00\x00" +
	"\x00\x00" +
	"\x00\x00" +
	"\x00\x00" +
	"\x1f\x00" +
	"\x00\x00\x00\x00" +
	"\x00\x00\x00\x00" +
	"\x00\x00\x00\x00" +
	"\x00\x00\x00\x00" +
	"\x00\x00\x00\x00" +
	"\x00\x00\x00\x00" +
	"\x00\x00\x00\x00" +
	"\x00\x00\x00\x00" +
	"\x00\x00\x00\x00" +
	"\x00\x00\x00\x00" +
	"\x00\x00\x00\x00" +
	"\x00\x00\x00\x00" +

	// [MS-SMB2]: SMB2 NEGOTIATE_CONTEXT
	// https://docs.microsoft.com/en-us/openspecs/windows_protocols/ms-smb2/15332256-522e-4a53-8cd7-0bd17678a2f7

	"$\x00" +
	"\x08\x00" +
	"\x01\x00" +
	"\x00\x00" +
	"\x7f\x00\x00\x00" +
	"\x00\x00\x00\x00" +
	"\x00\x00\x00\x00" +
	"\x00\x00\x00\x00" +
	"\x00\x00\x00\x00" +
	"x\x00" +
	"\x00\x00" +
	"\x02\x00" +
	"\x00\x00" +
	"\x02\x02" +
	"\x10\x02" +
	"\x22\x02" +
	"$\x02" +
	"\x00\x03" +
	"\x02\x03" +
	"\x10\x03" +
	"\x11\x03" +
	"\x00\x00\x00\x00" +

	// [MS-SMB2]: SMB2_PREAUTH_INTEGRITY_CAPABILITIES
	// https://docs.microsoft.com/en-us/openspecs/windows_protocols/ms-smb2/5a07bd66-4734-4af8-abcf-5a44ff7ee0e5

	"\x01\x00" +
	"&\x00" +
	"\x00\x00\x00\x00" +
	"\x01\x00" +
	"\x20\x00" +
	"\x01\x00" +
	"\x00\x00\x00\x00" +
	"\x00\x00\x00\x00" +
	"\x00\x00\x00\x00" +
	"\x00\x00\x00\x00" +
	"\x00\x00\x00\x00" +
	"\x00\x00\x00\x00" +
	"\x00\x00\x00\x00" +
	"\x00\x00\x00\x00" +
	"\x00\x00" +

	// [MS-SMB2]: SMB2_COMPRESSION_CAPABILITIES
	// https://docs.microsoft.com/en-us/openspecs/windows_protocols/ms-smb2/78e0c942-ab41-472b-b117-4a95ebe88271

	"\x03\x00" +
	"\x0e\x00" +
	"\x00\x00\x00\x00" +
	"\x01\x00" + //CompressionAlgorithmCount
	"\x00\x00" +
	"\x01\x00\x00\x00" +
	"\x01\x00" + //LZNT1
	"\x00\x00" +
	"\x00\x00\x00\x00"

// NewSmbGhostPlugin 创建SMB Ghost插件
func NewSmbGhostPlugin() *SmbGhostPlugin {
	return &SmbGhostPlugin{
		name:  "smbghost",
		ports: []int{445},
	}
}

// GetName 实现Plugin接口
func (p *SmbGhostPlugin) Name() string {
	return p.name
}

// GetPorts 实现Plugin接口
func (p *SmbGhostPlugin) GetPorts() []int {
	return p.ports
}

// Scan 执行SMB Ghost漏洞检测
func (p *SmbGhostPlugin) Scan(ctx context.Context, info *common.HostInfo) *ScanResult {
	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)

	// 检查端口
	if info.Ports != "445" {
		return &ScanResult{
			Success: false,
			Service: "smbghost",
			Error:   fmt.Errorf("SMB Ghost插件仅支持445端口"),
		}
	}

	// 执行漏洞检测
	vulnerable, err := p.detectVulnerability(ctx, info)
	if err != nil {
		return &ScanResult{
			Success: false,
			Service: "smbghost",
			Error:   fmt.Errorf("检测失败: %v", err),
		}
	}

	if vulnerable {
		common.LogSuccess(fmt.Sprintf("SMB Ghost %s CVE-2020-0796 漏洞存在", target))
		return &ScanResult{
			Success: true,
			Service: "smbghost",
			Banner:  "CVE-2020-0796 SMB Ghost漏洞",
		}
	}

	return &ScanResult{
		Success: false,
		Service: "smbghost",
		Error:   fmt.Errorf("未发现CVE-2020-0796漏洞"),
	}
}


// detectVulnerability 检测SMB Ghost漏洞
func (p *SmbGhostPlugin) detectVulnerability(ctx context.Context, info *common.HostInfo) (bool, error) {
	addr := fmt.Sprintf("%s:445", info.Host)
	timeout := time.Duration(common.Timeout) * time.Second

	// 建立TCP连接
	conn, err := common.WrapperTcpWithTimeout("tcp", addr, timeout)
	if err != nil {
		return false, fmt.Errorf("连接失败: %v", err)
	}
	defer conn.Close()

	// 设置连接超时
	if err = conn.SetDeadline(time.Now().Add(timeout)); err != nil {
		return false, fmt.Errorf("设置超时失败: %v", err)
	}

	// 发送SMB Ghost检测数据包
	if _, err = conn.Write([]byte(smbGhostPacket)); err != nil {
		return false, fmt.Errorf("发送数据包失败: %v", err)
	}

	// 读取响应
	buff := make([]byte, 1024)
	n, err := conn.Read(buff)
	if err != nil || n == 0 {
		return false, fmt.Errorf("读取响应失败: %v", err)
	}

	// 分析响应，检测CVE-2020-0796特征
	// 检查条件：
	// 1. 响应包含"Public"字符串
	// 2. 响应长度大于等于76字节
	// 3. 特征字节匹配 (0x11,0x03) 和 (0x02,0x00)
	if bytes.Contains(buff[:n], []byte("Public")) &&
		len(buff[:n]) >= 76 &&
		bytes.Equal(buff[72:74], []byte{0x11, 0x03}) &&
		bytes.Equal(buff[74:76], []byte{0x02, 0x00}) {
		return true, nil
	}

	return false, nil
}

// gatherSystemInfo 收集系统信息
func (p *SmbGhostPlugin) gatherSystemInfo(ctx context.Context, info *common.HostInfo) string {
	addr := fmt.Sprintf("%s:445", info.Host)
	timeout := time.Duration(common.Timeout) * time.Second

	conn, err := common.WrapperTcpWithTimeout("tcp", addr, timeout)
	if err != nil {
		return ""
	}
	defer conn.Close()

	// 发送标准SMB协商请求获取系统信息
	negotiateReq := []byte{
		0x00, 0x00, 0x00, 0x85,
		0xfe, 0x53, 0x4d, 0x42, 0x40, 0x00, 0x00, 0x00,
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		0x24, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00,
		0x00, 0x00, 0x00, 0x00, 0x02, 0x02, 0x10, 0x02,
		0x22, 0x02, 0x24, 0x02, 0x00, 0x03, 0x02, 0x03,
		0x10, 0x03, 0x11, 0x03,
	}

	conn.SetDeadline(time.Now().Add(timeout))
	conn.Write(negotiateReq)
	
	buff := make([]byte, 1024)
	n, err := conn.Read(buff)
	if err != nil || n < 65 {
		return ""
	}

	var sysInfo strings.Builder
	sysInfo.WriteString("SMB服务检测成功\n")
	
	// 解析SMB版本
	if n >= 72 {
		dialect := uint16(buff[70]) | uint16(buff[71])<<8
		switch dialect {
		case 0x0311:
			sysInfo.WriteString("SMB版本: 3.1.1 (存在CVE-2020-0796风险)\n")
		case 0x0302:
			sysInfo.WriteString("SMB版本: 3.0.2\n")
		case 0x0300:
			sysInfo.WriteString("SMB版本: 3.0.0\n")
		default:
			sysInfo.WriteString(fmt.Sprintf("SMB版本: 未知 (0x%04x)\n", dialect))
		}
	}

	return sysInfo.String()
}

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