HY2XS_flamy/apps/dao/sqlite.go

package dao

import (
	"errors"
	"fmt"
	"github.com/glebarez/sqlite"
	"github.com/sirupsen/logrus"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
	"gorm.io/gorm/schema"
	"hy2xs-admin/model/constant"
	"hy2xs-admin/model/entity"
	"hy2xs-admin/util"
	"log"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"time"
)

var sqliteDB *gorm.DB

func InitSqliteDB() error {
	var err error
	sqliteDB, err = gorm.Open(sqlite.Open(constant.SqliteDBPath), &gorm.Config{
		TranslateError: true,
		Logger: logger.New(
			log.New(os.Stdout, "\r\n", log.LstdFlags),
			logger.Config{
				SlowThreshold:             time.Second,
				LogLevel:                  logger.Silent,
				IgnoreRecordNotFoundError: true,
				ParameterizedQueries:      true,
				Colorful:                  false,
			},
		),
		NamingStrategy: schema.NamingStrategy{
			SingularTable: true,
		},
	})
	if err != nil {
		logrus.Errorf("sqlite open err: %v", err)
		return errors.New("sqlite open err")
	}
	return nil
}

func InitSql(port string) error {
	if err := InitSqliteDB(); err != nil {
		return err
	}
	if err := runMigrations(); err != nil {
		return err
	}
	if port != "" {
		if tx := sqliteDB.Exec("UPDATE config set value = ? where key = 'H_UI_WEB_PORT'", port); tx.Error != nil {
			logrus.Errorf("sqlite exec err: %v", tx.Error)
			return errors.New("sqlite exec err")
		}
	}
	if err := ensureSecureBootstrapAdmin(); err != nil {
		return err
	}
	if err := ensureTrafficStatsSecret(); err != nil {
		return err
	}

	return nil
}

func envInt(name string, fallback int) int {
	if raw := strings.TrimSpace(os.Getenv(name)); raw != "" {
		if value, err := strconv.Atoi(raw); err == nil {
			return value
		}
	}
	return fallback
}

func envBoolAsInt(name string, fallback int) int {
	raw := strings.TrimSpace(strings.ToLower(os.Getenv(name)))
	if raw == "" {
		return fallback
	}
	if raw == "true" || raw == "1" {
		return 1
	}
	if raw == "false" || raw == "0" {
		return 0
	}
	return fallback
}

func ensureSecureBootstrapAdmin() error {
	adminUser := strings.TrimSpace(os.Getenv("HY2XS_ADMIN_USER"))
	if adminUser == "" {
		adminUser = "hy2xsadmin"
	}
	adminPassword := strings.TrimSpace(os.Getenv("HY2XS_ADMIN_INITIAL_PASSWORD"))
	if adminPassword == "" {
		password, pwdErr := util.RandomString(18)
		if pwdErr != nil {
			return pwdErr
		}
		adminPassword = password
		logrus.Warnf("Initial admin username: %s", adminUser)
		logrus.Warnf("Initial admin password: %s", adminPassword)
	}
	forcePasswordChange := envBoolAsInt("HY2XS_FORCE_PASSWORD_CHANGE", 1)
	status := int64(1)
	tokenVersion := int64(1)
	passwordChangedAt := time.Now().UnixMilli()
	hash, hashErr := util.HashPassword(adminPassword)
	if hashErr != nil {
		return hashErr
	}

	admin, err := GetAdminUser("username = ?", adminUser)
	if err != nil {
		username := adminUser
		account := entity.AdminUser{
			Username:            &username,
			PasswordHash:        &hash,
			Status:              &status,
			TokenVersion:        &tokenVersion,
			PasswordChangedAt:   &passwordChangedAt,
			ForcePasswordChange: func() *int64 { v := int64(forcePasswordChange); return &v }(),
		}
		if _, saveErr := SaveAdminUser(account); saveErr != nil {
			return saveErr
		}
		return nil
	}

	if admin.PasswordHash == nil {
		return nil
	}
	return nil
}

