package services

import (
	"context"
	"database/sql"
	"fmt"
	"strings"
	"time"

	_ "github.com/denisenkom/go-mssqldb"
	"github.com/shadow1ng/fscan/common"
	"github.com/shadow1ng/fscan/common/i18n"
)

// MSSQLPlugin Microsoft SQL Server数据库扫描和利用插件 - 包含数据库查询利用功能
type MSSQLPlugin struct {
	name  string
	ports []int
}

// NewMSSQLPlugin 创建MSSQL插件
func NewMSSQLPlugin() *MSSQLPlugin {
	return &MSSQLPlugin{
		name:  "mssql",
		ports: []int{1433, 1434}, // MSSQL端口
	}
}

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

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

// Scan 执行MSSQL扫描 - 弱密码检测
func (p *MSSQLPlugin) Scan(ctx context.Context, info *common.HostInfo) *ScanResult {
	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)

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

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

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

		// 测试凭据
		if db := p.testCredential(ctx, info, cred); db != nil {
			db.Close() // 关闭测试连接
			
			// MSSQL认证成功
			common.LogSuccess(i18n.GetText("mssql_scan_success", target, cred.Username, cred.Password))

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

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

// Exploit 执行MSSQL利用操作 - 实现数据库查询功能
func (p *MSSQLPlugin) Exploit(ctx context.Context, info *common.HostInfo, creds Credential) *ExploitResult {
	// 建立MSSQL连接
	db := p.testCredential(ctx, info, creds)
	if db == nil {
		return &ExploitResult{
			Success: false,
			Error:   fmt.Errorf("MSSQL连接失败"),
		}
	}
	defer db.Close()

	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)
	common.LogSuccess(fmt.Sprintf("MSSQL利用开始: %s (用户: %s)", target, creds.Username))

	var output strings.Builder
	output.WriteString(fmt.Sprintf("=== MSSQL利用结果 - %s ===\n", target))

	// 获取版本信息
	if version := p.getVersion(db); version != "" {
		output.WriteString(fmt.Sprintf("\n[版本信息]\n%s\n", version))
	}

	// 获取数据库列表
	if databases := p.getDatabases(db); len(databases) > 0 {
		output.WriteString(fmt.Sprintf("\n[数据库列表] (共%d个)\n", len(databases)))
		for i, dbName := range databases {
			if i >= 10 { // 限制显示前10个
				output.WriteString("... (更多数据库)\n")
				break
			}
			output.WriteString(fmt.Sprintf("  %s\n", dbName))
		}
	}

	// 获取表列表（master数据库）
	if tables := p.getTables(db, "master"); len(tables) > 0 {
		output.WriteString(fmt.Sprintf("\n[master数据库表] (共%d个)\n", len(tables)))
		for i, table := range tables {
			if i >= 5 { // 限制显示前5个表
				output.WriteString("... (更多表)\n")
				break
			}
			output.WriteString(fmt.Sprintf("  %s\n", table))
		}
	}

	// 获取用户列表
	if users := p.getUsers(db); len(users) > 0 {
		output.WriteString(fmt.Sprintf("\n[用户列表] (共%d个)\n", len(users)))
		for i, user := range users {
			if i >= 10 { // 限制显示前10个用户
				output.WriteString("... (更多用户)\n")
				break
			}
			output.WriteString(fmt.Sprintf("  %s\n", user))
		}
	}

	// 获取权限信息
	if privileges := p.getPrivileges(db, creds.Username); privileges != "" {
		output.WriteString(fmt.Sprintf("\n[用户权限]\n%s\n", privileges))
	}

	// 检查xp_cmdshell状态
	if cmdshellStatus := p.checkXpCmdshell(db); cmdshellStatus != "" {
		output.WriteString(fmt.Sprintf("\n[系统命令执行]\n%s\n", cmdshellStatus))
	}

	common.LogSuccess(fmt.Sprintf("MSSQL利用完成: %s", target))

	return &ExploitResult{
		Success: true,
		Output:  output.String(),
	}
}

// testCredential 测试单个凭据 - 返回数据库连接或nil
func (p *MSSQLPlugin) testCredential(ctx context.Context, info *common.HostInfo, cred Credential) *sql.DB {
	// 构建连接字符串
	connStr := fmt.Sprintf("server=%s;user id=%s;password=%s;port=%s;database=master;connection timeout=%d",
		info.Host, cred.Username, cred.Password, info.Ports, common.Timeout)

	// 打开数据库连接
	db, err := sql.Open("mssql", connStr)
	if err != nil {
		return nil
	}

	// 设置连接参数
	db.SetConnMaxLifetime(time.Duration(common.Timeout) * time.Second)
	db.SetMaxOpenConns(1)
	db.SetMaxIdleConns(0)

	// 创建超时上下文
	pingCtx, cancel := context.WithTimeout(ctx, time.Duration(common.Timeout)*time.Second)
	defer cancel()

	// 测试连接
	err = db.PingContext(pingCtx)
	if err != nil {
		db.Close()
		return nil
	}

	return db
}

// getVersion 获取MSSQL版本信息
func (p *MSSQLPlugin) getVersion(db *sql.DB) string {
	var version string
	err := db.QueryRow("SELECT @@VERSION").Scan(&version)
	if err != nil {
		return ""
	}
	
	// 只返回第一行版本信息
	lines := strings.Split(version, "\n")
	if len(lines) > 0 {
		return strings.TrimSpace(lines[0])
	}
	
	return version
}

