Modularized operator code, so it does not look like a garbage
This commit is contained in:
Pavlo Khazov
@@ -215,107 +215,94 @@ func getCompleter() *readline.PrefixCompleter {
|
||||
return readline.NewPrefixCompleter(items...)
|
||||
}
|
||||
|
||||
func main() {
|
||||
// Define flags for username, password and address of server
|
||||
operatorPassword := flag.String("p", "", "Password")
|
||||
operatorID := flag.String("u", "", "Username")
|
||||
serverAddr := flag.String("ip", defaultServerAddr, "IP and port of server")
|
||||
flag.Parse()
|
||||
|
||||
// Check if password and username are provided
|
||||
if *operatorPassword == "" || *operatorID == "" {
|
||||
fmt.Println("Usage: client.exe -u <username> -p <password> -ip <ip:port>")
|
||||
return
|
||||
}
|
||||
fmt.Println("Client's username and password provided")
|
||||
|
||||
// Prepares the TLS configuration for secure connection
|
||||
func setupTLSConfig() (*tls.Config, error) {
|
||||
// Load CA certificate
|
||||
caCert, err := os.ReadFile(caCertPath)
|
||||
if err != nil {
|
||||
log.Fatalf("Failed to read CA certificate: %v", err)
|
||||
return nil, fmt.Errorf("failed to read CA certificate: %v", err)
|
||||
}
|
||||
|
||||
caCertPool := x509.NewCertPool()
|
||||
if !caCertPool.AppendCertsFromPEM(caCert) {
|
||||
log.Fatal("Failed to append CA certificate to pool")
|
||||
return nil, fmt.Errorf("failed to append CA certificate to pool")
|
||||
}
|
||||
|
||||
// Load client certificate and key
|
||||
clientCert, err := tls.LoadX509KeyPair(clientCertPath, clientKeyPath)
|
||||
if err != nil {
|
||||
log.Fatalf("Failed to load client certificate and key: %v", err)
|
||||
return nil, fmt.Errorf("failed to load client certificate and key: %v", err)
|
||||
}
|
||||
|
||||
// Configure TLS
|
||||
tlsConfig := &tls.Config{
|
||||
return &tls.Config{
|
||||
Certificates: []tls.Certificate{clientCert},
|
||||
RootCAs: caCertPool,
|
||||
InsecureSkipVerify: false, // Ensure server's certificate is validated
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Establishes a secure connection to the server
|
||||
func connectToServer(serverAddr string, tlsConfig *tls.Config) (*tls.Conn, error) {
|
||||
conn, err := tls.Dial("tcp", serverAddr, tlsConfig)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error connecting to server: %v", err)
|
||||
}
|
||||
|
||||
// Connect to the server
|
||||
conn, err := tls.Dial("tcp", *serverAddr, tlsConfig)
|
||||
if err != nil {
|
||||
fmt.Println("Error connecting to server:", err)
|
||||
return
|
||||
}
|
||||
defer conn.Close()
|
||||
fmt.Println("Connected to server")
|
||||
return conn, nil
|
||||
}
|
||||
|
||||
// Authentication
|
||||
// Pretty simple authentication of user with the server
|
||||
func authenticateUser(conn *tls.Conn, userID, password string) error {
|
||||
fmt.Println("Starting authentication")
|
||||
authMessage := fmt.Sprintf("AUTH:%s:%s", *operatorID, *operatorPassword)
|
||||
_, err = conn.Write([]byte(authMessage))
|
||||
|
||||
// Send authentication message
|
||||
authMessage := fmt.Sprintf("AUTH:%s:%s", userID, password)
|
||||
_, err := conn.Write([]byte(authMessage))
|
||||
if err != nil {
|
||||
fmt.Println("Error sending credentials:", err)
|
||||
return
|
||||
return fmt.Errorf("error sending credentials: %v", err)
|
||||
}
|
||||
|
||||
// Receive server response
|
||||
buffer := make([]byte, 4096)
|
||||
n, err := conn.Read(buffer)
|
||||
if err != nil {
|
||||
fmt.Println("Error reading from server:", err)
|
||||
return
|
||||
return fmt.Errorf("error reading from server: %v", err)
|
||||
}
|
||||
response := string(buffer[:n])
|
||||
|
||||
// Parse response message
|
||||
response := string(buffer[:n])
|
||||
if response == "AUTH_OK" {
|
||||
fmt.Println("Auth successful. You can now issue commands.")
|
||||
} else {
|
||||
fmt.Println("Auth failed.")
|
||||
return
|
||||
return nil
|
||||
}
|
||||
|
||||
// Create readline config with more explicit settings
|
||||
rlConfig := &readline.Config{
|
||||
Prompt: "\033[31mSigma >\033[0m ", // Colored prompt
|
||||
HistoryFile: "/tmp/operator_history.txt",
|
||||
AutoComplete: getCompleter(),
|
||||
InterruptPrompt: "^C",
|
||||
// EOFPrompt: "exit",
|
||||
return fmt.Errorf("authentication failed")
|
||||
}
|
||||
|
||||
// Enable tab completion settings
|
||||
// Configure readline interface
|
||||
func setupReadline() (*readline.Instance, error) {
|
||||
rlConfig := &readline.Config{
|
||||
Prompt: "\033[31mSigma >\033[0m ", // Colored prompt
|
||||
HistoryFile: "/tmp/operator_history.txt",
|
||||
AutoComplete: getCompleter(),
|
||||
InterruptPrompt: "^C",
|
||||
DisableAutoSaveHistory: false,
|
||||
HistorySearchFold: true,
|
||||
|
||||
// Add filters for better handling of input
|
||||
UniqueEditLine: true,
|
||||
UniqueEditLine: true,
|
||||
}
|
||||
|
||||
rl, err := readline.NewEx(rlConfig)
|
||||
if err != nil {
|
||||
fmt.Printf("Error initializing readline: %v\n", err)
|
||||
return
|
||||
return nil, fmt.Errorf("error initializing readline: %v", err)
|
||||
}
|
||||
defer rl.Close()
|
||||
|
||||
// Set up better handling of terminal
|
||||
rl.SetVimMode(false)
|
||||
return rl, nil
|
||||
}
|
||||
|
||||
exitChan := make(chan struct{})
|
||||
|
||||
// Goroutine for receiving messages from the server
|
||||
// Goroutine to listen for server messages
|
||||
func startServerListener(conn *tls.Conn, rl *readline.Instance, exitChan chan struct{}) {
|
||||
go func() {
|
||||
for {
|
||||
serverMsg := make([]byte, 4096)
|
||||
@@ -325,46 +312,137 @@ func main() {
|
||||
close(exitChan)
|
||||
return
|
||||
}
|
||||
|
||||
// Parse the server message
|
||||
message := string(serverMsg[:n])
|
||||
if strings.HasPrefix(message, "LISTENERS_UPDATE:") {
|
||||
// Extract the listener names from the message
|
||||
listeners := strings.Split(strings.TrimPrefix(message, "LISTENERS_UPDATE:"), ",")
|
||||
|
||||
// Update subcommands
|
||||
UpdateListenersSubCommands(listeners)
|
||||
// message = fmt.Sprintf("Listeners list updated: %s", listeners)
|
||||
// fmt.Print("\r\033[K") // Clears the line and the prompt
|
||||
// fmt.Println("\n" + message + "\n")
|
||||
|
||||
// Rebuild the tab completion
|
||||
rl.Config.AutoComplete = getCompleter()
|
||||
// Refresh readline input
|
||||
rl.Refresh()
|
||||
} else if strings.HasPrefix(message, "AGENTS_UPDATE:") {
|
||||
// Extract the agent names from the message
|
||||
agents = strings.Split(strings.TrimPrefix(message, "AGENTS_UPDATE:"), ",")
|
||||
|
||||
// Update subcommands
|
||||
UpdateAgentsSubCommands(agents)
|
||||
// message = fmt.Sprintf("Agents list updated: %s", agents)
|
||||
// fmt.Print("\r\033[K") // Clears the line and the prompt
|
||||
// fmt.Println("\n" + message + "\n")
|
||||
|
||||
// Rebuild the tab completion
|
||||
rl.Config.AutoComplete = getCompleter()
|
||||
// Refresh readline input
|
||||
rl.Refresh()
|
||||
} else {
|
||||
// Handle other messages
|
||||
fmt.Print("\r\033[K") // Clears the line and the prompt
|
||||
fmt.Println("\n" + message)
|
||||
rl.Refresh() // Refresh readline input
|
||||
}
|
||||
handleServerMessage(message, rl)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// Interactive terminal loop
|
||||
// Processes messages received from the server
|
||||
func handleServerMessage(message string, rl *readline.Instance) {
|
||||
// Tab-completion update with listeners
|
||||
if strings.HasPrefix(message, "LISTENERS_UPDATE:") {
|
||||
// Extract the listener names from the message
|
||||
listeners := strings.Split(strings.TrimPrefix(message, "LISTENERS_UPDATE:"), ",")
|
||||
|
||||
// Update subcommands
|
||||
UpdateListenersSubCommands(listeners)
|
||||
|
||||
// Rebuild the tab completion
|
||||
rl.Config.AutoComplete = getCompleter()
|
||||
rl.Refresh()
|
||||
|
||||
// Tab-completion update with new agents
|
||||
} else if strings.HasPrefix(message, "AGENTS_UPDATE:") {
|
||||
// Extract the agent names from the message
|
||||
agents = strings.Split(strings.TrimPrefix(message, "AGENTS_UPDATE:"), ",")
|
||||
|
||||
// Update subcommands
|
||||
UpdateAgentsSubCommands(agents)
|
||||
|
||||
// Rebuild the tab completion
|
||||
rl.Config.AutoComplete = getCompleter()
|
||||
rl.Refresh()
|
||||
|
||||
// Handle usual messages
|
||||
} else {
|
||||
fmt.Print("\r\033[K") // Clears the line and the prompt
|
||||
fmt.Println("\n" + message)
|
||||
rl.Refresh() // Refresh readline input
|
||||
}
|
||||
}
|
||||
|
||||
// Handles operator commands and context switching
|
||||
func processCommand(command string, conn *tls.Conn, rl *readline.Instance) bool {
|
||||
if command == "" {
|
||||
return false
|
||||
}
|
||||
|
||||
// Handle exit command
|
||||
if command == "exit" {
|
||||
if currentAgentContext != "" {
|
||||
// Exit from agent context
|
||||
currentAgentContext = ""
|
||||
rl.SetPrompt("\033[31mSigma >\033[0m ")
|
||||
|
||||
// Update tab completion after context change
|
||||
rl.Config.AutoComplete = getCompleter()
|
||||
return false
|
||||
} else {
|
||||
// Exit from application
|
||||
fmt.Println("Exiting...")
|
||||
conn.Write([]byte(command))
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
// Handle interact command
|
||||
if strings.HasPrefix(command, "interact") {
|
||||
handleInteractCommand(command, rl)
|
||||
return false
|
||||
}
|
||||
|
||||
// Handle command in agent context
|
||||
if currentAgentContext != "" && isControlCommand(command) {
|
||||
command = currentAgentContext + " " + command
|
||||
}
|
||||
|
||||
// Send command to server
|
||||
_, err := conn.Write([]byte(command))
|
||||
if err != nil {
|
||||
fmt.Printf("Error sending command: %v", err)
|
||||
return true
|
||||
}
|
||||
|
||||
return false
|
||||
}
|
||||
|
||||
// Check if command is a task for agent
|
||||
func isControlCommand(command string) bool {
|
||||
// Split the command into words
|
||||
parts := strings.Fields(command)
|
||||
if len(parts) == 0 {
|
||||
fmt.Println("Empty command :(")
|
||||
return false
|
||||
}
|
||||
|
||||
// Check if the first word is a context command
|
||||
if slices.Contains(contextCommands, parts[0]) {
|
||||
return true
|
||||
} else {
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// Processes the interact command for agent context switching
|
||||
func handleInteractCommand(command string, rl *readline.Instance) {
|
||||
cmd := strings.Fields(command)
|
||||
|
||||
if len(cmd) != 2 {
|
||||
fmt.Println("Usage: interact <agent_id>")
|
||||
return
|
||||
}
|
||||
|
||||
if !slices.Contains(agents, cmd[1]) {
|
||||
fmt.Println("Agent not found!")
|
||||
return
|
||||
}
|
||||
|
||||
// Set new agent context
|
||||
currentAgentContext = cmd[1]
|
||||
|
||||
// Update tab completion
|
||||
rl.Config.AutoComplete = getCompleter()
|
||||
|
||||
// Update prompt with agent context
|
||||
newPrompt := fmt.Sprintf("\033[31mSigma [%s] >\033[0m ", currentAgentContext)
|
||||
rl.SetPrompt(newPrompt)
|
||||
}
|
||||
|
||||
// Executes the main command processing loop
|
||||
func runCommandLoop(conn *tls.Conn, rl *readline.Instance, exitChan chan struct{}) {
|
||||
for {
|
||||
select {
|
||||
case <-exitChan:
|
||||
@@ -388,94 +466,60 @@ func main() {
|
||||
}
|
||||
|
||||
command := strings.TrimSpace(line)
|
||||
if command == "" {
|
||||
continue
|
||||
}
|
||||
|
||||
if command == "exit" {
|
||||
if currentAgentContext != "" {
|
||||
// fmt.Println("Clearing context")
|
||||
currentAgentContext = ""
|
||||
rlConfig.Prompt = "\033[31mSigma >\033[0m " // Reset to default prompt
|
||||
rl.SetPrompt(rlConfig.Prompt)
|
||||
// fmt.Println("Reversed to default prompt")
|
||||
|
||||
// Update tab completion after context change
|
||||
rl.Config.AutoComplete = getCompleter()
|
||||
|
||||
break
|
||||
} else {
|
||||
fmt.Println("Exiting...")
|
||||
_, err = conn.Write([]byte(command))
|
||||
if err != nil {
|
||||
fmt.Printf("Error sending command: %v", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
if strings.HasPrefix(command, "interact") {
|
||||
// fmt.Println("Changing context")
|
||||
cmd := strings.Fields(command)
|
||||
|
||||
if len(cmd) != 2 {
|
||||
fmt.Println("Usage: interact <agent_id>")
|
||||
break
|
||||
}
|
||||
|
||||
if !slices.Contains(agents, cmd[1]) {
|
||||
fmt.Println("Agent not found!")
|
||||
break
|
||||
} else {
|
||||
currentAgentContext = cmd[1]
|
||||
|
||||
// Update tab completion after context change
|
||||
rl.Config.AutoComplete = getCompleter()
|
||||
|
||||
// fmt.Printf("New context: %s\n", currentAgentContext)
|
||||
|
||||
// Update the prompt to reflect the current agent context
|
||||
if currentAgentContext == "" {
|
||||
rlConfig.Prompt = "\033[31mSigma >\033[0m " // Default prompt when no context
|
||||
} else {
|
||||
rlConfig.Prompt = fmt.Sprintf("\033[31mSigma [%s] >\033[0m ", currentAgentContext) // Prompt with agent context
|
||||
}
|
||||
rl.SetPrompt(rlConfig.Prompt) // Apply the updated prompt to the readline instance
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
// Simplified command handling pseudo-code
|
||||
if currentAgentContext != "" && isControlCommand(command) {
|
||||
// In agent context, prepend agent ID to any non-control command
|
||||
command = currentAgentContext + " " + command
|
||||
// fmt.Printf("Context command: %s", command)
|
||||
}
|
||||
|
||||
_, err = conn.Write([]byte(command))
|
||||
if err != nil {
|
||||
fmt.Printf("Error sending command: %v", err)
|
||||
shouldExit := processCommand(command, conn, rl)
|
||||
if shouldExit {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Check if command is a task for agent
|
||||
func isControlCommand(command string) bool {
|
||||
// Split the command into words
|
||||
parts := strings.Fields(command)
|
||||
if len(parts) == 0 {
|
||||
fmt.Println("Empty command :(")
|
||||
return false
|
||||
func main() {
|
||||
// Define flags for username, password and address of server
|
||||
operatorPassword := flag.String("p", "", "Password")
|
||||
operatorID := flag.String("u", "", "Username")
|
||||
serverAddr := flag.String("ip", defaultServerAddr, "IP and port of server")
|
||||
flag.Parse()
|
||||
|
||||
// Check if password and username are provided
|
||||
if *operatorPassword == "" || *operatorID == "" {
|
||||
fmt.Println("Usage: client.exe -u <username> -p <password> -ip <ip:port>")
|
||||
return
|
||||
}
|
||||
fmt.Println("Client's username and password provided")
|
||||
|
||||
// Setup TLS
|
||||
tlsConfig, err := setupTLSConfig()
|
||||
if err != nil {
|
||||
log.Fatalf("TLS setup failed: %v", err)
|
||||
}
|
||||
|
||||
// Check if the first word is a context command
|
||||
if slices.Contains(contextCommands, parts[0]) {
|
||||
// fmt.Println("Context command found!")
|
||||
return true
|
||||
} else {
|
||||
// fmt.Println("Not a context command :(")
|
||||
return false
|
||||
// Connect to server
|
||||
conn, err := connectToServer(*serverAddr, tlsConfig)
|
||||
if err != nil {
|
||||
log.Fatalf("Server connection failed: %v", err)
|
||||
}
|
||||
defer conn.Close()
|
||||
|
||||
// Authenticate
|
||||
err = authenticateUser(conn, *operatorID, *operatorPassword)
|
||||
if err != nil {
|
||||
log.Fatalf("Authentication failed: %v", err)
|
||||
}
|
||||
|
||||
// Setup readline
|
||||
rl, err := setupReadline()
|
||||
if err != nil {
|
||||
log.Fatalf("Readline setup failed: %v", err)
|
||||
}
|
||||
defer rl.Close()
|
||||
|
||||
// Create exit channel
|
||||
exitChan := make(chan struct{})
|
||||
|
||||
// Start server message listener
|
||||
startServerListener(conn, rl, exitChan)
|
||||
|
||||
// Run main command loop
|
||||
runCommandLoop(conn, rl, exitChan)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user