Merge pull request #367 from vulncheck-oss/fix/unify-broken-httpserveshell

Fix OS signal handling on HTTPServeShell with Active Sessions

Author: j-baines

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
@@ -77,46 +79,112 @@ func (serveShell *Server) CreateFlags() {
 }
 
 // load the provided file into memory. Generate the random filename.
-func (serveShell *Server) Init(channel *channel.Channel) bool {
-	if channel.Shutdown == nil {
+func (serveShell *Server) Init(ch *channel.Channel) bool {
+	if ch.Shutdown == nil {
 		// Initialize the shutdown atomic. This lets us not have to define it if the C2 is manually
 		// configured.
 		var shutdown atomic.Bool
 		shutdown.Store(false)
-		channel.Shutdown = &shutdown
+		ch.Shutdown = &shutdown
 	}
-	serveShell.channel = channel
+	serveShell.pastTimeout.Store(false)
+	serveShell.channel = ch
 	if len(serveShell.HTTPAddr) == 0 {
 		output.PrintFrameworkError("User must specify -httpServeFile.BindAddr")
 
 		return false
 	}
-	channel.HTTPAddr = serveShell.HTTPAddr
-	channel.HTTPPort = serveShell.HTTPPort
+	ch.HTTPAddr = serveShell.HTTPAddr
+	ch.HTTPPort = serveShell.HTTPPort
 
-	if !httpservefile.GetInstance().Init(channel) {
+	if !httpservefile.GetInstance().Init(ch) {
 		return false
 	}
 
+	// Initialize the shell server channels with variables from upstream
+	var shutdown atomic.Bool
+	shutdown.Store(false)
+	shellChannel := &channel.Channel{
+		IPAddr:   ch.IPAddr,
+		Port:     ch.Port,
+		IsClient: false,
+		Shutdown: &shutdown,
+	}
 	if serveShell.SSLShell {
-		return sslshell.GetInstance().Init(channel)
+		return sslshell.GetInstance().Init(shellChannel)
 	}
 
-	return simpleshell.GetServerInstance().Init(channel)
+	return simpleshell.GetServerInstance().Init(shellChannel)
 }
 
 // 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:
+	//
+	//  - 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.
 	//
-	// Each of the shell sessions and httpservefile sessions should be handled by the channel Shutdown
-	// atomic.
+	// 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 {
-		sslshell.GetInstance().Channel().Shutdown.Store(true)
+		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 {
-		simpleshell.GetServerInstance().Channel().Shutdown.Store(true)
+		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
 }
@@ -141,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() {
@@ -150,7 +222,7 @@ func (serveShell *Server) Run(timeout int) {
 			go func() {
 				for {
 					if sslshell.GetInstance().Channel().Shutdown.Load() {
-						sslshell.GetInstance().Shutdown()
+						serveShell.Channel().Shutdown.Store(true)
 						wg.Done()
 
 						break
@@ -163,7 +235,7 @@ func (serveShell *Server) Run(timeout int) {
 			go func() {
 				for {
 					if simpleshell.GetServerInstance().Channel().Shutdown.Load() {
-						simpleshell.GetServerInstance().Shutdown()
+						serveShell.Channel().Shutdown.Store(true)
 						wg.Done()
 
 						break
@@ -175,9 +247,7 @@ func (serveShell *Server) Run(timeout int) {
 
 	// Spin up the http server
 	wg.Add(1)
-	go func() {
-		httpservefile.GetInstance().Run(timeout)
-	}()
+	go func() { httpservefile.GetInstance().Run(timeout) }()
 
 	// wait until the go routines are clean up
 	wg.Wait()

c2/simpleshell/simpleshellserver.go

@@ -48,7 +48,7 @@ func (shellServer *Server) Shutdown() bool {
 	output.PrintFrameworkStatus("C2 received shutdown, killing server and client sockets for shell server")
 	if len(shellServer.Channel().Sessions) > 0 {
 		for k, session := range shellServer.Channel().Sessions {
-			output.PrintfFrameworkStatus("Connection closed: %s", session.RemoteAddr)
+			output.PrintfFrameworkStatus("Connection closed for shell server: %s", session.RemoteAddr)
 			shellServer.Channel().RemoveSession(k)
 		}
 	}