| 1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102 |
- package raft
- import (
- "bufio"
- "bytes"
- "fmt"
- "net"
- "os"
- "os/exec"
- "path/filepath"
- "regexp"
- "runtime"
- "sort"
- "strconv"
- "strings"
- "sync"
- "text/tabwriter"
- "time"
- )
- const (
- ColorReset = "\033[0m"
- ColorDim = "\033[90m" // Dark Gray
- ColorRed = "\033[31m"
- ColorGreen = "\033[32m"
- ColorYellow = "\033[33m"
- ColorBlue = "\033[34m"
- ColorCyan = "\033[36m"
- )
- // CLICommand defines a CLI command handler
- type CLICommand struct {
- Name string
- Description string
- Handler func(parts []string, server *KVServer)
- }
- // CLI manages the command line interface
- type CLI struct {
- server *KVServer
- commands map[string]CLICommand
- mu sync.RWMutex
- token string
- }
- // NewCLI creates a new CLI instance
- func NewCLI(server *KVServer) *CLI {
- cli := &CLI{
- server: server,
- commands: make(map[string]CLICommand),
- }
- cli.registerDefaultCommands()
- return cli
- }
- // Helper to calculate visible length of string ignoring ANSI codes
- var ansiRegex = regexp.MustCompile(`\x1b\[[0-9;]*m`)
- func visibleLen(s string) int {
- clean := ansiRegex.ReplaceAllString(s, "")
- return len([]rune(clean))
- }
- func printBoxed(content string) {
- lines := strings.Split(strings.TrimSpace(content), "\n")
- maxWidth := 0
- for _, line := range lines {
- l := visibleLen(line)
- if l > maxWidth {
- maxWidth = l
- }
- }
-
- // Min width to look decent
- if maxWidth < 20 {
- maxWidth = 20
- }
- // Add padding
- contentWidth := maxWidth + 2
-
- fmt.Println() // Start new line before box
- // Top Border
- fmt.Printf("%s╭%s╮%s\n", ColorDim, strings.Repeat("─", contentWidth), ColorReset)
-
- // Content
- for _, line := range lines {
- visLen := visibleLen(line)
- padding := contentWidth - visLen - 1 // -1 for the space after │
- // Ensure padding is not negative
- if padding < 0 { padding = 0 }
-
- fmt.Printf("%s│%s %s%s%s│%s\n", ColorDim, ColorReset, line, strings.Repeat(" ", padding), ColorDim, ColorReset)
- }
-
- // Bottom Border
- fmt.Printf("%s╰%s╯%s\n", ColorDim, strings.Repeat("─", contentWidth), ColorReset)
- fmt.Println() // End new line after box
- }
- func formatBytes(bytes int64) string {
- const unit = 1024
- if bytes < unit {
- return fmt.Sprintf("%d B", bytes)
- }
- div, exp := int64(unit), 0
- for n := bytes / unit; n >= unit; n /= unit {
- div *= unit
- exp++
- }
- return fmt.Sprintf("%.1f %cB", float64(bytes)/float64(div), "KMGTPE"[exp])
- }
- // scanCRLF is a split function for a Scanner that returns each line of
- // text, stripped of any trailing end-of-line marker. The returned line may
- // be empty. The end-of-line marker is one optional carriage return followed
- // by one mandatory newline, OR just a carriage return.
- func scanCRLF(data []byte, atEOF bool) (advance int, token []byte, err error) {
- if atEOF && len(data) == 0 {
- return 0, nil, nil
- }
- if i := bytes.IndexByte(data, '\n'); i >= 0 {
- // We have a full newline-terminated line.
- return i + 1, dropCR(data[0:i]), nil
- }
- if i := bytes.IndexByte(data, '\r'); i >= 0 {
- // We have a CR-terminated line.
- return i + 1, data[0:i], nil
- }
- // If we're at EOF, we have a final, non-terminated line. Return it.
- if atEOF {
- return len(data), dropCR(data), nil
- }
- // Request more data.
- return 0, nil, nil
- }
- // dropCR drops a terminal \r from the data.
- func dropCR(data []byte) []byte {
- if len(data) > 0 && data[len(data)-1] == '\r' {
- return data[0 : len(data)-1]
- }
- return data
- }
- // Start starts the command line interface
- func (c *CLI) Start() {
- // Attempt to fix terminal settings if they are messed up (e.g. missing ICRNL)
- // This helps when Enter sends \r but the terminal driver expects \n
- cmd := exec.Command("stty", "sane")
- cmd.Stdin = os.Stdin
- cmd.Stdout = os.Stdout
- _ = cmd.Run()
- // Boxed startup message
- var sb strings.Builder
- sb.WriteString(fmt.Sprintf("Node %s%s%s CLI Started\n", ColorGreen, c.server.Raft.nodeID, ColorReset))
- sb.WriteString("Type 'help' for commands.")
- printBoxed(sb.String())
- // Initialize history handling
- historyPath := filepath.Join(os.TempDir(), fmt.Sprintf(".raft_history_%s", c.server.Raft.nodeID))
- historyFile, err := os.OpenFile(historyPath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
- if err != nil {
- // Just warn, don't fail
- fmt.Printf("%sWarning: Could not open history file: %v%s\n", ColorYellow, err, ColorReset)
- }
- defer func() {
- if historyFile != nil {
- historyFile.Close()
- }
- }()
- // State Monitor Loop
- go func() {
- var lastState string
- var lastTerm uint64
- stats := c.server.GetStats()
- lastState = stats.State
- lastTerm = stats.Term
- ticker := time.NewTicker(100 * time.Millisecond)
- defer ticker.Stop()
- for {
- select {
- case <-c.server.stopCh:
- return
- case <-ticker.C:
- stats := c.server.GetStats()
- if stats.State != lastState || stats.Term != lastTerm {
- msg := fmt.Sprintf("%s[State Change] %s (Term %d) -> %s (Term %d)%s",
- ColorYellow, lastState, lastTerm, stats.State, stats.Term, ColorReset)
- fmt.Printf("\n%s\n> ", msg)
-
- lastState = stats.State
- lastTerm = stats.Term
- }
- }
- }
- }()
- scanner := bufio.NewScanner(os.Stdin)
- scanner.Split(scanCRLF)
- fmt.Print("> ")
- for scanner.Scan() {
- text := strings.TrimSpace(scanner.Text())
- if text == "" {
- fmt.Print("> ")
- continue
- }
- // Save command to history
- if historyFile != nil {
- if _, err := historyFile.WriteString(text + "\n"); err != nil {
- // Ignore write error
- }
- }
- parts := strings.Fields(text)
- cmdName := strings.ToLower(parts[0])
- c.mu.RLock()
- cmd, exists := c.commands[cmdName]
- c.mu.RUnlock()
- if exists {
- cmd.Handler(parts, c.server)
- } else {
- // Check aliases
- switch cmdName {
- case "binlog", "db", "stats":
- c.mu.RLock()
- if infoCmd, ok := c.commands["info"]; ok {
- infoCmd.Handler(parts, c.server)
- }
- c.mu.RUnlock()
- case "history":
- c.printHistory(historyPath)
- default:
- printBoxed("Unknown command. Type 'help'.")
- }
- }
- fmt.Print("> ")
- }
- }
- // printHistory prints the command history
- func (c *CLI) printHistory(path string) {
- content, err := os.ReadFile(path)
- if err != nil {
- printBoxed("No history found.")
- return
- }
-
- lines := strings.Split(string(content), "\n")
- // Show last 20 lines
- start := 0
- if len(lines) > 20 {
- start = len(lines) - 20
- }
-
- var sb strings.Builder
- sb.WriteString(fmt.Sprintf("%sCommand History (Last 20):%s\n", ColorCyan, ColorReset))
- for i := start; i < len(lines); i++ {
- line := strings.TrimSpace(lines[i])
- if line != "" {
- sb.WriteString(fmt.Sprintf("%d: %s\n", i+1, line))
- }
- }
- printBoxed(sb.String())
- }
- // RegisterCommand registers a new CLI command
- func (c *CLI) RegisterCommand(name, description string, handler func(parts []string, server *KVServer)) {
- c.mu.Lock()
- defer c.mu.Unlock()
- c.commands[strings.ToLower(name)] = CLICommand{
- Name: name,
- Description: description,
- Handler: handler,
- }
- }
- func (c *CLI) registerDefaultCommands() {
- c.RegisterCommand("set", "Set value (set <key> <val>)", func(parts []string, server *KVServer) {
- if len(parts) != 3 {
- printBoxed("Usage: set <key> <value>")
- return
- }
- key, val := parts[1], parts[2]
-
- // Use Authenticated Set
- c.mu.RLock()
- token := c.token
- c.mu.RUnlock()
- if err := server.SetAuthenticated(key, val, token); err != nil {
- printBoxed(fmt.Sprintf("%sError:%s %v", ColorRed, ColorReset, err))
- } else {
- printBoxed(fmt.Sprintf("%sOK%s", ColorGreen, ColorReset))
- }
- })
- c.RegisterCommand("get", "Get value (get <key>)", func(parts []string, server *KVServer) {
- if len(parts) != 2 {
- printBoxed("Usage: get <key>")
- return
- }
- key := parts[1]
-
- c.mu.RLock()
- token := c.token
- c.mu.RUnlock()
- if val, ok, err := server.GetLinearAuthenticated(key, token); err != nil {
- printBoxed(fmt.Sprintf("%sError:%s %v", ColorRed, ColorReset, err))
- } else if !ok {
- printBoxed(fmt.Sprintf("%sNot Found%s", ColorYellow, ColorReset))
- } else {
- printBoxed(fmt.Sprintf("%s%s%s = %s%s%s", ColorCyan, key, ColorReset, ColorYellow, val, ColorReset))
- }
- })
- c.RegisterCommand("del", "Delete key (del <key>)", func(parts []string, server *KVServer) {
- if len(parts) != 2 {
- printBoxed("Usage: del <key>")
- return
- }
- key := parts[1]
-
- c.mu.RLock()
- token := c.token
- c.mu.RUnlock()
- if err := server.DelAuthenticated(key, token); err != nil {
- printBoxed(fmt.Sprintf("%sError:%s %v", ColorRed, ColorReset, err))
- } else {
- printBoxed(fmt.Sprintf("%sDeleted%s", ColorGreen, ColorReset))
- }
- })
- // User Command Group
- c.RegisterCommand("user", "User management commands (type 'user help' for more)", func(parts []string, server *KVServer) {
- if len(parts) < 2 {
- printBoxed("Usage: user <subcommand> [args...]\nType 'user help' for list of subcommands.")
- return
- }
-
- subCmd := strings.ToLower(parts[1])
- args := parts[2:]
-
- // Common vars
- c.mu.RLock()
- token := c.token
- c.mu.RUnlock()
- switch subCmd {
- case "init": // Was auth-init
- if server.AuthManager.IsEnabled() {
- fmt.Printf("%sAuth system is already initialized.%s\n", ColorYellow, ColorReset)
- fmt.Print("Do you want to reset the root user? (y/N): ")
- var resp string
- fmt.Scanln(&resp)
- if strings.ToLower(resp) != "y" {
- return
- }
- }
-
- fmt.Print("Enter root password: ")
- password, err := readPassword()
- fmt.Println()
- if err != nil || password == "" {
- printBoxed("Error reading password")
- return
- }
- if err := server.AuthManager.CreateRootUser(password); err != nil {
- printBoxed(fmt.Sprintf("Failed to create root user: %v", err))
- return
- }
- if err := server.AuthManager.EnableAuth(); err != nil {
- printBoxed(fmt.Sprintf("Failed to enable auth: %v", err))
- return
- }
- printBoxed("Auth Initialized. Please login as 'root'.")
- case "login": // Was login
- if len(args) < 1 {
- printBoxed("Usage: user login <username> [code]")
- return
- }
- username := args[0]
- code := ""
- if len(args) > 1 {
- code = args[1]
- }
- fmt.Print("Password: ")
- password, err := readPassword()
- fmt.Println()
- if err != nil {
- printBoxed("Error reading password")
- return
- }
- t, err := server.AuthManager.Login(username, password, code, "cli")
- if err != nil && err.Error() == "mfa code required" {
- fmt.Printf("%sTwo-Factor Authentication Required%s\n", ColorYellow, ColorReset)
- fmt.Print("Enter MFA Code: ")
- code, err = readPassword()
- fmt.Println()
- if err == nil {
- t, err = server.AuthManager.Login(username, password, code, "cli")
- }
- }
- if err != nil {
- printBoxed(fmt.Sprintf("%sLogin Failed:%s %v", ColorRed, ColorReset, err))
- } else {
- c.mu.Lock()
- c.token = t
- c.mu.Unlock()
- printBoxed(fmt.Sprintf("%sLogin Success!%s", ColorGreen, ColorReset))
- }
- case "mfa":
- if len(args) < 1 {
- printBoxed("Usage: user mfa <enable|disable>")
- return
- }
- // Require Login
- c.mu.RLock()
- token := c.token
- c.mu.RUnlock()
- if token == "" {
- printBoxed("Error: You must be logged in.")
- return
- }
- switch args[0] {
- case "enable":
- // Generate
- secret, err := GenerateMFASecret()
- if err != nil {
- printBoxed(fmt.Sprintf("Error generating secret: %v", err))
- return
- }
-
- // Get Username
- session, err := server.AuthManager.GetSession(token)
- if err != nil {
- printBoxed("Session expired")
- return
- }
-
- url := GenerateOTPAuthURL(session.Username, secret, "RaftKV")
-
- // Show info
- fmt.Println()
- fmt.Printf("%sSetup 2FA%s\n", ColorCyan, ColorReset)
- fmt.Println("1. Install Google Authenticator or compatible app.")
- fmt.Println("2. Add account manually using this secret:")
- fmt.Printf(" %s%s%s\n", ColorGreen, secret, ColorReset)
- fmt.Printf(" (URL: %s)\n", url)
- fmt.Println("3. Enter the 6-digit code to verify.")
- fmt.Println()
-
- fmt.Print("Verification Code: ")
- code, err := readPassword()
- fmt.Println()
-
- if !ValidateTOTP(secret, code) {
- printBoxed(fmt.Sprintf("%sError: Invalid Code%s", ColorRed, ColorReset))
- return
- }
-
- if err := server.SetUserMFA(session.Username, secret, true, token); err != nil {
- printBoxed(fmt.Sprintf("Error saving user: %v", err))
- } else {
- printBoxed(fmt.Sprintf("%sMFA Enabled Successfully!%s", ColorGreen, ColorReset))
- }
- case "disable":
- fmt.Print("Are you sure you want to disable MFA? (y/N): ")
- var resp string
- fmt.Scanln(&resp)
- if strings.ToLower(resp) != "y" {
- return
- }
-
- session, err := server.AuthManager.GetSession(token)
- if err != nil {
- printBoxed("Session expired")
- return
- }
-
- if err := server.SetUserMFA(session.Username, "", false, token); err != nil {
- printBoxed(fmt.Sprintf("Error saving user: %v", err))
- } else {
- printBoxed(fmt.Sprintf("%sMFA Disabled%s", ColorGreen, ColorReset))
- }
- }
- case "logout": // Was logout
- if token != "" {
- server.Logout(token)
- c.mu.Lock()
- c.token = ""
- c.mu.Unlock()
- printBoxed("Logged out.")
- } else {
- printBoxed("Not logged in.")
- }
- case "who": // Was whoami
- if token == "" {
- printBoxed("Not logged in.")
- return
- }
- session, err := server.GetSessionInfo(token)
- if err != nil {
- printBoxed(fmt.Sprintf("Error: %v", err))
- return
- }
- printBoxed(fmt.Sprintf("Logged in as: %s%s%s", ColorCyan, session.Username, ColorReset))
- case "add": // Was user-add
- if len(args) < 1 {
- printBoxed("Usage: user add <username> [role1,role2...]")
- return
- }
- username := args[0]
- var roles []string
- if len(args) > 1 {
- roles = strings.Split(args[1], ",")
- }
- fmt.Printf("Enter password for %s: ", username)
- password, err := readPassword()
- fmt.Println()
- if err != nil || password == "" {
- printBoxed("Error reading password")
- return
- }
- if err := server.CreateUser(username, password, roles, token); err != nil {
- printBoxed(fmt.Sprintf("%sError:%s %v", ColorRed, ColorReset, err))
- } else {
- printBoxed(fmt.Sprintf("%sUser %s created%s", ColorGreen, username, ColorReset))
- }
- case "del": // Was user-del
- if len(args) < 1 {
- printBoxed("Usage: user del <username>")
- return
- }
- username := args[0]
- fmt.Printf("Delete user '%s'? (y/N): ", username)
- var resp string
- fmt.Scanln(&resp)
- if strings.ToLower(resp) != "y" {
- return
- }
- if err := server.DeleteUser(username, token); err != nil {
- printBoxed(fmt.Sprintf("%sError:%s %v", ColorRed, ColorReset, err))
- } else {
- printBoxed(fmt.Sprintf("%sUser deleted%s", ColorGreen, ColorReset))
- }
- case "list":
- users, err := server.ListUsers(token)
- if err != nil {
- printBoxed(fmt.Sprintf("%sError:%s %v", ColorRed, ColorReset, err))
- return
- }
- var buf bytes.Buffer
- w := tabwriter.NewWriter(&buf, 0, 0, 2, ' ', 0)
- fmt.Fprintln(w, "Username\tRoles\tStatus")
- fmt.Fprintln(w, "--------\t-----\t------")
- for _, u := range users {
- status := "Offline"
- if u.IsOnline { status = "Online" }
- fmt.Fprintf(w, "%s\t%s\t%s\n", u.Username, strings.Join(u.Roles, ","), status)
- }
- w.Flush()
- printBoxed(buf.String())
- case "update":
- // user update <username> [roles=r1,r2]
- // Currently only password update via prompt, maybe roles later
- // The prompt implies a generic update but also mentions "user update"
- // Let's implement password update here or allow changing roles.
- // Ideally arguments: user update <username> roles=admin,dev
- // If no args (other than username), prompt for password?
- if len(args) < 1 {
- printBoxed("Usage: user update <username> [roles=r1,r2...]")
- return
- }
- username := args[0]
-
- // If args present, update roles
- if len(args) > 1 {
- // Parse options
- updated := false
- // user, err := server.AuthManager.GetUser(username) // Use internal GetUser helper or logic?
- // We need authenticated way. server.ListUsers gives us access if admin.
- // But we need the User object to modify.
- // Let's rely on server.UpdateUser taking a struct.
- // We first need to GET the user to avoid overwriting fields.
- // Since we don't have GetUser exposed in server (only List), let's add one or iterate List.
- // Iterating List is inefficient but works for now.
- users, _ := server.ListUsers(token)
- var targetUser *User
- for _, u := range users {
- if u.Username == username {
- targetUser = u
- break
- }
- }
- if targetUser == nil {
- printBoxed("User not found or permission denied")
- return
- }
-
- // Apply updates
- for _, arg := range args[1:] {
- if strings.HasPrefix(arg, "roles=") {
- roleStr := strings.TrimPrefix(arg, "roles=")
- targetUser.Roles = strings.Split(roleStr, ",")
- updated = true
- }
- }
-
- if updated {
- if err := server.UpdateUser(*targetUser, token); err != nil {
- printBoxed(fmt.Sprintf("%sError:%s %v", ColorRed, ColorReset, err))
- } else {
- printBoxed(fmt.Sprintf("%sUser updated%s", ColorGreen, ColorReset))
- }
- }
- } else {
- // No args -> Change Password
- fmt.Printf("Enter new password for %s: ", username)
- pass, err := readPassword()
- fmt.Println()
- if err != nil || pass == "" {
- printBoxed("Error reading password")
- return
- }
- if err := server.ChangeUserPassword(username, pass, token); err != nil {
- printBoxed(fmt.Sprintf("%sError:%s %v", ColorRed, ColorReset, err))
- } else {
- printBoxed(fmt.Sprintf("%sPassword updated%s", ColorGreen, ColorReset))
- }
- }
- case "passwd": // Keep alias or subcommand
- if len(args) < 1 {
- printBoxed("Usage: user passwd <username>")
- return
- }
- // Reuse logic
- username := args[0]
- fmt.Printf("Enter new password for %s: ", username)
- pass, err := readPassword()
- fmt.Println()
- if err != nil || pass == "" {
- printBoxed("Error reading password")
- return
- }
- if err := server.ChangeUserPassword(username, pass, token); err != nil {
- printBoxed(fmt.Sprintf("%sError:%s %v", ColorRed, ColorReset, err))
- } else {
- printBoxed(fmt.Sprintf("%sPassword updated%s", ColorGreen, ColorReset))
- }
- case "role":
- if len(args) < 1 {
- printBoxed("Usage: user role <list|add|delete|update> ...")
- return
- }
- roleCmd := strings.ToLower(args[0])
- roleArgs := args[1:]
-
- switch roleCmd {
- case "list":
- roles, err := server.ListRoles(token)
- if err != nil {
- printBoxed(fmt.Sprintf("%sError:%s %v", ColorRed, ColorReset, err))
- return
- }
- var buf bytes.Buffer
- w := tabwriter.NewWriter(&buf, 0, 0, 2, ' ', 0)
- fmt.Fprintln(w, "Role\tParents\tPerms")
- fmt.Fprintln(w, "----\t-------\t-----")
- for _, r := range roles {
- parents := strings.Join(r.ParentRoles, ",")
- permCount := len(r.Permissions)
- fmt.Fprintf(w, "%s\t%s\t%d rules\n", r.Name, parents, permCount)
- }
- w.Flush()
- printBoxed(buf.String())
-
- case "add":
- if len(roleArgs) < 1 {
- printBoxed("Usage: user role add <name>")
- return
- }
- if err := server.CreateRole(roleArgs[0], token); err != nil {
- printBoxed(fmt.Sprintf("%sError:%s %v", ColorRed, ColorReset, err))
- } else {
- printBoxed(fmt.Sprintf("%sRole %s created%s", ColorGreen, roleArgs[0], ColorReset))
- }
-
- case "delete":
- if len(roleArgs) < 1 {
- printBoxed("Usage: user role delete <name>")
- return
- }
- if err := server.DeleteRole(roleArgs[0], token); err != nil {
- printBoxed(fmt.Sprintf("%sError:%s %v", ColorRed, ColorReset, err))
- } else {
- printBoxed(fmt.Sprintf("%sRole %s deleted%s", ColorGreen, roleArgs[0], ColorReset))
- }
-
- case "update":
- if len(roleArgs) < 1 {
- printBoxed("Usage: user role update <name> [parents=p1,p2] [perms=action:pattern,...]")
- return
- }
- name := roleArgs[0]
- roles, _ := server.ListRoles(token)
- var targetRole *Role
- for _, r := range roles {
- if r.Name == name {
- targetRole = r
- break
- }
- }
- if targetRole == nil {
- printBoxed("Role not found")
- return
- }
-
- updated := false
- for _, arg := range roleArgs[1:] {
- if strings.HasPrefix(arg, "parents=") {
- pStr := strings.TrimPrefix(arg, "parents=")
- targetRole.ParentRoles = strings.Split(pStr, ",")
- updated = true
- }
- if strings.HasPrefix(arg, "perms=") {
- permStr := strings.TrimPrefix(arg, "perms=")
- rules := strings.Split(permStr, ",")
- var newPerms []Permission
- for _, r := range rules {
- parts := strings.SplitN(r, ":", 2)
- if len(parts) == 2 {
- action := parts[0]
- pattern := parts[1]
- newPerms = append(newPerms, Permission{
- KeyPattern: pattern,
- Actions: []string{action},
- })
- }
- }
- targetRole.Permissions = newPerms
- updated = true
- }
- }
-
- if updated {
- if err := server.UpdateRole(*targetRole, token); err != nil {
- printBoxed(fmt.Sprintf("%sError:%s %v", ColorRed, ColorReset, err))
- } else {
- printBoxed(fmt.Sprintf("%sRole updated%s", ColorGreen, ColorReset))
- }
- } else {
- printBoxed("No changes specified.")
- }
-
- default:
- printBoxed("Unknown sub-command. Options: list, add, delete, update")
- }
- case "help":
- var sb strings.Builder
- sb.WriteString(fmt.Sprintf("%sUser Management Commands:%s\n", ColorCyan, ColorReset))
-
- // Define subcommands for formatting
- type subCmd struct {
- Name string
- Desc string
- }
- cmds := []subCmd{
- {"user init", "Initialize auth system (Root)"},
- {"user login <user> [code]", "Login"},
- {"user logout", "Logout"},
- {"user who", "Current user info"},
- {"user add <user> [roles]", "Create user (Root)"},
- {"user del <user>", "Delete user (Root)"},
- {"user update <user> [opts]", "Update user/pass (Root)"},
- {"user list", "List users (Root)"},
- {"user role list", "List roles (Root)"},
- {"user role add <name>", "Create role (Root)"},
- {"user role delete <name>", "Delete role (Root)"},
- {"user role update <name> ...", "Update role (Root)"},
- }
-
- // Find max length
- maxLen := 0
- for _, c := range cmds {
- if len(c.Name) > maxLen {
- maxLen = len(c.Name)
- }
- }
-
- for _, c := range cmds {
- padding := strings.Repeat(" ", maxLen-len(c.Name))
- sb.WriteString(fmt.Sprintf(" %s%s%s%s %s%s%s\n", ColorGreen, c.Name, ColorReset, padding, ColorDim, c.Desc, ColorReset))
- }
-
- printBoxed(sb.String())
- default:
- printBoxed(fmt.Sprintf("Unknown subcommand '%s'. Type 'user help'.", subCmd))
- }
- })
- c.RegisterCommand("join", "Add node to cluster (join <id> <addr>)", func(parts []string, server *KVServer) {
- if len(parts) != 3 {
- printBoxed("Usage: join <nodeID> <address>")
- return
- }
-
- // Admin only
- c.mu.RLock()
- token := c.token
- c.mu.RUnlock()
- if server.AuthManager.IsEnabled() && !server.IsRoot(token) {
- printBoxed(fmt.Sprintf("%sPermission Denied: Root access required%s", ColorRed, ColorReset))
- return
- }
- nodeID := parts[1]
- addr := parts[2]
- if err := server.Join(nodeID, addr); err != nil {
- printBoxed(fmt.Sprintf("%sError joining node:%s %v", ColorRed, ColorReset, err))
- } else {
- printBoxed(fmt.Sprintf("%sJoin request sent for node %s%s", ColorGreen, nodeID, ColorReset))
- }
- })
- c.RegisterCommand("leave", "Remove node from cluster (leave <id>)", func(parts []string, server *KVServer) {
- if len(parts) != 2 {
- printBoxed("Usage: leave <nodeID>")
- return
- }
- // Admin only
- c.mu.RLock()
- token := c.token
- c.mu.RUnlock()
- if server.AuthManager.IsEnabled() && !server.IsRoot(token) {
- printBoxed(fmt.Sprintf("%sPermission Denied: Root access required%s", ColorRed, ColorReset))
- return
- }
- nodeID := parts[1]
- if err := server.Leave(nodeID); err != nil {
- printBoxed(fmt.Sprintf("%sError leaving node:%s %v", ColorRed, ColorReset, err))
- } else {
- printBoxed(fmt.Sprintf("%sLeave request sent for node %s%s", ColorGreen, nodeID, ColorReset))
- }
- })
- c.RegisterCommand("search", "Search keys (search <pat> [lim] [off])", func(parts []string, server *KVServer) {
- if len(parts) < 2 {
- printBoxed("Usage: search <pattern> [limit] [offset]")
- return
- }
- pattern := parts[1]
- limit := 20 // Default limit
- offset := 0
-
- if len(parts) >= 3 {
- if l, err := strconv.Atoi(parts[2]); err == nil {
- limit = l
- }
- }
- if len(parts) >= 4 {
- if o, err := strconv.Atoi(parts[3]); err == nil {
- offset = o
- }
- }
- c.mu.RLock()
- token := c.token
- c.mu.RUnlock()
- // Use Authenticated Search
- results, err := server.SearchAuthenticated(pattern, limit, offset, token)
- if err != nil {
- printBoxed(fmt.Sprintf("%sError:%s %v", ColorRed, ColorReset, err))
- return
- }
- // Format Table to Buffer
- var buf bytes.Buffer
- w := tabwriter.NewWriter(&buf, 0, 0, 2, ' ', 0)
- fmt.Fprintln(w, "Key\tValue\tCommitIndex")
- fmt.Fprintln(w, "---\t-----\t-----------")
- for _, r := range results {
- fmt.Fprintf(w, "%s\t%s\t%d\n", r.Key, r.Value, r.CommitIndex)
- }
- w.Flush()
-
- summary := fmt.Sprintf("\n%sFound %d records (showing max %d)%s", ColorDim, len(results), limit, ColorReset)
- printBoxed(buf.String() + summary)
- })
- c.RegisterCommand("count", "Count keys (count <pat>)", func(parts []string, server *KVServer) {
- if len(parts) < 2 {
- printBoxed("Usage: count <pattern> (e.g. 'count *', 'count user.*')")
- return
- }
- pattern := parts[1]
-
- c.mu.RLock()
- token := c.token
- c.mu.RUnlock()
-
- start := time.Now()
- count, err := server.CountAuthenticated(pattern, token)
- duration := time.Since(start)
-
- if err != nil {
- printBoxed(fmt.Sprintf("%sError:%s %v", ColorRed, ColorReset, err))
- } else {
- printBoxed(fmt.Sprintf("%sCount: %d%s (took %v)", ColorGreen, count, ColorReset, duration))
- }
- })
- c.RegisterCommand("info", "Show all system stats", func(parts []string, server *KVServer) {
- // Admin only
- c.mu.RLock()
- token := c.token
- c.mu.RUnlock()
- if server.AuthManager.IsEnabled() && !server.IsRoot(token) {
- printBoxed(fmt.Sprintf("%sPermission Denied: Root access required%s", ColorRed, ColorReset))
- return
- }
- // 1. Gather all data
- var m runtime.MemStats
- runtime.ReadMemStats(&m)
-
- dbSize := server.GetDBSize()
- logSize := server.GetLogSize()
-
- stats := server.GetStats()
- health := server.HealthCheck()
- dbApplied := server.DB.GetLastAppliedIndex()
- // 2. Build Output
- var sb strings.Builder
-
- // Section: Node Status
- sb.WriteString(fmt.Sprintf("%sNode Status:%s\n", ColorCyan, ColorReset))
- sb.WriteString(fmt.Sprintf(" ID: %s%s%s\n", ColorGreen, health.NodeID, ColorReset))
- sb.WriteString(fmt.Sprintf(" State: %s\n", health.State))
- sb.WriteString(fmt.Sprintf(" Term: %d\n", health.Term))
- sb.WriteString(fmt.Sprintf(" Leader: %s\n", health.LeaderID))
- sb.WriteString(fmt.Sprintf(" Healthy: %v\n", health.IsHealthy))
- sb.WriteString("\n")
- // Section: Storage Stats
- sb.WriteString(fmt.Sprintf("%sStorage & Memory:%s\n", ColorCyan, ColorReset))
- sb.WriteString(fmt.Sprintf(" DB Size: %s\n", formatBytes(dbSize)))
- sb.WriteString(fmt.Sprintf(" Raft Log: %s\n", formatBytes(logSize)))
- sb.WriteString(fmt.Sprintf(" Mem Alloc: %s\n", formatBytes(int64(m.Alloc))))
- sb.WriteString(fmt.Sprintf(" Mem Sys: %s\n", formatBytes(int64(m.Sys))))
- sb.WriteString(fmt.Sprintf(" NumGC: %d\n", m.NumGC))
- sb.WriteString("\n")
- // Section: Raft & DB Indices
- sb.WriteString(fmt.Sprintf("%sConsensus Indices:%s\n", ColorCyan, ColorReset))
- sb.WriteString(fmt.Sprintf(" Commit Index: %d\n", stats.CommitIndex))
- sb.WriteString(fmt.Sprintf(" Applied Index:%d\n", stats.LastApplied))
- sb.WriteString(fmt.Sprintf(" Last Log Idx: %d\n", stats.LastLogIndex))
- sb.WriteString(fmt.Sprintf(" DB Applied: %d\n", dbApplied))
- sb.WriteString("\n")
- // Section: Cluster Members
- sb.WriteString(fmt.Sprintf("%sCluster Members (%d):%s", ColorCyan, stats.ClusterSize, ColorReset))
-
- for id, addr := range stats.ClusterNodes {
- // Determine status color
- statusColor := ColorDim
- nodeID := server.Raft.nodeID // Assuming local nodeID is accessible
- if id == nodeID {
- statusColor = ColorCyan // Self
- } else {
- // Quick check
- conn, err := net.DialTimeout("tcp", addr, 50*time.Millisecond)
- if err == nil {
- statusColor = ColorGreen
- conn.Close()
- }
- }
-
- // Marker for Leader
- marker := ""
- if id == health.LeaderID {
- marker = fmt.Sprintf(" %s(Leader)%s", ColorYellow, ColorReset)
- }
-
- sb.WriteString(fmt.Sprintf("\n - %s%s%s: %s%s", statusColor, id, ColorReset, addr, marker))
- }
-
- printBoxed(sb.String())
- })
- c.RegisterCommand("quit", "Gracefully shutdown the node", func(parts []string, server *KVServer) {
- // Admin only
- c.mu.RLock()
- token := c.token
- c.mu.RUnlock()
- if server.AuthManager.IsEnabled() && !server.IsRoot(token) {
- printBoxed(fmt.Sprintf("%sPermission Denied: Root access required%s", ColorRed, ColorReset))
- return
- }
- printBoxed("Shutting down node...")
- if err := server.Stop(); err != nil {
- printBoxed(fmt.Sprintf("%sError stopping server:%s %v", ColorRed, ColorReset, err))
- } else {
- printBoxed(fmt.Sprintf("%sNode stopped successfully.%s", ColorGreen, ColorReset))
- }
- os.Exit(0)
- })
- c.RegisterCommand("clear", "Clear screen", func(parts []string, server *KVServer) {
- fmt.Print("\033[H\033[2J")
- })
- c.RegisterCommand("help", "Show commands", func(parts []string, server *KVServer) {
- var sb strings.Builder
- sb.WriteString(fmt.Sprintf("%sCommands:%s\n", ColorGreen, ColorReset))
-
- // Sort commands
- c.mu.RLock()
- var cmds []CLICommand
- for _, cmd := range c.commands {
- cmds = append(cmds, cmd)
- }
- c.mu.RUnlock()
-
- sort.Slice(cmds, func(i, j int) bool {
- return cmds[i].Name < cmds[j].Name
- })
- // Find max length for padding
- maxNameLen := 0
- for _, cmd := range cmds {
- if len(cmd.Name) > maxNameLen {
- maxNameLen = len(cmd.Name)
- }
- }
- for _, cmd := range cmds {
- padding := strings.Repeat(" ", maxNameLen-len(cmd.Name))
- sb.WriteString(fmt.Sprintf(" %s%s%s%s %s%s%s\n", ColorGreen, cmd.Name, ColorReset, padding, ColorDim, cmd.Description, ColorReset))
- }
-
- printBoxed(sb.String())
- })
- }
- // readPassword securely reads a password from stdin without echoing.
- func readPassword() (string, error) {
- // 1. Disable echo
- cmd := exec.Command("stty", "-echo")
- cmd.Stdin = os.Stdin
- cmd.Stdout = os.Stdout
- if err := cmd.Run(); err != nil {
- return "", fmt.Errorf("failed to disable echo: %v", err)
- }
- defer func() {
- // 3. Re-enable echo
- cmd := exec.Command("stty", "echo")
- cmd.Stdin = os.Stdin
- cmd.Stdout = os.Stdout
- _ = cmd.Run()
- }()
- // 2. Read input
- reader := bufio.NewReader(os.Stdin)
- pass, err := reader.ReadString('\n')
- if err != nil {
- return "", err
- }
- return strings.TrimSpace(pass), nil
- }
|