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 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, 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 DBUpdateUserType (userID int, newType string) error { _, err := db.Exec("UPDATE userLogin SET type = ? WHERE userID = ?", newType, userID) if err != nil { return fmt.Errorf("failed to update user type: %w", err) } return nil } func DBAddFrpcInstance(instance FrpcInstance) error { _, err := frpcDB.Exec("INSERT INTO frpcInstances (userID, name, bootAtStart, runUser, configPath, createdAt) VALUES (?, ?, ?, ?, ?, ?)", instance.UserID, instance.Name, 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, bootAtStart, runUser, configPath, createdAt FROM frpcInstances WHERE id = ?", instanceID).Scan( &instance.ID, &instance.UserID, &instance.Name, &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, bootAtStart, runUser, configPath, createdAt FROM frpcInstances WHERE userID = ? AND name = ?", userID, instanceName).Scan( &instance.ID, &instance.UserID, &instance.Name, &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 DBRemoveFrpcInstanceByID(instanceID int) error { _, err := frpcDB.Exec("DELETE FROM frpcInstances WHERE id = ?", instanceID) 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 bootAtStart = ?, runUser = ?, configPath = ? WHERE id = ?", 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.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.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 } func GetServiceNameByInstanceID(instanceID int) (string, error) { instance, err := DBQueryFrpcInstanceByID(instanceID) if err != nil { return "", fmt.Errorf("failed to query frpc instance: %w", err) } user, err := GetUserByID(instance.UserID) if err != nil { return "", fmt.Errorf("failed to get user info: %w", err) } serviceName := fmt.Sprintf("superfrpc_%s_%s", user.Username, instance.Name) postLog.Debug(fmt.Sprintf("[GetServiceNameByInstanceID] instanceID: %d, serviceName: %s", instanceID, serviceName)) return serviceName, nil }