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 )", func(parts []string, server *KVServer) { if len(parts) != 3 { printBoxed("Usage: set ") 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 )", func(parts []string, server *KVServer) { if len(parts) != 2 { printBoxed("Usage: get ") 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 )", func(parts []string, server *KVServer) { if len(parts) != 2 { printBoxed("Usage: del ") 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)) } }) c.RegisterCommand("login", "Login to system (login [code])", func(parts []string, server *KVServer) { if len(parts) < 3 { printBoxed("Usage: login [mfa_code]") return } username := parts[1] password := parts[2] code := "" if len(parts) > 3 { code = parts[3] } token, 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 = token c.mu.Unlock() printBoxed(fmt.Sprintf("%sLogin Success! Token saved.%s", ColorGreen, ColorReset)) } }) c.RegisterCommand("logout", "Logout from system", func(parts []string, server *KVServer) { c.mu.Lock() c.token = "" c.mu.Unlock() printBoxed("Logged out.") }) c.RegisterCommand("auth-init", "Initialize Auth System (auth-init )", func(parts []string, server *KVServer) { if len(parts) != 2 { printBoxed("Usage: auth-init ") return } // 1. Create Root User if err := server.AuthManager.CreateRootUser(parts[1]); err != nil { printBoxed(fmt.Sprintf("Failed to create root user: %v", err)) return } // 2. Enable Auth if err := server.AuthManager.EnableAuth(); err != nil { printBoxed(fmt.Sprintf("Failed to enable auth: %v", err)) return } printBoxed("Auth Initialized. Please login as 'root'.") }) c.RegisterCommand("whoami", "Show current user info", func(parts []string, server *KVServer) { c.mu.RLock() token := c.token c.mu.RUnlock() if token == "" { printBoxed("Not logged in.") return } // Decode token or check session // Accessing AuthManager internal session map is protected. // We can add a helper in AuthManager to get session info or just show token. printBoxed(fmt.Sprintf("Logged in (Token: %s...)", token[:8])) }) c.RegisterCommand("join", "Add node to cluster (join )", func(parts []string, server *KVServer) { if len(parts) != 3 { printBoxed("Usage: join
") 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 )", func(parts []string, server *KVServer) { if len(parts) != 2 { printBoxed("Usage: leave ") 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 [lim] [off])", func(parts []string, server *KVServer) { if len(parts) < 2 { printBoxed("Usage: search [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 } } // Construct SQL for DB Engine sql := fmt.Sprintf("key like \"%s\" LIMIT %d OFFSET %d", pattern, limit, offset) results, err := server.DB.Query(sql) 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 )", func(parts []string, server *KVServer) { if len(parts) < 2 { printBoxed("Usage: count (e.g. 'count *', 'count user.*')") return } pattern := parts[1] sql := "" if pattern == "*" { sql = "*" } else { sql = fmt.Sprintf("key like \"%s\"", pattern) } start := time.Now() count, err := server.DB.Count(sql) 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) { // 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) { 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()) }) }