package services

import (
	"context"
	"fmt"
	"time"

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

// SmbPlugin SMB弱密码检测插件
type SmbPlugin struct {
	name  string
	ports []int
}

// NewSmbPlugin 创建SMB插件
func NewSmbPlugin() *SmbPlugin {
	return &SmbPlugin{
		name:  "smb",
		ports: []int{445},
	}
}

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

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

// Scan 执行SMB扫描
func (p *SmbPlugin) 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: "smb",
			Error:   fmt.Errorf("SMB插件仅支持445端口"),
		}
	}

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

	// 生成测试凭据
	credentials := GenerateCredentials("smb")
	if len(credentials) == 0 {
		// SMB默认凭据
		credentials = []Credential{
			{Username: "", Password: ""},
			{Username: "administrator", Password: ""},
			{Username: "administrator", Password: "admin"},
			{Username: "administrator", Password: "password"},
			{Username: "administrator", Password: "123456"},
			{Username: "admin", Password: "admin"},
			{Username: "guest", Password: ""},
			{Username: "root", Password: ""},
			{Username: "root", Password: "root"},
		}
	}

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

		// 测试凭据
		if p.testCredential(ctx, info, cred) {
			// SMB认证成功
			var successMsg string
			if common.Domain != "" {
				successMsg = fmt.Sprintf("SMB %s 弱密码 %s\\%s:%s", target, common.Domain, cred.Username, cred.Password)
			} else {
				successMsg = fmt.Sprintf("SMB %s 弱密码 %s:%s", target, cred.Username, cred.Password)
			}
			common.LogSuccess(successMsg)

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

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


// testCredential 测试单个凭据
func (p *SmbPlugin) testCredential(ctx context.Context, info *common.HostInfo, cred Credential) bool {
	options := smb.Options{
		Host:        info.Host,
		Port:        445,
		User:        cred.Username,
		Password:    cred.Password,
		Domain:      common.Domain,
		Workstation: "",
	}

	// 设置超时
	timeoutCtx, cancel := context.WithTimeout(ctx, time.Duration(common.Timeout)*time.Second)
	defer cancel()

	// 在协程中执行连接测试
	resultChan := make(chan bool, 1)
	go func() {
		session, err := smb.NewSession(options, false)
		if err == nil {
			defer session.Close()
			resultChan <- session.IsAuthenticated
		} else {
			resultChan <- false
		}
	}()

	// 等待结果或超时
	select {
	case result := <-resultChan:
		return result
	case <-timeoutCtx.Done():
		return false
	case <-ctx.Done():
		return false
	}
}

// getShares 获取共享列表
func (p *SmbPlugin) getShares(ctx context.Context, info *common.HostInfo, creds Credential) []string {
	options := smb.Options{
		Host:        info.Host,
		Port:        445,
		User:        creds.Username,
		Password:    creds.Password,
		Domain:      common.Domain,
		Workstation: "",
	}

	session, err := smb.NewSession(options, false)
	if err != nil {
		return nil
	}
	defer session.Close()

	if !session.IsAuthenticated {
		return nil
	}

	// 简化实现，返回常见共享列表
	// 原SMB库可能不支持ListShares，这里使用模拟实现
	commonShares := []string{"ADMIN$", "C$", "IPC$", "Users", "Public"}
	return commonShares
}

// testShareAccess 测试共享访问
func (p *SmbPlugin) testShareAccess(ctx context.Context, info *common.HostInfo, creds Credential, shareName string) bool {
	options := smb.Options{
		Host:        info.Host,
		Port:        445,
		User:        creds.Username,
		Password:    creds.Password,
		Domain:      common.Domain,
		Workstation: "",
	}

	session, err := smb.NewSession(options, false)
	if err != nil {
		return false
	}
	defer session.Close()

	if !session.IsAuthenticated {
		return false
	}

	// 简化实现，假设管理员用户可以访问管理员共享
	if creds.Username == "administrator" && (shareName == "ADMIN$" || shareName == "C$") {
		return true
	}
	// 其他情况返回false
	return false
}

// identifyService 服务识别
func (p *SmbPlugin) identifyService(ctx context.Context, info *common.HostInfo) *ScanResult {
	if p.testCredential(ctx, info, Credential{Username: "", Password: ""}) {
		target := fmt.Sprintf("%s:%s", info.Host, info.Ports)
		banner := "SMB文件共享服务"
		common.LogSuccess(fmt.Sprintf("SMB %s %s", target, banner))

		return &ScanResult{
			Success: true,
			Service: "smb",
			Banner:  banner,
		}
	}

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

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