func runMigrations() error {
	if tx := sqliteDB.Exec(`CREATE TABLE IF NOT EXISTS schema_migrations (
		version TEXT PRIMARY KEY,
		applied_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
	)`); tx.Error != nil {
		logrus.Errorf("sqlite migration init err: %v", tx.Error)
		return errors.New("sqlite migration init err")
	}

	type migration struct {
		version string
		apply   func() error
	}

	migrations := []migration{
		{version: "000_base_config", apply: migrateBaseConfig},
		{version: "001_admin_peer_split", apply: migrateAdminPeerSplit},
		{version: "002_migrate_legacy_accounts", apply: migrateLegacyAccounts},
		{version: "003_archive_legacy_account", apply: archiveLegacyAccount},
		{version: "004_traffic_samples_and_aggregates", apply: migrateTrafficTables},
		{version: "005_metric_sample", apply: migrateMetricSampleTable},
	}

	for _, m := range migrations {
		if isApplied, err := migrationApplied(m.version); err != nil {
			return err
		} else if isApplied {
			continue
		}
		if err := m.apply(); err != nil {
			return err
		}
		if tx := sqliteDB.Exec("INSERT INTO schema_migrations(version) VALUES(?)", m.version); tx.Error != nil {
			logrus.Errorf("sqlite migration mark err: %v", tx.Error)
			return errors.New("sqlite migration mark err")
		}
	}

	return nil
}

func seedBaseConfig() error {
	defaults := map[string]string{
		constant.HUIWebPort:                  "8080",
		constant.HUIWebContext:               "/",
		constant.HUICrtPath:                  "",
		constant.HUIKeyPath:                  "",
		constant.JwtSecret:                   "",
		constant.Hysteria2Enable:             "0",
		constant.Hysteria2Config:             "",
		constant.Hysteria2TrafficTime:        "10",
		constant.Hysteria2ConfigRemark:       "",
		constant.ResetTrafficCron:            "0 0 * * *",
		constant.Hysteria2TrafficStatsSecret: "",
		constant.PeerSecretKey:               "",
		constant.PeerSecretEncryptionKey:     "",
	}
	for k, v := range defaults {
		if tx := sqliteDB.Exec("INSERT OR IGNORE INTO config(key, value, remark) VALUES(?, ?, ?)", k, v, k); tx.Error != nil {
			logrus.Errorf("sqlite seed config err: %v", tx.Error)
			return errors.New("sqlite seed config err")
		}
	}
	return nil
}

func migrateBaseConfig() error {
	stmts := []string{
		`CREATE TABLE IF NOT EXISTS config (
			id INTEGER PRIMARY KEY AUTOINCREMENT,
			key TEXT NOT NULL UNIQUE DEFAULT '',
			value TEXT NOT NULL DEFAULT '',
			remark TEXT NOT NULL DEFAULT '',
			create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
			update_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP
		)`,
		`CREATE INDEX IF NOT EXISTS config_key_index ON config (key)`,
	}
	for _, stmt := range stmts {
		if tx := sqliteDB.Exec(stmt); tx.Error != nil {
			logrus.Errorf("sqlite migration exec err: %v", tx.Error)
			return errors.New("sqlite migration exec err")
		}
	}
	return seedBaseConfig()
}

func migrationApplied(version string) (bool, error) {
	var count int64
	if tx := sqliteDB.Raw("SELECT COUNT(1) FROM schema_migrations WHERE version = ?", version).Scan(&count); tx.Error != nil {
		logrus.Errorf("sqlite migration query err: %v", tx.Error)
		return false, errors.New("sqlite migration query err")
	}
	return count > 0, nil
}

