frps-console/db.go

package main

import (
	"crypto/rand"
	"database/sql"
	"encoding/hex"
	"fmt"
	"io"
	"log"
	"os"
	"time"

	_ "modernc.org/sqlite"
)

var DB *sql.DB

// ============================================================
// 数据库版本常量
// ============================================================

const SchemaVersion = "2.0.0"

// ============================================================
// 数据模型
// ============================================================

// GlobalConfig frps 全局配置
type GlobalConfig struct {
	ID         int    `json:"id"`
	BindPort   int    `json:"bindPort"`
	Token      string `json:"token"`
	LogLevel   string `json:"logLevel"`
	LogMaxDays int    `json:"logMaxDays"`
	TcpMux     bool   `json:"tcpMux"`
}

// Client 客户端连接记录
type Client struct {
	ID        int    `json:"id"`
	Name      string `json:"name"`
	OS        string `json:"os"`
	Arch      string `json:"arch"`
	Version   string `json:"version"`
	Status    string `json:"status"`
	ConnTime  string `json:"connTime"`
	UpdatedAt string `json:"updatedAt"`
}

// User 用户表
type User struct {
	ID           int    `json:"id"`
	Username     string `json:"username"`
	PasswordHash string `json:"-"`
	CreatedAt    string `json:"createdAt"`
}

// ============================================================
// 数据库初始化
// ============================================================

func InitDB() error {
	var err error
	DB, err = sql.Open("sqlite", "./frps-console.db")
	if err != nil {
		return err
	}

	if err := createTables(); err != nil {
		return err
	}

	if err := runMigrations(); err != nil {
		return err
	}

	if err := ensureJwtSecret(); err != nil {
		return err
	}

	log.Println("✅ 数据库初始化完成 (Schema v" + SchemaVersion + ")")
	return nil
}

// ============================================================
// 建表
// ============================================================

func createTables() error {
	// 用户表
	_, err := DB.Exec(`
        CREATE TABLE IF NOT EXISTS users (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            username TEXT UNIQUE NOT NULL,
            password_hash TEXT NOT NULL,
            created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
            updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
        )
    `)
	if err != nil {
		return err
	}

	// 全局配置表（frps 专用）
	_, err = DB.Exec(`
        CREATE TABLE IF NOT EXISTS global_config (
            id INTEGER PRIMARY KEY CHECK (id = 1),
            bind_port INTEGER NOT NULL DEFAULT 9358,
            token TEXT NOT NULL DEFAULT 'CHANGE_ME',
            log_level TEXT NOT NULL DEFAULT 'info',
            log_max_days INTEGER NOT NULL DEFAULT 3,
            tcp_mux INTEGER NOT NULL DEFAULT 1,
            updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
        )
    `)
	if err != nil {
		return err
	}

	// 客户端记录表
	_, err = DB.Exec(`
        CREATE TABLE IF NOT EXISTS clients (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            name TEXT,
            os TEXT,
            arch TEXT,
            version TEXT,
            status TEXT,
            conn_time DATETIME,
            updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
        )
    `)
	if err != nil {
		return err
	}

	// 应用配置表
	_, err = DB.Exec(`
        CREATE TABLE IF NOT EXISTS app_config (
            key TEXT PRIMARY KEY,
            value TEXT NOT NULL,
            updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
        )
    `)
	if err != nil {
		return err
	}

	// 初始化默认配置
	var count int
	DB.QueryRow("SELECT COUNT(*) FROM global_config").Scan(&count)
	if count == 0 {
		_, err = DB.Exec(`
            INSERT INTO global_config (
                id, bind_port, token, log_level, log_max_days, tcp_mux
            ) VALUES (1, 9358, 'CHANGE_ME', 'info', 3, 1)
        `)
		if err != nil {
			return err
		}
		log.Println("✅ 全局配置初始化完成")
	}

	return nil
}

// ============================================================
// 迁移引擎
// ============================================================

func getSchemaVersion() string {
	var version string
	err := DB.QueryRow("SELECT value FROM app_config WHERE key = 'schema_version'").Scan(&version)
	if err != nil {
		if err == sql.ErrNoRows {
			var count int
			DB.QueryRow("SELECT COUNT(*) FROM users").Scan(&count)
			if count > 0 {
				return "1.0.0"
			}
			return SchemaVersion
		}
		return "1.0.0"
	}
	return version
}

func setSchemaVersion(version string) error {
	_, err := DB.Exec(`
        INSERT INTO app_config (key, value) VALUES ('schema_version', ?)
        ON CONFLICT(key) DO UPDATE SET value = excluded.value, updated_at = CURRENT_TIMESTAMP
    `, version)
	return err
}

func backupDatabase() (string, error) {
	src := "./frps-console.db"
	if _, err := os.Stat(src); os.IsNotExist(err) {
		return "", nil
	}

	timestamp := time.Now().Format("20060102_150405")
	dst := fmt.Sprintf("./frps-console.db.pre-v%s.%s", SchemaVersion, timestamp)

	srcFile, err := os.Open(src)
	if err != nil {
		return "", fmt.Errorf("打开源数据库失败: %w", err)
	}
	defer srcFile.Close()

	dstFile, err := os.Create(dst)
	if err != nil {
		return "", fmt.Errorf("创建备份文件失败: %w", err)
	}
	defer dstFile.Close()

	if _, err := io.Copy(dstFile, srcFile); err != nil {
		return "", fmt.Errorf("复制数据库失败: %w", err)
	}

	log.Printf("✅ 数据库备份完成: %s", dst)
	return dst, nil
}

