Unify clean shutdown and update C2 channel with session tracking

c2/channel/channel.go

@@ -3,11 +3,107 @@
 // components in order to support extracting components such as lhost and lport.
 package channel
 
+import (
+	"crypto/rand"
+	"encoding/base32"
+	"io"
+	"net"
+	"sync/atomic"
+	"time"
+
+	"github.com/vulncheck-oss/go-exploit/output"
+)
+
 type Channel struct {
 	IPAddr   string
 	HTTPAddr string
 	Port     int
 	HTTPPort int
 	Timeout  int
 	IsClient bool
+	Shutdown *atomic.Bool
+	Sessions map[string]Session
+	Input    io.Reader
+	Output   io.Writer // Currently unused but figured we'd add it ahead of time
+}
+
+type Session struct {
+	RemoteAddr     string
+	ConnectionTime time.Time
+	conn           *net.Conn
+}
+
+// HasSessions checks if a channel has any tracked sessions. This can be used to lookup if a C2
+// successfully received callbacks:
+//
+//	c, ok := c2.GetInstance(conf.C2Type)
+//	c.Channel().HasSessions()
+func (c *Channel) HasSessions() bool {
+	return len(c.Sessions) > 0
+}
+
+// AddSession adds a remote connection for session tracking. If a network connection is being
+// tracked it can be added here and will be cleaned up and closed automatically by the C2 on
+// shutdown.
+func (c *Channel) AddSession(conn *net.Conn, addr string) bool {
+	if len(c.Sessions) == 0 {
+		c.Sessions = make(map[string]Session)
+	}
+	// This is my session randomizing logic. The theory is that it keeps us dependency free while
+	// also creating the same 16bit strength of UUIDs. If we only plan on using the random UUIDs
+	// anyway this should meet the same goals while also being URL safe and no special characters.
+	k := make([]byte, 16)
+	_, err := rand.Read(k)
+	if err != nil {
+		output.PrintfFrameworkError("Could not add session: %s", err.Error())
+
+		return false
+	}
+	id := base32.StdEncoding.WithPadding(base32.NoPadding).EncodeToString(k)
+	c.Sessions[id] = Session{
+		// Add the time of now to the current connection time
+		ConnectionTime: time.Now(),
+		conn:           conn,
+		RemoteAddr:     addr,
+	}
+
+	return true
+}
+
+// RemoveSession removes a specific session ID and if a connection exists, closes it.
+func (c *Channel) RemoveSession(id string) bool {
+	if len(c.Sessions) == 0 {
+		output.PrintFrameworkDebug("No sessions exist")
+
+		return false
+	}
+	_, ok := c.Sessions[id]
+	if !ok {
+		output.PrintFrameworkError("Session ID does not exist")
+
+		return false
+	}
+	if c.Sessions[id].conn != nil {
+		(*c.Sessions[id].conn).Close()
+	}
+	delete(c.Sessions, id)
+
+	return true
+}
+
+// RemoveSessions removes all tracked sessions and closes any open connections if applicable.
+func (c *Channel) RemoveSessions() bool {
+	if len(c.Sessions) == 0 {
+		output.PrintFrameworkDebug("No sessions exist")
+
+		return false
+	}
+	for k := range c.Sessions {
+		if c.Sessions[k].conn != nil {
+			(*c.Sessions[k].conn).Close()
+		}
+		delete(c.Sessions, k)
+	}
+
+	return true
 }

c2/cli/basic.go

@@ -4,81 +4,115 @@ import (
 	"bufio"
 	"net"
 	"os"
-	"strings"
 	"sync"
+	"testing"
 	"time"
 
+	"github.com/vulncheck-oss/go-exploit/c2/channel"
 	"github.com/vulncheck-oss/go-exploit/output"
 	"github.com/vulncheck-oss/go-exploit/protocol"
 )
 