func migrateAdminPeerSplit() error {
	stmts := []string{
		`CREATE TABLE IF NOT EXISTS admin_user (
			id INTEGER PRIMARY KEY AUTOINCREMENT,
			username TEXT NOT NULL UNIQUE DEFAULT '',
			password_hash TEXT NOT NULL DEFAULT '',
			status INTEGER NOT NULL DEFAULT 1,
			force_password_change INTEGER NOT NULL DEFAULT 0,
			last_login_at INTEGER NOT NULL DEFAULT 0,
			password_changed_at INTEGER NOT NULL DEFAULT 0,
			token_version INTEGER NOT NULL DEFAULT 1,
			create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
			update_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP
		)`,
		`CREATE INDEX IF NOT EXISTS admin_user_username_index ON admin_user (username)`,
		`CREATE TABLE IF NOT EXISTS peer (
			id INTEGER PRIMARY KEY AUTOINCREMENT,
			name TEXT NOT NULL UNIQUE DEFAULT '',
			remark TEXT NOT NULL DEFAULT '',
			auth_id TEXT NOT NULL UNIQUE DEFAULT '',
			secret_digest TEXT NOT NULL UNIQUE DEFAULT '',
			secret_ciphertext TEXT NOT NULL DEFAULT '',
			quota_bytes INTEGER NOT NULL DEFAULT 0,
			download_bytes INTEGER NOT NULL DEFAULT 0,
			upload_bytes INTEGER NOT NULL DEFAULT 0,
			expires_at INTEGER NOT NULL DEFAULT 0,
			max_devices INTEGER NOT NULL DEFAULT 3,
			disabled INTEGER NOT NULL DEFAULT 0,
			banned_until INTEGER NOT NULL DEFAULT 0,
			last_connection_at INTEGER NOT NULL DEFAULT 0,
			create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
			update_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP
		)`,
		`CREATE INDEX IF NOT EXISTS peer_name_index ON peer (name)`,
		`CREATE INDEX IF NOT EXISTS peer_auth_id_index ON peer (auth_id)`,
		`CREATE INDEX IF NOT EXISTS peer_secret_digest_index ON peer (secret_digest)`,
	}
	for _, stmt := range stmts {
		if tx := sqliteDB.Exec(stmt); tx.Error != nil {
			logrus.Errorf("sqlite migration exec err: %v", tx.Error)
			return errors.New("sqlite migration exec err")
		}
	}
	return nil
}

func migrateLegacyAccounts() error {
	if !tableExists("account") {
		return nil
	}
	var accounts []entity.Account
	if tx := sqliteDB.Model(&entity.Account{}).Order("id asc").Find(&accounts); tx.Error != nil {
		logrus.Errorf("sqlite legacy account query err: %v", tx.Error)
		return errors.New("sqlite legacy account query err")
	}
	nowMs := time.Now().UnixMilli()
	for _, acc := range accounts {
		if acc.Role != nil && *acc.Role == "admin" {
			if acc.Username == nil || acc.Pass == nil {
				continue
			}
			_, getErr := GetAdminUser("username = ?", *acc.Username)
			if getErr == nil {
				continue
			}
			status := int64(1)
			if acc.Deleted != nil && *acc.Deleted == 1 {
				status = 0
			}
			lastLogin := int64(0)
			if acc.LoginAt != nil {
				lastLogin = *acc.LoginAt
			}
			passwordChangedAt := nowMs
			admin := entity.AdminUser{
				Username:            acc.Username,
				PasswordHash:        acc.Pass,
				Status:              &status,
				ForcePasswordChange: acc.ForcePasswordChange,
				LastLoginAt:         &lastLogin,
				PasswordChangedAt:   &passwordChangedAt,
				TokenVersion:        func() *int64 { v := int64(1); return &v }(),
			}
			if _, saveErr := SaveAdminUser(admin); saveErr != nil {
				return saveErr
			}
			continue
		}

		if acc.Username == nil || acc.ConPass == nil {
			continue
		}
		_, getPeerErr := GetPeer("name = ?", *acc.Username)
		if getPeerErr == nil {
			continue
		}

		authId, authErr := util.RandomString(18)
		if authErr != nil {
			return authErr
		}
		secretDigest := util.PeerSecretDigest(*acc.ConPass)
		secretEncrypted := *acc.ConPass
		quota := int64(0)
		if acc.Quota != nil {
			quota = *acc.Quota
		}
		download := int64(0)
		if acc.Download != nil {
			download = *acc.Download
		}
		upload := int64(0)
		if acc.Upload != nil {
			upload = *acc.Upload
		}
		expires := int64(0)
		if acc.ExpireTime != nil {
			expires = *acc.ExpireTime
		}
		maxDevices := int64(3)
		if acc.DeviceNo != nil {
			maxDevices = *acc.DeviceNo
		}
		disabled := int64(0)
		if acc.Deleted != nil {
			disabled = *acc.Deleted
		}
		bannedUntil := int64(0)
		if acc.KickUtilTime != nil {
			bannedUntil = *acc.KickUtilTime
		}
		lastConnection := int64(0)
		if acc.ConAt != nil {
			lastConnection = *acc.ConAt
		}
		remark := ""
		if acc.Remark != nil {
			remark = *acc.Remark
		}
		peer := entity.Peer{
			Name:             acc.Username,
			Remark:           &remark,
			AuthId:           &authId,
			SecretDigest:     &secretDigest,
			SecretEncrypted:  &secretEncrypted,
			QuotaBytes:       &quota,
			DownloadBytes:    &download,
			UploadBytes:      &upload,
			ExpiresAt:        &expires,
			MaxDevices:       &maxDevices,
			Disabled:         &disabled,
			BannedUntil:      &bannedUntil,
			LastConnectionAt: &lastConnection,
		}
		if _, saveErr := SavePeer(peer); saveErr != nil {
			return saveErr
		}
	}
	return nil
}

