package main import ( "database/sql" "errors" "fmt" "strconv" "strings" "super-frpc/postLog" "time" _ "modernc.org/sqlite" ) var db *sql.DB var logsDB *sql.DB func GetLogsDatabase() *sql.DB { return logsDB } type User struct { UserID int Username string Passwd string Type string CreatedAt string } type FrpcInstance struct { ID int UserID int Name string ServerAddr string ServerPort string AuthMethod string BootAtStart bool RunUser string ConfigPath string CreatedAt time.Time CreatedBy string } func InitDatabase(dbPath_data string, dbPath_log string) error { InitUserDatabase(dbPath_data) InitFrpcDatabase(dbPath_data) postLog.InitLogsDatabase(dbPath_log) return nil } func CloseDatabase() error { if db != nil { if err := db.Close(); err != nil { return err } } if logsDB != nil { if err := logsDB.Close(); err != nil { return err } } 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, userType string) (int, error) { // New user registration with specified type if !isValidInput(username) || !isValidInput(passwd) || !isValidInput(userType) { 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, userType) 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 RemoveUser(userID int) error { if !isValidInput(strconv.Itoa(userID)) { return errors.New("invalid input: contains illegal characters") } 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 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, createdAt FROM userLogin WHERE username = ?", username). Scan(&user.UserID, &user.Username, &user.Passwd, &user.Type, &user.CreatedAt) 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, createdAt FROM userLogin WHERE userID = ?", userID). Scan(&user.UserID, &user.Username, &user.Passwd, &user.Type, &user.CreatedAt) 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 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', createdAt TEXT NOT NULL DEFAULT (datetime('now')) ); ` _, err = db.Exec(createTableSQL) if err != nil { return fmt.Errorf("failed to create table: %w", err) } return nil } func DBQueryUsers() ([]User, error) { // List all users rows, err := db.Query("SELECT userID, username, type, createdAt 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, &user.CreatedAt); 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 DBQuerySpecificUser(userID int) (User, error) { // Query user by ID var user User err := db.QueryRow("SELECT userID, username, type, createdAt FROM userLogin WHERE userID = ?", userID).Scan(&user.UserID, &user.Username, &user.Type, &user.CreatedAt) if err != nil { return user, fmt.Errorf("failed to query user: %w", err) } return user, 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 DBQueryFrpcInstanceByID(instanceID int) (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 id = ?", instanceID).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 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 name = ?, serverAddr = ?, serverPort = ?, auth_method = ?, bootAtStart = ?, runUser = ?, configPath = ? WHERE id = ?", instance.Name, 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() ([]FrpcInstance, error) { rows, err := frpcDB.Query(` SELECT fi.id, fi.userID, fi.name, fi.serverAddr, fi.serverPort, fi.auth_method, fi.bootAtStart, fi.runUser, fi.configPath, fi.createdAt, u.username FROM frpcInstances fi JOIN userLogin u ON fi.userID = u.userID `) if err != nil { return nil, fmt.Errorf("failed to query all 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, &instance.CreatedBy); 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 }