package smb2

import (
	"context"
	"fmt"
	"os"
	"time"
	
	"github.com/shadow1ng/fscan/plugins/legacy/smb/common"
	fscanCommon "github.com/shadow1ng/fscan/common"
	"github.com/hirochachacha/go-smb2"
)

// Smb2Connector SMB2连接器实现
type Smb2Connector struct{}

// NewSmb2Connector 创建SMB2连接器
func NewSmb2Connector() *Smb2Connector {
	return &Smb2Connector{}
}

// Connect 建立SMB2连接并进行认证
func (c *Smb2Connector) Connect(ctx context.Context, target *common.TargetInfo, 
	cred *common.Credential) (*common.ConnectionResult, error) {
	
	// 建立TCP连接，使用socks代理支持
	conn, err := fscanCommon.WrapperTcpWithTimeout("tcp", 
		fmt.Sprintf("%s:%d", target.Host, target.Port), 
		time.Duration(fscanCommon.Timeout)*time.Second)
	if err != nil {
		return nil, fmt.Errorf("连接失败: %v", err)
	}
	defer conn.Close()
	
	// 配置NTLM认证
	initiator := smb2.NTLMInitiator{
		User:   cred.Username,
		Domain: target.Domain,
	}
	
	// 设置认证方式(哈希或密码)
	if cred.IsHash && len(cred.Hash) > 0 {
		initiator.Hash = cred.Hash
	} else {
		initiator.Password = cred.Password
	}
	
	// 创建SMB2会话
	dialer := &smb2.Dialer{
		Initiator: &initiator,
	}
	
	// 建立会话
	session, err := dialer.Dial(conn)
	if err != nil {
		classifiedError := common.ClassifySmbError(err)
		return &common.ConnectionResult{
			Success: false,
			Error:   classifiedError,
		}, nil
	}
	defer session.Logoff()
	
	// 检查上下文是否已取消
	select {
	case <-ctx.Done():
		return &common.ConnectionResult{
			Success: false,
			Error:   ctx.Err(),
		}, nil
	default:
	}
	
	// 获取共享列表
	shares, err := session.ListSharenames()
	if err != nil {
		return &common.ConnectionResult{
			Success: false,
			Error:   fmt.Errorf("获取共享列表失败: %v", err),
		}, nil
	}
	
	// 检查上下文是否已取消
	select {
	case <-ctx.Done():
		return &common.ConnectionResult{
			Success: false,
			Error:   ctx.Err(),
			Shares:  shares,
		}, nil
	default:
	}
	
	// 尝试验证管理员权限
	hasAdminAccess := c.validateAdminAccess(ctx, session)
	
	return &common.ConnectionResult{
		Success:       true,
		Shares:        shares,
		HasAdminAccess: hasAdminAccess,
	}, nil
}

// GetProtocolName 获取协议名称
func (c *Smb2Connector) GetProtocolName() string {
	return "SMB2"
}

// GetDefaultPort 获取默认端口
func (c *Smb2Connector) GetDefaultPort() int {
	return 445
}

// validateAdminAccess 验证管理员权限
func (c *Smb2Connector) validateAdminAccess(ctx context.Context, session *smb2.Session) bool {
	// 检查上下文
	select {
	case <-ctx.Done():
		return false
	default:
	}
	
	// 尝试挂载C$共享
	fs, err := session.Mount("C$")
	if err != nil {
		return false
	}
	defer fs.Umount()
	
	// 检查上下文
	select {
	case <-ctx.Done():
		return false
	default:
	}
	
	// 尝试读取系统文件以验证权限
	path := `Windows\win.ini`
	f, err := fs.OpenFile(path, os.O_RDONLY, 0666)
	if err != nil {
		return false
	}
	defer f.Close()
	
	return true
}