func archiveLegacyAccount() error {
	if !tableExists("account") {
		return nil
	}
	backupName := fmt.Sprintf("legacy_account_backup_%d", time.Now().Unix())
	if tx := sqliteDB.Exec("ALTER TABLE account RENAME TO " + backupName); tx.Error != nil {
		logrus.Errorf("sqlite legacy archive err: %v", tx.Error)
		return errors.New("sqlite legacy archive err")
	}
	return nil
}

func migrateTrafficTables() error {
	stmts := []string{
		`CREATE TABLE IF NOT EXISTS traffic_sample (
			id INTEGER PRIMARY KEY AUTOINCREMENT,
			peer_id INTEGER NOT NULL,
			auth_id TEXT NOT NULL,
			rx_bytes INTEGER NOT NULL DEFAULT 0,
			tx_bytes INTEGER NOT NULL DEFAULT 0,
			sampled_at INTEGER NOT NULL,
			create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
			update_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP
		)`,
		`CREATE INDEX IF NOT EXISTS idx_traffic_sample_time ON traffic_sample(sampled_at)`,
		`CREATE INDEX IF NOT EXISTS idx_traffic_sample_peer_time ON traffic_sample(peer_id, sampled_at)`,

		`CREATE TABLE IF NOT EXISTS traffic_aggregate_hourly (
			peer_id INTEGER NOT NULL,
			hour_start INTEGER NOT NULL,
			rx_bytes INTEGER NOT NULL DEFAULT 0,
			tx_bytes INTEGER NOT NULL DEFAULT 0,
			create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
			update_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
			PRIMARY KEY(peer_id, hour_start)
		)`,
		`CREATE INDEX IF NOT EXISTS idx_traffic_hourly_hour_start ON traffic_aggregate_hourly(hour_start)`,

		`CREATE TABLE IF NOT EXISTS traffic_aggregate_daily (
			peer_id INTEGER NOT NULL,
			day_start INTEGER NOT NULL,
			rx_bytes INTEGER NOT NULL DEFAULT 0,
			tx_bytes INTEGER NOT NULL DEFAULT 0,
			create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
			update_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
			PRIMARY KEY(peer_id, day_start)
		)`,
		`CREATE INDEX IF NOT EXISTS idx_traffic_daily_day_start ON traffic_aggregate_daily(day_start)`,
	}
	for _, stmt := range stmts {
		if tx := sqliteDB.Exec(stmt); tx.Error != nil {
			logrus.Errorf("sqlite migration exec err: %v", tx.Error)
			return errors.New("sqlite migration exec err")
		}
	}
	return nil
}

