createMongoUser.go

/*
export MONGO_URI="mongodb://admin:passwd@localhost:27017"
export USER_YAML="filename"
createMongoUsers

yaml-file like this

users:
  - username: exampleUser
    password: examplePassword
    roles:
      - role: readAnyDatabase
      - role: readWrite
        db: exampleDB

*/

package main

import (
	"context"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"time"

	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
	"gopkg.in/yaml.v3"
)

type Role struct {
	Role string `yaml:"role" bson:"role"`
	DB   string `yaml:"db,omitempty" bson:"db"`
}

type User struct {
	Username string `yaml:"username"`
	Password string `yaml:"password"`
	Roles    []Role `yaml:"roles"`
}

type UserList struct {
	Users []User `yaml:"users"`
}

func normalizeRoles(roles []Role) []Role {
	normalized := make([]Role, len(roles))
	for i, role := range roles {
		if role.DB == "" {
			role.DB = "admin"
		}
		normalized[i] = role
	}
	return normalized
}

func syncUserRoles(ctx context.Context, client *mongo.Client, username, password string, targetRoles []Role) error {
	adminDB := client.Database("admin")
	targetRoles = normalizeRoles(targetRoles)

	result := adminDB.RunCommand(ctx, bson.D{{Key: "usersInfo", Value: username}})
	var userInfo bson.M
	if err := result.Decode(&userInfo); err != nil {
		return fmt.Errorf("failed to fetch user info: %v", err)
	}

	users, ok := userInfo["users"].(bson.A)
	if !ok || len(users) == 0 {
		createCmd := bson.D{
			{Key: "createUser", Value: username},
			{Key: "pwd", Value: password},
			{Key: "roles", Value: targetRoles},
		}
		if err := adminDB.RunCommand(ctx, createCmd).Err(); err != nil {
			return fmt.Errorf("failed to create user: %v", err)
		}
		fmt.Printf("User %s created with roles: %v\n", username, targetRoles)
		return nil
	}

	user := users[0].(bson.M)
	currentRolesRaw, _ := user["roles"].(bson.A)
	currentRoles := make([]Role, 0, len(currentRolesRaw))
	for _, r := range currentRolesRaw {
		roleMap := r.(bson.M)
		currentRoles = append(currentRoles, Role{
			Role: roleMap["role"].(string),
			DB:   roleMap["db"].(string),
		})
	}

	missingRoles := []Role{}
	existingRolesMap := map[string]bool{}
	for _, role := range currentRoles {
		key := fmt.Sprintf("%s@%s", role.Role, role.DB)
		existingRolesMap[key] = true
	}
	for _, role := range targetRoles {
		key := fmt.Sprintf("%s@%s", role.Role, role.DB)
		if !existingRolesMap[key] {
			missingRoles = append(missingRoles, role)
		}
	}

	requiredRolesMap := map[string]bool{}
	for _, role := range targetRoles {
		key := fmt.Sprintf("%s@%s", role.Role, role.DB)
		requiredRolesMap[key] = true
	}
	extraRoles := []Role{}
	for _, role := range currentRoles {
		key := fmt.Sprintf("%s@%s", role.Role, role.DB)
		if !requiredRolesMap[key] {
			extraRoles = append(extraRoles, role)
		}
	}

	if len(missingRoles) > 0 {
		grantCmd := bson.D{
			{Key: "grantRolesToUser", Value: username},
			{Key: "roles", Value: missingRoles},
		}
		if err := adminDB.RunCommand(ctx, grantCmd).Err(); err != nil {
			return fmt.Errorf("failed to grant roles: %v", err)
		}
		fmt.Printf("Granted roles to %s: %v\n", username, missingRoles)
	}

	if len(extraRoles) > 0 {
		revokeCmd := bson.D{
			{Key: "revokeRolesFromUser", Value: username},
			{Key: "roles", Value: extraRoles},
		}
		if err := adminDB.RunCommand(ctx, revokeCmd).Err(); err != nil {
			return fmt.Errorf("failed to revoke roles: %v", err)
		}
		fmt.Printf("Revoked roles from %s: %v\n", username, extraRoles)
	}

	if len(missingRoles) == 0 && len(extraRoles) == 0 {
		fmt.Printf("User %s already has the correct roles. No changes needed.\n", username)
	}

	return nil
}

func loadUsersFromYAML(filePath string) ([]User, error) {
	data, err := ioutil.ReadFile(filePath)
	if err != nil {
		return nil, fmt.Errorf("failed to read YAML file: %v", err)
	}

	var userList UserList
	if err := yaml.Unmarshal(data, &userList); err != nil {
		return nil, fmt.Errorf("failed to parse YAML: %v", err)
	}

	return userList.Users, nil
}

func main() {
	mongoURI := os.Getenv("MONGO_URI")
	if mongoURI == "" {
		log.Fatal("MONGO_URI environment variable is required")
	}

	yamlPath := os.Getenv("USER_YAML")
	if yamlPath == "" {
		log.Fatal("USER_YAML environment variable is required")
	}

	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
	defer cancel()

	client, err := mongo.Connect(ctx, options.Client().ApplyURI(mongoURI))
	if err != nil {
		log.Fatalf("Failed to connect to MongoDB: %v", err)
	}
	defer client.Disconnect(ctx)

	users, err := loadUsersFromYAML(yamlPath)
	if err != nil {
		log.Fatalf("Failed to load users: %v", err)
	}

	for _, user := range users {
		err := syncUserRoles(ctx, client, user.Username, user.Password, user.Roles)
		if err != nil {
			log.Printf("Error syncing user %s: %v", user.Username, err)
		}
	}
}