package main import ( "flag" "fmt" "log" "os" "strings" "github.com/jfraeys/fetch_ml/internal/auth" "gopkg.in/yaml.v3" ) type ConfigWithAuth struct { Auth auth.AuthConfig `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 --cmd [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 [--admin] [--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) } // Generate API key apiKey := auth.GenerateAPIKey() // Setup user if config.Auth.APIKeys == nil { config.Auth.APIKeys = make(map[auth.Username]auth.APIKeyEntry) } 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 } } // Save user config.Auth.APIKeys[auth.Username(*username)] = auth.APIKeyEntry{ Hash: auth.APIKeyHash(auth.HashAPIKey(apiKey)), Admin: adminStatus, Roles: roles, Permissions: permissions, } 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 ") } hash := auth.HashAPIKey(*apiKey) fmt.Printf("Hash: %s\n", hash) 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) }