func restoreDatabase(backupPath string) error {
	srcFile, err := os.Open(backupPath)
	if err != nil {
		return err
	}
	defer srcFile.Close()

	dstFile, err := os.Create("./frps-console.db")
	if err != nil {
		return err
	}
	defer dstFile.Close()

	if _, err := io.Copy(dstFile, srcFile); err != nil {
		return err
	}

	log.Printf("✅ 数据库已从备份恢复: %s", backupPath)
	return nil
}

func runMigrations() error {
	currentVersion := getSchemaVersion()
	log.Printf("📌 当前数据库 Schema: %s, 目标版本: %s", currentVersion, SchemaVersion)

	if currentVersion != SchemaVersion {
		log.Printf("🔄 检测到版本变更 (%s → %s)，开始迁移...", currentVersion, SchemaVersion)

		backupPath, err := backupDatabase()
		if err != nil {
			return fmt.Errorf("备份数据库失败: %w", err)
		}
		if backupPath != "" {
			log.Printf("📦 备份文件: %s", backupPath)
		}

		if err := setSchemaVersion(SchemaVersion); err != nil {
			return fmt.Errorf("更新 Schema 版本失败: %w", err)
		}

		log.Printf("✅ 版本迁移完成")
	} else {
		log.Println("✅ Schema 已是最新，无需迁移")
	}

	return nil
}

// ============================================================
// JWT 密钥管理
// ============================================================

func ensureJwtSecret() error {
	var value string
	err := DB.QueryRow("SELECT value FROM app_config WHERE key = 'jwt_secret'").Scan(&value)
	if err == nil && value != "" {
		return nil
	}

	bytes := make([]byte, 32)
	if _, err := rand.Read(bytes); err != nil {
		return err
	}
	secret := hex.EncodeToString(bytes)

	_, err = DB.Exec(`
        INSERT INTO app_config (key, value) VALUES ('jwt_secret', ?)
    `, secret)
	if err != nil {
		return err
	}
	log.Printf("✅ JWT 密钥已生成")
	return nil
}

func GetJwtSecret() (string, error) {
	var secret string
	err := DB.QueryRow("SELECT value FROM app_config WHERE key = 'jwt_secret'").Scan(&secret)
	if err != nil {
		return "", err
	}
	return secret, nil
}

// ============================================================
// 全局配置 CRUD
// ============================================================

func GetGlobalConfig() (*GlobalConfig, error) {
	var cfg GlobalConfig
	err := DB.QueryRow(`
        SELECT id, bind_port, token, log_level, log_max_days, tcp_mux
        FROM global_config WHERE id = 1
    `).Scan(
		&cfg.ID, &cfg.BindPort, &cfg.Token,
		&cfg.LogLevel, &cfg.LogMaxDays,
		&cfg.TcpMux,
	)
	if err != nil {
		return nil, err
	}
	cfg.TcpMux = true
	return &cfg, nil
}

func UpdateGlobalConfig(cfg *GlobalConfig) error {
	_, err := DB.Exec(`
        UPDATE global_config SET
            bind_port = ?, token = ?, log_level = ?, log_max_days = ?,
            tcp_mux = 1,
            updated_at = CURRENT_TIMESTAMP
        WHERE id = 1
    `, cfg.BindPort, cfg.Token, cfg.LogLevel, cfg.LogMaxDays)
	return err
}

// ============================================================
// 客户端记录 CRUD
// ============================================================

func SaveClient(client *Client) error {
	_, err := DB.Exec(`
        INSERT INTO clients (name, os, arch, version, status, conn_time)
        VALUES (?, ?, ?, ?, ?, ?)
    `, client.Name, client.OS, client.Arch, client.Version, client.Status, client.ConnTime)
	return err
}

func GetClients() ([]Client, error) {
	rows, err := DB.Query(`
        SELECT id, name, os, arch, version, status, conn_time, updated_at
        FROM clients ORDER BY id DESC LIMIT 100
    `)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var clients []Client
	for rows.Next() {
		var c Client
		err := rows.Scan(&c.ID, &c.Name, &c.OS, &c.Arch, &c.Version, &c.Status, &c.ConnTime, &c.UpdatedAt)
		if err != nil {
			return nil, err
		}
		clients = append(clients, c)
	}
	return clients, rows.Err()
}

// ============================================================
// 用户 CRUD
// ============================================================

func GetUserByUsername(username string) (*User, error) {
	var u User
	err := DB.QueryRow(`
        SELECT id, username, password_hash, created_at
        FROM users WHERE username = ?
    `, username).Scan(&u.ID, &u.Username, &u.PasswordHash, &u.CreatedAt)
	if err != nil {
		return nil, err
	}
	return &u, nil
}

func CreateUser(username, passwordHash string) error {
	_, err := DB.Exec(`
        INSERT INTO users (username, password_hash) VALUES (?, ?)
    `, username, passwordHash)
	return err
}

func UpdatePassword(username, passwordHash string) error {
	_, err := DB.Exec(`
        UPDATE users SET password_hash = ?, updated_at = CURRENT_TIMESTAMP
        WHERE username = ?
    `, passwordHash, username)
	return err
}

func CountUsers() (int, error) {
	var count int
	err := DB.QueryRow("SELECT COUNT(*) FROM users").Scan(&count)
	return count, err
}