backend/database.go

package main

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

	_ "modernc.org/sqlite"
)

var db *sql.DB

type User struct {
	UserID   int
	Username string
	Passwd   string
	Type     string
}

func InitDatabase(dbPath string) error {
	InitUserDatabase(dbPath)
	InitFrpcDatabase(dbPath)
	return nil
}

func CloseDatabase() error {
	if db != nil {
		return db.Close()
	}
	return nil
}

func isValidInput(input string) bool {
	invalidChars := []string{"'", "\"", ";", "--", "/*", "*/", "xp_", "sp_", "EXEC", "EXECUTE", "DROP", "INSERT", "UPDATE", "DELETE", "SELECT"}
	lowerInput := strings.ToLower(input)
	for _, chars := range invalidChars {
		if strings.Contains(lowerInput, chars) {
			return false
		}
	}
	return true
}

func AddUser(username, passwd string) (int, error) { // New user registration with default type "visitor"
	if !isValidInput(username) || !isValidInput(passwd) {
		return 0, errors.New("invalid input: contains illegal characters")
	}

	if !isValidPassword(passwd) {
		return 0, errors.New("password does not meet complexity requirements")
	}

	hashedPasswd, err := hashPassword(passwd)
	if err != nil {
		return 0, fmt.Errorf("failed to hash password: %w", err)
	}

	result, err := db.Exec("INSERT INTO userLogin (username, passwd, type) VALUES (?, ?, ?)",
		username, hashedPasswd, "visitor")
	if err != nil {
		if strings.Contains(err.Error(), "UNIQUE constraint failed") {
			return 0, errors.New("username already exists")
		}
		return 0, fmt.Errorf("failed to insert user: %w", err)
	}

	lastID, err := result.LastInsertId()
	if err != nil {
		return 0, fmt.Errorf("failed to get last insert id: %w", err)
	}

	var count int
	err = db.QueryRow("SELECT COUNT(*) FROM userLogin WHERE userID = ?", lastID).Scan(&count)
	if err != nil {
		return 0, fmt.Errorf("failed to verify user insertion: %w", err)
	}

	if count == 0 {
		return 0, errors.New("user insertion verification failed")
	}

	return int(lastID), nil
}

func GetUserByUsername(username string) (*User, error) {
	if !isValidInput(username) {
		return nil, errors.New("invalid input: contains illegal characters")
	}

	var user User
	err := db.QueryRow("SELECT userID, username, passwd, type FROM userLogin WHERE username = ?", username).
		Scan(&user.UserID, &user.Username, &user.Passwd, &user.Type)
	if err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return nil, errors.New("user not found")
		}
		return nil, fmt.Errorf("failed to query user: %w", err)
	}

	return &user, nil
}

func GetUserByID(userID int) (*User, error) {
	var user User
	err := db.QueryRow("SELECT userID, username, passwd, type FROM userLogin WHERE userID = ?", userID).
		Scan(&user.UserID, &user.Username, &user.Passwd, &user.Type)
	if err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return nil, errors.New("user not found")
		}
		return nil, fmt.Errorf("failed to query user: %w", err)
	}

	return &user, nil
}

func UpdateUserPassword(userID int, newPasswd string) error {
	if !isValidInput(newPasswd) {
		return errors.New("invalid input: contains illegal characters")
	}

	if !isValidPassword(newPasswd) {
		return errors.New("password does not meet complexity requirements")
	}

	hashedPasswd, err := hashPassword(newPasswd)
	if err != nil {
		return fmt.Errorf("failed to hash password: %w", err)
	}

	result, err := db.Exec("UPDATE userLogin SET passwd = ? WHERE userID = ?", hashedPasswd, userID)
	if err != nil {
		return fmt.Errorf("failed to update password: %w", err)
	}

	rowsAffected, err := result.RowsAffected()
	if err != nil {
		return fmt.Errorf("failed to get rows affected: %w", err)
	}

	if rowsAffected == 0 {
		return errors.New("user not found")
	}

	return nil
}

func UpdateUserType(userID int, newType string) error {
	validTypes := map[string]bool{
		"superuser": true,
		"admin":     true,
		"visitor":   true,
	}

	if !validTypes[newType] {
		return errors.New("invalid user type")
	}

	result, err := db.Exec("UPDATE userLogin SET type = ? WHERE userID = ?", newType, userID)
	if err != nil {
		return fmt.Errorf("failed to update user type: %w", err)
	}

	rowsAffected, err := result.RowsAffected()
	if err != nil {
		return fmt.Errorf("failed to get rows affected: %w", err)
	}

	if rowsAffected == 0 {
		return errors.New("user not found")
	}

	return nil
}

func DeleteUser(userID int) error {
	result, err := db.Exec("DELETE FROM userLogin WHERE userID = ?", userID)
	if err != nil {
		return fmt.Errorf("failed to delete user: %w", err)
	}

	rowsAffected, err := result.RowsAffected()
	if err != nil {
		return fmt.Errorf("failed to get rows affected: %w", err)
	}

	if rowsAffected == 0 {
		return errors.New("user not found")
	}

	return nil
}

func GetNextAvailableUserID() (int, error) {
	var maxID int
	err := db.QueryRow("SELECT COALESCE(MAX(userID), 0) FROM userLogin").Scan(&maxID)
	if err != nil {
		return 0, fmt.Errorf("failed to get max userID: %w", err)
	}
	return maxID + 1, nil
}

