fetch_ml/cmd/user_manager/main.go

// Package main implements the fetch_ml user management CLI
package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"strings"

	"github.com/jfraeys/fetch_ml/internal/auth"
	"gopkg.in/yaml.v3"
)

// ConfigWithAuth wraps auth configuration for user management.
type ConfigWithAuth struct {
	Auth auth.Config `yaml:"auth"`
}

func main() {
	var (
		configFile = flag.String("config", "", "Configuration file path")
		command    = flag.String("cmd", "", "Command: generate-key, list-users, hash-key")
		username   = flag.String("username", "", "Username for generate-key")
		role       = flag.String("role", "", "Role for generate-key")
		admin      = flag.Bool("admin", false, "Admin flag for generate-key")
		apiKey     = flag.String("key", "", "API key to hash")
	)
	flag.Parse()

	if *configFile == "" || *command == "" {
		fmt.Println("Usage: user_manager --config <config.yaml> --cmd <command> [options]")
		fmt.Println("Commands: generate-key, list-users, hash-key")
		os.Exit(1)
	}

	switch *command {
	case "generate-key":
		if *username == "" {
			log.Fatal("Usage: --cmd generate-key --username <name> [--admin] [--role <role>]")
		}

		// Load config
		data, err := os.ReadFile(*configFile)
		if err != nil {
			log.Fatalf("Failed to read config: %v", err)
		}

		var config ConfigWithAuth
		if err := yaml.Unmarshal(data, &config); err != nil {
			log.Fatalf("Failed to parse config: %v", err)
		}

		// Determine admin status and roles
		adminStatus := *admin
		roles := []string{"viewer"}
		permissions := make(map[string]bool)

		if !adminStatus && *role == "" {
			fmt.Printf("Make user '%s' an admin? (y/N): ", *username)
			var response string
			_, _ = fmt.Scanln(&response)
			adminStatus = strings.ToLower(strings.TrimSpace(response)) == "y"
		}

		if adminStatus {
			roles = []string{"admin"}
			permissions["*"] = true
		} else if *role != "" {
			roles = []string{*role}
			rolePerms := getRolePermissions(*role)
			for perm, value := range rolePerms {
				permissions[perm] = value
			}
		}

		// Generate API key using new Argon2id method
		apiKey, apiKeyEntry, err := auth.GenerateNewAPIKey(adminStatus, roles, permissions)
		if err != nil {
			log.Fatalf("Failed to generate API key: %v", err)
		}

		// Setup user
		if config.Auth.APIKeys == nil {
			config.Auth.APIKeys = make(map[auth.Username]auth.APIKeyEntry)
		}

		config.Auth.APIKeys[auth.Username(*username)] = apiKeyEntry

		data, err = yaml.Marshal(config)
		if err != nil {
			log.Fatalf("Failed to marshal config: %v", err)
		}

		if err := os.WriteFile(*configFile, data, 0600); err != nil {
			log.Fatalf("Failed to write config: %v", err)
		}

		fmt.Printf("Generated API key for user '%s':\nKey: %s\n", *username, apiKey)

	case "list-users":
		data, err := os.ReadFile(*configFile)
		if err != nil {
			log.Fatalf("Failed to read config: %v", err)
		}

		var config ConfigWithAuth
		if err := yaml.Unmarshal(data, &config); err != nil {
			log.Fatalf("Failed to parse config: %v", err)
		}

		fmt.Println("Configured Users:")
		fmt.Println("=================")
		for username, entry := range config.Auth.APIKeys {
			fmt.Printf("User: %s\n", string(username))
			fmt.Printf("   Admin: %v\n", entry.Admin)
			if len(entry.Roles) > 0 {
				fmt.Printf("   Roles: %v\n", entry.Roles)
			}
			if len(entry.Permissions) > 0 {
				fmt.Printf("   Permissions: %d\n", len(entry.Permissions))
			}
			fmt.Printf("   Key Hash: %s...\n\n", string(entry.Hash)[:8])
		}

	case "hash-key":
		if *apiKey == "" {
			log.Fatal("Usage: --cmd hash-key --key <api-key>")
		}
		// Hash the provided key with Argon2id
		hashed, err := auth.HashAPIKeyArgon2id(*apiKey)
		if err != nil {
			log.Fatalf("Failed to hash key: %v", err)
		}
		fmt.Printf("Hash: %s\nSalt: %s\nAlgorithm: %s\n", hashed.Hash, hashed.Salt, hashed.Algorithm)

	default:
		log.Fatalf("Unknown command: %s", *command)
	}
}

// getRolePermissions returns permissions for a role
func getRolePermissions(role string) map[string]bool {
	rolePermissions := map[string]map[string]bool{
		"admin": {
			"*": true,
		},
		"data_scientist": {
			"jobs:create": true,
			"jobs:read":   true,
			"jobs:update": true,
			"data:read":   true,
			"models:read": true,
		},
		"data_engineer": {
			"data:create": true,
			"data:read":   true,
			"data:update": true,
			"data:delete": true,
		},
		"viewer": {
			"jobs:read":    true,
			"data:read":    true,
			"models:read":  true,
			"metrics:read": true,
		},
		"operator": {
			"jobs:read":    true,
			"jobs:update":  true,
			"metrics:read": true,
			"system:read":  true,
		},
	}

	if perms, exists := rolePermissions[role]; exists {
		return perms
	}
	return make(map[string]bool)
}