+// backgroundResponse handles the network connection reading for response data and contains a
+// trigger to the shutdown of the channel to ensure cleanup happens on socket close.
+func backgroundResponse(ch *channel.Channel, wg *sync.WaitGroup, conn net.Conn, responseCh chan string) {
+	defer wg.Done()
+	defer func(channel *channel.Channel) {
+		// Signals for both routines to stop, this should get triggered when socket is closed
+		// and causes it to fail the read
+		channel.Shutdown.Store(true)
+	}(ch)
+	responseBuffer := make([]byte, 1024)
+	for {
+		if ch.Shutdown.Load() {
+			return
+		}
+
+		err := conn.SetReadDeadline(time.Now().Add(1 * time.Second))
+		if err != nil {
+			output.PrintfFrameworkError("Error setting read deadline: %s, exiting.", err)
+
+			return
+		}
+
+		bytesRead, err := conn.Read(responseBuffer)
+		if err != nil && !os.IsTimeout(err) {
+			// things have gone sideways, but the command line won't know that
+			// until they attempt to execute a command and the socket fails.
+			// i think that's largely okay.
+			return
+		}
+
+		if bytesRead > 0 {
+			// I think there is technically a race condition here where the socket
+			// could have move data to write, but the user has already called exit
+			// below. I that that's tolerable for now.
+			responseCh <- string(responseBuffer[:bytesRead])
+		}
+	}
+}
+
 // A very basic reverse/bind shell handler.
