Create ServeShell logic for the c2 lacking sessions.

Author: terrorbyte

c2/httpserveshell/httpserveshell.go

@@ -34,6 +34,7 @@ import (
 	"flag"
 	"sync"
 	"sync/atomic"
+	"time"
 
 	"github.com/vulncheck-oss/go-exploit/c2/channel"
 	"github.com/vulncheck-oss/go-exploit/c2/httpservefile"
@@ -50,7 +51,8 @@ type Server struct {
 	// The HTTP port to bind to
 	HTTPPort int
 	// The underlying C2 channel with metadata and session information
-	channel *channel.Channel
+	channel     *channel.Channel
+	pastTimeout atomic.Bool
 }
 
 var singleton *Server
@@ -85,6 +87,7 @@ func (serveShell *Server) Init(ch *channel.Channel) bool {
 		shutdown.Store(false)
 		ch.Shutdown = &shutdown
 	}
+	serveShell.pastTimeout.Store(false)
 	serveShell.channel = ch
 	if len(serveShell.HTTPAddr) == 0 {
 		output.PrintFrameworkError("User must specify -httpServeFile.BindAddr")
@@ -107,7 +110,6 @@ func (serveShell *Server) Init(ch *channel.Channel) bool {
 		IsClient: false,
 		Shutdown: &shutdown,
 	}
-	shellChannel.Shutdown = &shutdown
 	if serveShell.SSLShell {
 		return sslshell.GetInstance().Init(shellChannel)
 	}
@@ -117,20 +119,72 @@ func (serveShell *Server) Init(ch *channel.Channel) bool {
 
 // Shutdown triggers the shutdown for all running C2s.
 func (serveShell *Server) Shutdown() bool {
-	// Since no logic actually handles client interactions here, this is stub.
+	// This is a bit confusing at first glance, but it solves the fact that this c2 doesn't directly
+	// keep track of sessions and we can't differentiate between a timeout "done" and a signal "done".
+	// What this means is that if a underlying shell server has sessions that we want to keep open for
+	// use after callbacks occur we have to account for a few things:
 	//
-	// Each of the shell sessions and httpservefile sessions should be handled by the channel Shutdown
-	// atomic.
+	//  - If serveShell shutdown is called and there are no sessions, just trigger shutdown on the
+	//    underlying shell server (easy).
+	//  - If the server does have sessions, we have a few issues. The serveShell shutdown is triggered
+	//    from a shutdown call, meaning that it's already in a closing state. There is no way to tell
+	//    if it was an OS signal or a timeout anymore and now if a background shell is running and we
+	//    are closing serveShell we cannot catch the signal and pass the shutdown to the shell server.
+	//
+	// In order to solve the second, we added a `pastTimeout` atomic that only signals if we are past
+	// timeout. Now, when timeout is reached and there's a background shell (the positive case) we
+	// reset the Shutdown atomic to false and then begin looping to check if it closes again, making
+	// the server in a state that it knows it's past timeout and reactivating the server until a
+	// signal is hit or the underlying server also shuts down.
+
+	httpservefile.GetInstance().Channel().Shutdown.Store(true)
 	if serveShell.SSLShell {
 		if !sslshell.GetInstance().Channel().HasSessions() {
 			sslshell.GetInstance().Channel().Shutdown.Store(true)
+		} else {
+			// Session exist, reset the shutdown atomic and loop until a second shutdown occurs.
+			serveShell.Channel().Shutdown.Store(false)
+			for {
+				if serveShell.Channel().Shutdown.Load() {
+					// The the shutdown happens and it is past timeout, that means that
+					// we have sessions but timeout has passed, so reset all atomics.
+					// Now when the loop happens we will be able to check for other
+					// shutdown signals of any kind.
+					if serveShell.pastTimeout.Load() {
+						serveShell.Channel().Shutdown.Store(false)
+						serveShell.pastTimeout.Store(false)
+					} else {
+						sslshell.GetInstance().Channel().Shutdown.Store(true)
+
+						break
+					}
+				}
+			}
 		}
 	} else {
 		if !simpleshell.GetServerInstance().Channel().HasSessions() {
 			simpleshell.GetServerInstance().Channel().Shutdown.Store(true)
+		} else {
+			// Session exist, reset the shutdown atomic and loop until a second shutdown occurs.
+			serveShell.Channel().Shutdown.Store(false)
+			for {
+				if serveShell.Channel().Shutdown.Load() {
+					// The the shutdown happens and it is past timeout, that means that
+					// we have sessions but timeout has passed, so reset all atomics.
+					// Now when the loop happens we will be able to check for other
+					// shutdown signals of any kind.
+					if serveShell.pastTimeout.Load() {
+						serveShell.Channel().Shutdown.Store(false)
+						serveShell.pastTimeout.Store(false)
+					} else {
+						simpleshell.GetServerInstance().Channel().Shutdown.Store(true)
+
+						break
+					}
+				}
+			}
 		}
 	}
-	httpservefile.GetInstance().Channel().Shutdown.Store(true)
 
 	return true
 }
@@ -155,6 +209,10 @@ func (serveShell *Server) Run(timeout int) {
 			}
 		}
 	}()
+	go func() {
+		time.Sleep(time.Duration(timeout) * time.Second)
+		serveShell.pastTimeout.Store(true)
+	}()
 	// Spin up the shell
 	wg.Add(1)
 	go func() {
@@ -189,16 +247,7 @@ func (serveShell *Server) Run(timeout int) {
 
 	// Spin up the http server
 	wg.Add(1)
-	go func() {
-		httpservefile.GetInstance().Run(timeout)
-		for {
-			if serveShell.Channel().Shutdown.Load() {
-				serveShell.Shutdown()
-
-				break
-			}
-		}
-	}()
+	go func() { httpservefile.GetInstance().Run(timeout) }()
 
 	// wait until the go routines are clean up
 	wg.Wait()