package services

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

	_ "github.com/lib/pq"
	"github.com/shadow1ng/fscan/common"
	"github.com/shadow1ng/fscan/common/i18n"
)

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

// NewPostgreSQLPlugin 创建PostgreSQL插件
func NewPostgreSQLPlugin() *PostgreSQLPlugin {
	return &PostgreSQLPlugin{
		name:  "postgresql",
		ports: []int{5432, 5433, 5434}, // PostgreSQL端口
	}
}

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

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

// Scan 执行PostgreSQL扫描 - 弱密码检测
func (p *PostgreSQLPlugin) 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("postgresql")
	if len(credentials) == 0 {
		// PostgreSQL默认凭据
		credentials = []Credential{
			{Username: "postgres", Password: ""},
			{Username: "postgres", Password: "postgres"},
			{Username: "postgres", Password: "password"},
			{Username: "postgres", Password: "admin"},
			{Username: "postgres", Password: "123456"},
			{Username: "admin", Password: "admin"},
			{Username: "root", Password: "root"},
		}
	}

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

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

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

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


// testCredential 测试单个凭据 - 返回数据库连接或nil
func (p *PostgreSQLPlugin) testCredential(ctx context.Context, info *common.HostInfo, cred Credential) *sql.DB {
	// 构建连接字符串
	connStr := fmt.Sprintf("postgres://%s:%s@%s:%s/postgres?sslmode=disable&connect_timeout=%d",
		cred.Username, cred.Password, info.Host, info.Ports, common.Timeout)

	// 打开数据库连接
	db, err := sql.Open("postgres", 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 获取PostgreSQL版本信息
func (p *PostgreSQLPlugin) getVersion(db *sql.DB) string {
	var version string
	err := db.QueryRow("SELECT version()").Scan(&version)
	if err != nil {
		return ""
	}
	return version
}

// getDatabases 获取数据库列表
func (p *PostgreSQLPlugin) getDatabases(db *sql.DB) []string {
	query := "SELECT datname FROM pg_database WHERE datistemplate = false"
	rows, err := db.Query(query)
	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 *PostgreSQLPlugin) getTables(db *sql.DB) []string {
	query := `SELECT tablename FROM pg_tables WHERE schemaname = 'public' 
			  UNION SELECT viewname FROM pg_views WHERE schemaname = 'public'
			  ORDER BY 1 LIMIT 20`
			  
	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 *PostgreSQLPlugin) getUsers(db *sql.DB) []string {
	query := "SELECT usename FROM pg_user ORDER BY usename"
	rows, err := db.Query(query)
	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 *PostgreSQLPlugin) getPrivileges(db *sql.DB, username string) string {
	var privileges strings.Builder
	
	// 检查是否为超级用户
	var isSuperUser bool
	err := db.QueryRow("SELECT usesuper FROM pg_user WHERE usename = $1", username).Scan(&isSuperUser)
	if err == nil {
		if isSuperUser {
			privileges.WriteString("超级用户权限: YES\n")
		} else {
			privileges.WriteString("超级用户权限: NO\n")
		}
	}

	// 获取用户属性
	var createDB, createRole bool
	err = db.QueryRow("SELECT usecreatedb, usecreaterole FROM pg_user WHERE usename = $1", 
		username).Scan(&createDB, &createRole)
	if err == nil {
		privileges.WriteString(fmt.Sprintf("创建数据库权限: %v\n", createDB))
		privileges.WriteString(fmt.Sprintf("创建角色权限: %v\n", createRole))
	}

	return privileges.String()
}

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

	// 尝试连接PostgreSQL服务
	connStr := fmt.Sprintf("postgres://invalid:invalid@%s:%s/postgres?sslmode=disable&connect_timeout=%d",
		info.Host, info.Ports, common.Timeout)

	db, err := sql.Open("postgres", connStr)
	if err != nil {
		return &ScanResult{
			Success: false,
			Service: "postgresql",
			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()), "postgres") {
		banner = "PostgreSQL数据库服务"
	} else if err == nil {
		banner = "PostgreSQL数据库 (连接成功)"
	} else {
		return &ScanResult{
			Success: false,
			Service: "postgresql",
			Error:   fmt.Errorf("无法识别为PostgreSQL服务"),
		}
	}

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

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

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