func migrateMetricSampleTable() error {
	stmts := []string{
		`CREATE TABLE IF NOT EXISTS metric_sample (
			id INTEGER PRIMARY KEY AUTOINCREMENT,
			sampled_at INTEGER NOT NULL,
			cpu_percent REAL NOT NULL DEFAULT 0,
			load1 REAL NOT NULL DEFAULT 0,
			mem_used_bytes INTEGER NOT NULL DEFAULT 0,
			mem_total_bytes INTEGER NOT NULL DEFAULT 0,
			mem_percent REAL NOT NULL DEFAULT 0,
			disk_path TEXT NOT NULL DEFAULT '/',
			disk_used_bytes INTEGER NOT NULL DEFAULT 0,
			disk_total_bytes INTEGER NOT NULL DEFAULT 0,
			disk_percent REAL NOT NULL DEFAULT 0,
			hysteria_running INTEGER NOT NULL DEFAULT 0,
			online_peers INTEGER NOT NULL DEFAULT 0,
			online_devices INTEGER NOT NULL DEFAULT 0,
			create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
			update_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP
		)`,
		`CREATE INDEX IF NOT EXISTS idx_metric_sample_time ON metric_sample(sampled_at)`,
	}
	for _, stmt := range stmts {
		if tx := sqliteDB.Exec(stmt); tx.Error != nil {
			logrus.Errorf("sqlite migration exec err: %v", tx.Error)
			return errors.New("sqlite migration exec err")
		}
	}
	return nil
}

func tableExists(tableName string) bool {
	var count int64
	if tx := sqliteDB.Raw("SELECT COUNT(1) FROM sqlite_master WHERE type='table' AND name=?", tableName).Scan(&count); tx.Error != nil {
		return false
	}
	return count > 0
}

func listSQLMigrationFiles(dir string) ([]string, error) {
	entries, err := os.ReadDir(dir)
	if err != nil {
		return nil, err
	}
	files := make([]string, 0)
	for _, entry := range entries {
		if entry.IsDir() {
			continue
		}
		name := entry.Name()
		if strings.HasSuffix(strings.ToLower(name), ".sql") {
			files = append(files, filepath.Join(dir, name))
		}
	}
	sort.Strings(files)
	return files, nil
}

func ensureTrafficStatsSecret() error {
	envSecret := strings.TrimSpace(os.Getenv("HY2XS_HYSTERIA_TRAFFIC_STATS_SECRET"))
	if envSecret != "" {
		if existing, err := GetConfig("key = ?", constant.Hysteria2TrafficStatsSecret); err == nil {
			if existing.Value != nil && *existing.Value != envSecret {
				return UpdateConfig([]string{constant.Hysteria2TrafficStatsSecret}, map[string]interface{}{"value": envSecret})
			}
			return nil
		}
		key := constant.Hysteria2TrafficStatsSecret
		remark := "Hysteria2 trafficStats API secret"
		if _, saveErr := SaveConfig(entity.Config{Key: &key, Value: &envSecret, Remark: &remark}); saveErr != nil {
			return saveErr
		}
		return nil
	}

	if _, err := GetConfig("key = ?", constant.Hysteria2TrafficStatsSecret); err == nil {
		return nil
	}
	secret, secErr := util.RandomString(32)
	if secErr != nil {
		return secErr
	}
	key := constant.Hysteria2TrafficStatsSecret
	remark := "Hysteria2 trafficStats API secret"
	if _, saveErr := SaveConfig(entity.Config{Key: &key, Value: &secret, Remark: &remark}); saveErr != nil {
		return saveErr
	}
	return nil
}

func CloseSqliteDB() error {
	if sqliteDB != nil {
		db, err := sqliteDB.DB()
		if err != nil {
			logrus.Errorf("sqlite err: %v", err)
			return errors.New("sqlite err")
		}
		if err = db.Close(); err != nil {
			logrus.Errorf("sqlite close err: %v", err)
			return errors.New("sqlite close err")
		}
	}
	return nil
}

func IsSqliteReady() bool {
  if sqliteDB == nil {
    return false
  }
  db, err := sqliteDB.DB()
  if err != nil {
    return false
  }
  if err := db.Ping(); err != nil {
    return false
  }
  return true
}

func Paginate(pageNum *int64, pageSize *int64) func(db *gorm.DB) *gorm.DB {
	var num int64 = 1
	var size int64 = 10
	if pageNum != nil && *pageNum > 0 {
		num = *pageNum
	}
	if pageSize != nil && *pageSize > 0 {
		size = *pageSize
	}
	return func(db *gorm.DB) *gorm.DB {
		return db.Offset(int((num - 1) * size)).Limit(int(size))
	}
}