-func Basic(conn net.Conn) {
+func Basic(conn net.Conn, ch *channel.Channel) {
 	// Create channels for communication between goroutines.
 	responseCh := make(chan string)
-	quit := make(chan struct{})
 
 	// Use a WaitGroup to wait for goroutines to finish.
 	var wg sync.WaitGroup
 
 	// Goroutine to read responses from the server.
 	wg.Add(1)
-	go func() {
-		defer wg.Done()
-		responseBuffer := make([]byte, 1024)
-		for {
-			select {
-			case <-quit:
-				return
-			default:
-				_ = conn.SetReadDeadline(time.Now().Add(1 * time.Second))
-				bytesRead, err := conn.Read(responseBuffer)
-				if err != nil && !strings.Contains(err.Error(), "i/o timeout") {
-					// things have gone sideways, but the command line won't know that
-					// until they attempt to execute a command and the socket fails.
-					// i think that's largely okay.
-					return
-				}
-				if bytesRead > 0 {
-					// I think there is technically a race condition here where the socket
-					// could have move data to write, but the user has already called exit
-					// below. I that that's tolerable for now.
-					responseCh <- string(responseBuffer[:bytesRead])
-				}
-			}
-		}
-	}()
+
+	// If running in the test context inherit the channel input setting, this will let us control the
+	// input of the shell programmatically.
+	if !testing.Testing() {
+		ch.Input = os.Stdin
+	}
+	go backgroundResponse(ch, &wg, conn, responseCh)
 
 	// Goroutine to handle responses and print them.
 	wg.Add(1)
-	go func() {
+	go func(channel *channel.Channel) {
 		defer wg.Done()
-		for response := range responseCh {
-			select {
-			case <-quit:
+		for {
+			if channel.Shutdown.Load() {
 				return
-			default:
+			}
+			select {
+			case response := <-responseCh:
 				output.PrintShell(response)
+			default:
 			}
 		}
-	}()
+	}(ch)
 
-	for {
-		// read user input until they type 'exit\n' or the socket breaks
-		// note that ReadString is blocking, so they won't know the socket
-		// is broken until they attempt to write something
-		reader := bufio.NewReader(os.Stdin)
-		command, _ := reader.ReadString('\n')
-		ok := protocol.TCPWrite(conn, []byte(command))
-		if !ok || command == "exit\n" {
-			break
-		}
-	}
+	go func(channel *channel.Channel) {
+		// no waitgroup for this one because blocking IO, but this should not matter
+		// since we are intentionally not trying to be a multi-implant C2 framework.
+		// There still remains the issue that you would need to hit enter to find out
+		// that the socket is dead but at least we can stop Basic() regardless of this fact.
+		// This issue of unblocking stdin is discussed at length here https://github.com/golang/go/issues/24842
+		for {
+			reader := bufio.NewReader(ch.Input)
+			command, _ := reader.ReadString('\n')
+			if channel.Shutdown.Load() {
+				break
+			}
+			if command == "exit\n" {
+				channel.Shutdown.Store(true)
 
-	// signal for everyone to shutdown
-	quit <- struct{}{}
-	close(responseCh)
+				break
+			}
+			ok := protocol.TCPWrite(conn, []byte(command))
+			if !ok {
+				channel.Shutdown.Store(true)
+
+				break
+			}
+		}
+	}(ch)
 
 	// wait until the go routines are clean up
 	wg.Wait()
+	close(responseCh)
 }

c2/external/external.go

@@ -163,12 +163,13 @@ import (
 
 // The Server struct holds the declared external modules internal functions and channel data.
 type Server struct {
-	flags   func()
-	init    func()
-	run     func(int) bool
-	meta    func(*channel.Channel)
-	Channel *channel.Channel
-	name    string
+	flags    func()
+	init     func()
+	run      func(int) bool
+	shutdown func() bool
+	meta     func(*channel.Channel)
+	channel  *channel.Channel
+	name     string
 }
 
 // The External interface defines which functions are required to be defined in an external C2
@@ -180,6 +181,7 @@ type External interface {
 	SetFlags(func())
 	SetInit(func())
 	SetRun(func(int) bool)
+	SetShutdown(func() bool)
 }
 
 var serverSingletons map[string]*Server
@@ -245,14 +247,14 @@ func (externalServer *Server) SetChannel(f func(*channel.Channel)) {
 }
 
 // Init triggers the set C2 initialization and passes the channel to the external module.
-func (externalServer *Server) Init(channel channel.Channel) bool {
+func (externalServer *Server) Init(channel *channel.Channel) bool {
 	if channel.IsClient {
 		output.PrintFrameworkError("Called ExternalServer as a client.")
 
 		return false
 	}
 	externalServer.init()
-	externalServer.meta(&channel)
+	externalServer.meta(channel)
 
 	return true
 }
@@ -268,3 +270,21 @@ func (externalServer *Server) SetRun(f func(int) bool) {
 func (externalServer *Server) Run(timeout int) {
 	externalServer.run(timeout)
 }
+
+// SetShutdown sets the function for server shutdown handling and session cleanup logic. This
+// function is what gets called when the framework receives a OS signal, a shell is closed, or
+// manually invoked.
+func (externalServer *Server) SetShutdown(f func() bool) {
+	externalServer.shutdown = f
+}
+
+// Shutdown triggers the set shutdown function.
+func (externalServer *Server) Shutdown() bool {
+	return externalServer.shutdown()
+}
+
+// Return the underlying C2 channel containing channel metadata and session tracking.
+func (externalServer *Server) Channel() *channel.Channel {
+	// I'd much rather have just exposed a `Server.Channel`, but we are interface bound
+	return externalServer.channel
+}

c2/factory.go

@@ -15,8 +15,10 @@ import (
 // An interface used by both reverse shells, bind shells, and stagers.
 type Interface interface {
 	CreateFlags()
-	Init(channel channel.Channel) bool
+	Init(channel *channel.Channel) bool
 	Run(timeout int)
+	Shutdown() bool
+	Channel() *channel.Channel
 }
 
 // Internal representation of a C2 implementation. Each C2 is managed by
@@ -65,7 +67,7 @@ var internalSupported = map[string]Impl{
 	"HTTPServeFile":     {Name: "HTTPServeFile", Category: HTTPServeFileCategory},
 	"HTTPServeShell":    {Name: "HTTPServeShell", Category: HTTPServeShellCategory},
 	"HTTPShellServer":   {Name: "HTTPShellServer", Category: HTTPShellServerCategory},
-	// Insure the internal supported External module name is an error if used
+	// Ensure the internal supported External module name is an error if used
 	// directly.
 	"External":    {Name: "", Category: InvalidCategory},
 	"ShellTunnel": {Name: "ShellTunnel", Category: ShellTunnelCategory},
@@ -150,6 +152,37 @@ func CreateFlags(implementation Impl) {
 	}
 }
 
+// HasSessions returns if the underlying channel has active sessions. This is useful for code that
+// needs to validate if callbacks have occurred and is a helper wrapper around the channel package
+// function of the same name.
+func HasSessions(implementation Impl) bool {
+	switch implementation.Category {
+	case SimpleShellServerCategory:
+		return simpleshell.GetServerInstance().Channel().HasSessions()
+	case SimpleShellClientCategory:
+		return simpleshell.GetClientInstance().Channel().HasSessions()
+	case SSLShellServerCategory:
+		return sslshell.GetInstance().Channel().HasSessions()
+	case HTTPServeFileCategory:
+		return httpservefile.GetInstance().Channel().HasSessions()
+	case HTTPServeShellCategory:
+		return httpserveshell.GetInstance().Channel().HasSessions()
+	case ExternalCategory:
+		if implementation.Name != "" {
+			return external.GetInstance(implementation.Name).Channel().HasSessions()
+		}
+	case HTTPShellServerCategory:
+		return httpshellserver.GetInstance().Channel().HasSessions()
+	case ShellTunnelCategory:
+		return shelltunnel.GetInstance().Channel().HasSessions()
+	case InvalidCategory:
+	default:
+	}
+	output.PrintFrameworkError("Invalid C2 Server")
+
+	return false
+}
+
 // Return the internal representation of a C2 from a string.
 func StringToImpl(c2Name string) (Impl, bool) {
 	for _, value := range internalSupported {