func GetAllUsers() ([]User, error) {
	rows, err := db.Query("SELECT userID, username, type FROM userLogin")
	if err != nil {
		return nil, fmt.Errorf("failed to query users: %w", err)
	}
	defer rows.Close()

	var users []User
	for rows.Next() {
		var user User
		if err := rows.Scan(&user.UserID, &user.Username, &user.Type); err != nil {
			return nil, fmt.Errorf("failed to scan user: %w", err)
		}
		users = append(users, user)
	}

	if err = rows.Err(); err != nil {
		return nil, fmt.Errorf("rows error: %w", err)
	}

	return users, nil
}

func InitFrpcDatabase(dbPath string) error {
	var err error
	frpcDB, err = sql.Open("sqlite", dbPath)
	if err != nil {
		return fmt.Errorf("failed to open frpc database: %w", err)
	}

	if err = frpcDB.Ping(); err != nil {
		return fmt.Errorf("failed to ping frpc database: %w", err)
	}

	createTableSQL := `
	CREATE TABLE IF NOT EXISTS frpcInstances (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		userID INTEGER NOT NULL,
		name TEXT NOT NULL,
		serverAddr TEXT NOT NULL,
		serverPort TEXT NOT NULL,
		auth_method TEXT NOT NULL,
		bootAtStart INTEGER NOT NULL DEFAULT 0,
		runUser TEXT NOT NULL DEFAULT 'root',
		configPath TEXT NOT NULL,
		createdAt TEXT NOT NULL,
		UNIQUE(userID, name)
	);
	`
	_, err = frpcDB.Exec(createTableSQL)
	if err != nil {
		return fmt.Errorf("failed to create frpcInstances table: %w", err)
	}

	return nil
}

func InitUserDatabase(dbPath string) error {
	var err error
	db, err = sql.Open("sqlite", dbPath)
	if err != nil {
		return fmt.Errorf("failed to open database: %w", err)
	}

	if err = db.Ping(); err != nil {
		return fmt.Errorf("failed to ping database: %w", err)
	}

	createTableSQL := `
		CREATE TABLE IF NOT EXISTS userLogin (
			userID INTEGER PRIMARY KEY AUTOINCREMENT,
			username TEXT UNIQUE NOT NULL,
			passwd TEXT NOT NULL,
			type TEXT NOT NULL DEFAULT 'visitor'
		);
		`
	_, err = db.Exec(createTableSQL)
	if err != nil {
		return fmt.Errorf("failed to create table: %w", err)
	}

	return nil
}

func DBAddFrpcInstance(instance FrpcInstance) error {
	_, err := frpcDB.Exec("INSERT INTO frpcInstances (userID, name, serverAddr, serverPort, auth_method, bootAtStart, runUser, configPath, createdAt) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
		instance.UserID, instance.Name, instance.ServerAddr, instance.ServerPort, instance.AuthMethod, instance.BootAtStart, instance.RunUser, instance.ConfigPath, time.Now().Format(time.RFC3339))
	if err != nil {
		return fmt.Errorf("failed to insert frpc instance: %w", err)
	}
	return nil
}

func DBQueryFrpcInstance(userID int, instanceName string) (FrpcInstance, error) {
	var instance FrpcInstance
	var createdAtStr string
	err := frpcDB.QueryRow("SELECT id, userID, name, serverAddr, serverPort, auth_method, bootAtStart, runUser, configPath, createdAt FROM frpcInstances WHERE userID = ? AND name = ?", userID, instanceName).Scan(
		&instance.ID, &instance.UserID, &instance.Name, &instance.ServerAddr, &instance.ServerPort, &instance.AuthMethod, &instance.BootAtStart, &instance.RunUser, &instance.ConfigPath, &createdAtStr)
	if err != nil {
		return instance, fmt.Errorf("failed to query frpc instance: %w", err)
	}
	instance.CreatedAt, _ = time.Parse(time.RFC3339, createdAtStr)
	return instance, nil
}

func DBRemoveFrpcInstance(userID int, instanceName string) error {
	_, err := frpcDB.Exec("DELETE FROM frpcInstances WHERE userID = ? AND name = ?", userID, instanceName)
	if err != nil {
		return fmt.Errorf("failed to delete frpc instance: %w", err)
	}
	return nil
}

func DBUpdateFrpcInstance(instance FrpcInstance) error {
	_, err := frpcDB.Exec("UPDATE frpcInstances SET serverAddr = ?, serverPort = ?, auth_method = ?, bootAtStart = ?, runUser = ?, configPath = ? WHERE id = ?",
		instance.ServerAddr, instance.ServerPort, instance.AuthMethod, instance.BootAtStart, instance.RunUser, instance.ConfigPath, instance.ID)
	if err != nil {
		return fmt.Errorf("failed to update frpc instance: %w", err)
	}
	return nil
}

func DBListFrpcInstances(userID int) ([]FrpcInstance, error) {
	rows, err := frpcDB.Query("SELECT id, userID, name, serverAddr, serverPort, auth_method, bootAtStart, runUser, configPath, createdAt FROM frpcInstances WHERE userID = ?", userID)
	if err != nil {
		return nil, fmt.Errorf("failed to query frpc instances: %w", err)
	}
	defer rows.Close()

	var instances []FrpcInstance
	for rows.Next() {
		var instance FrpcInstance
		var createdAtStr string
		if err := rows.Scan(&instance.ID, &instance.UserID, &instance.Name, &instance.ServerAddr, &instance.ServerPort, &instance.AuthMethod, &instance.BootAtStart, &instance.RunUser, &instance.ConfigPath, &createdAtStr); err != nil {
			return nil, fmt.Errorf("failed to scan frpc instance: %w", err)
		}
		instance.CreatedAt, _ = time.Parse(time.RFC3339, createdAtStr)
		instances = append(instances, instance)
	}

	if err = rows.Err(); err != nil {
		return nil, fmt.Errorf("rows error: %w", err)
	}

	return instances, nil
}