// getDatabases 获取数据库列表
func (p *MSSQLPlugin) getDatabases(db *sql.DB) []string {
	query := "SELECT name FROM sys.databases WHERE database_id > 4" // 排除系统数据库
	rows, err := db.Query(query)
	if err != nil {
		// 如果查询失败，尝试查询所有数据库
		rows, err = db.Query("SELECT name FROM sys.databases")
		if err != nil {
			return nil
		}
	}
	defer rows.Close()

	var databases []string
	for rows.Next() {
		var dbName string
		if err := rows.Scan(&dbName); err == nil {
			databases = append(databases, dbName)
		}
	}

	return databases
}

// getTables 获取指定数据库的表列表
func (p *MSSQLPlugin) getTables(db *sql.DB, database string) []string {
	query := fmt.Sprintf("SELECT TABLE_NAME FROM %s.INFORMATION_SCHEMA.TABLES WHERE TABLE_TYPE = 'BASE TABLE'", database)
	rows, err := db.Query(query)
	if err != nil {
		return nil
	}
	defer rows.Close()

	var tables []string
	for rows.Next() {
		var tableName string
		if err := rows.Scan(&tableName); err == nil {
			tables = append(tables, tableName)
		}
	}

	return tables
}

// getUsers 获取用户列表
func (p *MSSQLPlugin) getUsers(db *sql.DB) []string {
	query := "SELECT name FROM sys.server_principals WHERE type IN ('S', 'U') AND is_disabled = 0"
	rows, err := db.Query(query)
	if err != nil {
		// 尝试备用查询
		rows, err = db.Query("SELECT loginname FROM master.dbo.syslogins")
		if err != nil {
			return nil
		}
	}
	defer rows.Close()

	var users []string
	for rows.Next() {
		var userName string
		if err := rows.Scan(&userName); err == nil {
			users = append(users, userName)
		}
	}

	return users
}

// getPrivileges 获取用户权限信息
func (p *MSSQLPlugin) getPrivileges(db *sql.DB, username string) string {
	var privileges strings.Builder
	
	// 检查是否为sysadmin
	var isSysadmin int
	err := db.QueryRow("SELECT IS_SRVROLEMEMBER('sysadmin', ?)", username).Scan(&isSysadmin)
	if err == nil {
		if isSysadmin == 1 {
			privileges.WriteString("sysadmin权限: YES\n")
		} else {
			privileges.WriteString("sysadmin权限: NO\n")
		}
	}

	// 检查其他服务器角色
	serverRoles := []string{"securityadmin", "serveradmin", "setupadmin", "processadmin", "diskadmin", "dbcreator", "bulkadmin"}
	for _, role := range serverRoles {
		var hasRole int
		err := db.QueryRow(fmt.Sprintf("SELECT IS_SRVROLEMEMBER('%s', ?)", role), username).Scan(&hasRole)
		if err == nil && hasRole == 1 {
			privileges.WriteString(fmt.Sprintf("%s权限: YES\n", role))
		}
	}

	return privileges.String()
}

// checkXpCmdshell 检查xp_cmdshell状态
func (p *MSSQLPlugin) checkXpCmdshell(db *sql.DB) string {
	// 检查xp_cmdshell是否启用
	var configValue int
	err := db.QueryRow("SELECT CAST(value as int) FROM sys.configurations WHERE name = 'xp_cmdshell'").Scan(&configValue)
	if err != nil {
		return "无法检查xp_cmdshell状态"
	}

	if configValue == 1 {
		// 尝试执行一个简单的命令测试
		var result sql.NullString
		err = db.QueryRow("EXEC xp_cmdshell 'echo test'").Scan(&result)
		if err == nil && result.Valid {
			return "✅ xp_cmdshell已启用，支持系统命令执行"
		}
		return "⚠️  xp_cmdshell已启用但无法执行命令"
	}

	return "❌ xp_cmdshell未启用"
}

// identifyService 服务识别 - 检测MSSQL服务
func (p *MSSQLPlugin) identifyService(ctx context.Context, info *common.HostInfo) *ScanResult {
	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)

	// 尝试连接MSSQL服务
	connStr := fmt.Sprintf("server=%s;user id=invalid;password=invalid;port=%s;database=master;connection timeout=%d",
		info.Host, info.Ports, common.Timeout)

	db, err := sql.Open("mssql", connStr)
	if err != nil {
		return &ScanResult{
			Success: false,
			Service: "mssql",
			Error:   err,
		}
	}
	defer db.Close()

	// 尝试连接（即使认证失败，也能识别服务）
	pingCtx, cancel := context.WithTimeout(ctx, time.Duration(common.Timeout)*time.Second)
	defer cancel()

	err = db.PingContext(pingCtx)
	
	var banner string
	if err != nil && (strings.Contains(strings.ToLower(err.Error()), "login failed") ||
					 strings.Contains(strings.ToLower(err.Error()), "mssql") ||
					 strings.Contains(strings.ToLower(err.Error()), "sql server")) {
		banner = "Microsoft SQL Server数据库服务"
	} else if err == nil {
		banner = "Microsoft SQL Server (连接成功)"
	} else {
		return &ScanResult{
			Success: false,
			Service: "mssql",
			Error:   fmt.Errorf("无法识别为MSSQL服务"),
		}
	}

	common.LogSuccess(i18n.GetText("mssql_service_identified", target, banner))

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

// init 自动注册插件
func init() {
	RegisterPlugin("mssql", func() Plugin {
		return NewMSSQLPlugin()
	})
}