// httpservefile c2 spawns an HTTP or HTTPS server and hosts arbitrary user-provided files. The normal use case
// is for an exploit to curl/wget the file and execute it. This is useful to spawn connections to other
// tools (e.g. Metasploit, nc, etc.) or to go-exploit. This is not a traditional "c2" but serves as a useful
// backend that logically plugs into our c2 design.
//
// Files are provided on the command line as a comma delimited string. For example:
//
//	-httpServeFile.FilesToServe ./build/reverse_shell_windows-arm64.exe,./build/reverse_shell_linux-amd64
//
// The above will load two files: a windows reverse shell and a linux reverse shell. This c2 will then
// generate random names for the files and host them on an HTTP / HTTPS server. To interact with the
// files from an implementing exploit, you can fetch the filename to random name mapping using
// GetRandomName(). For example:
//
//	httpservefile.GetInstance().GetRandomName(linux64)
//
// Where linux64 is a variable that contains "reverse_shell_linux-amd64".
//
// If you are only hosting one file, then GetRandomName("") will also return your one file.
//
// Files can also be provided programmatically via the AddFile function (must be called before run).
package httpservefile

import (
	"bytes"
	"crypto/tls"
	"flag"
	"fmt"
	"net/http"
	"os"
	"path"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/vulncheck-oss/go-exploit/c2/channel"
	"github.com/vulncheck-oss/go-exploit/encryption"
	"github.com/vulncheck-oss/go-exploit/output"
	"github.com/vulncheck-oss/go-exploit/random"
)

type HostedFile struct {
	// The user provided filename
	RealName string
	// A randomly generated filename to serve
	RandomName string
	// The file's data
	FileData []byte
}

type Server struct {
	// The HTTP address to bind to
	HTTPAddr string
	// The HTTP port to bind to
	HTTPPort int
	// Set to the Server field in HTTP response
	ServerField string
	// Indicates if TLS should be enabled
	TLS bool
	// The file path to the user provided private key (if provided)
	PrivateKeyFile string
	// The file path to the user provided certificate (if provided)
	CertificateFile string
	// Loaded certificate
	Certificate tls.Certificate
	// A map of hosted files
	HostedFiles map[string]HostedFile // RealName -> struct
	// A comma delimited list of all the files to serve
	FilesToServe string
	// C2 channel and session metadata
	channel *channel.Channel
}

var singleton *Server

// A basic singleton interface for the c2.
func GetInstance() *Server {
	if singleton == nil {
		singleton = new(Server)
		// init hosted files map
		singleton.HostedFiles = make(map[string]HostedFile)
	}

	return singleton
}

// User options for serving a file over HTTP as the "c2".
func (httpServer *Server) CreateFlags() {
	// some c2 are really just chained implementations (e.g. httpserveshell is httpservefile plus simpleshell or sslshell).
	// so first check if these values have already been generated
	if flag.Lookup("httpServeFile.FilesToServe") == nil {
		flag.StringVar(&httpServer.FilesToServe, "httpServeFile.FilesToServe", "", "A comma delimited list of all the files to serve")
		flag.StringVar(&httpServer.ServerField, "httpServeFile.ServerField", "Apache", "The value to insert in the HTTP server field")
		flag.BoolVar(&httpServer.TLS, "httpServeFile.TLS", false, "Indicates if the HTTP server should use encryption")
		flag.StringVar(&httpServer.PrivateKeyFile, "httpServeFile.PrivateKeyFile", "", "A private key to use with the HTTPS server")
		flag.StringVar(&httpServer.CertificateFile, "httpServeFile.CertificateFile", "", "The certificate to use with the HTTPS server")
	}
}

// Return the C2 specific channel.
func (httpServer *Server) Channel() *channel.Channel {
	return httpServer.channel
}

// Shutdown the C2 server and cleanup all the sessions.
func (httpServer *Server) Shutdown() bool {
	// Account for non-running case
	if httpServer.Channel() == nil {
		return true
	}
	output.PrintFrameworkStatus("Shutting down the HTTP Server")
	if len(httpServer.Channel().Sessions) > 0 {
		for k := range httpServer.Channel().Sessions {
			httpServer.Channel().RemoveSession(k)
		}
	}

	return true
}

// load the provided files into memory, stored in a map, and loads the tls cert if needed.
func (httpServer *Server) Init(channel *channel.Channel) bool {
	if channel.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
	}
	httpServer.channel = channel
	if channel.IsClient {
		output.PrintFrameworkError("Called C2HTTPServer as a client. Use lhost and lport.")

		return false
	}

	switch {
	case channel.HTTPPort == 0 && channel.Port != 0:
		// must be stand-alone invocation of HTTPServeFile
		httpServer.HTTPAddr = channel.IPAddr
		httpServer.HTTPPort = channel.Port
	case channel.HTTPPort != 0:
		// must be used with another C2
		httpServer.HTTPAddr = channel.HTTPAddr
		httpServer.HTTPPort = channel.HTTPPort
	default:
		output.PrintFrameworkError("Called HTTPServeFile without specifying a bind port.")

		return false
	}

	// split the provided files, read them in, and store them in the map
	files := strings.Split(httpServer.FilesToServe, ",")
	for _, file := range files {
		if len(file) == 0 {
			continue
		}
		output.PrintfFrameworkStatus("Loading the provided file: %s", file)
		fileData, err := os.ReadFile(file)
		if err != nil {
			output.PrintFrameworkError(err.Error())

			return false
		}

		// remove the path from the name (check for / and \)
		shortName := file
		pathSepIndex := strings.LastIndex(shortName, "/")
		if pathSepIndex != -1 {
			shortName = shortName[pathSepIndex+1:]
		}
		pathSepIndex = strings.LastIndex(shortName, `\`)
		if pathSepIndex != -1 {
			shortName = shortName[pathSepIndex+1:]
		}

		hosted := HostedFile{
			RealName:   shortName,
			RandomName: random.RandLetters(12),
			FileData:   fileData,
		}

		output.PrintfFrameworkDebug("Added %s as %s", hosted.RealName, hosted.RandomName)
		httpServer.HostedFiles[shortName] = hosted
	}

	if httpServer.TLS {
		var ok bool
		var err error
		if len(httpServer.CertificateFile) != 0 && len(httpServer.PrivateKeyFile) != 0 {
			httpServer.Certificate, err = tls.LoadX509KeyPair(httpServer.CertificateFile, httpServer.PrivateKeyFile)
			if err != nil {
				output.PrintfFrameworkError("Error loading certificate: %s", err.Error())

				return false
			}
		} else {
			output.PrintFrameworkStatus("Certificate not provided. Generating a TLS Certificate")
			httpServer.Certificate, ok = encryption.GenerateCertificate()
			if !ok {
				return false
			}
		}
	}

	return true
}

// Adds a file to the server. A route will be created for "randomName" when run() is executed.
func (httpServer *Server) AddFile(realName string, randomName string, data []byte) {
	hostMe := HostedFile{RealName: realName, RandomName: randomName, FileData: data}
	httpServer.HostedFiles[realName] = hostMe
}

// start the HTTP server and listen for incoming requests for `httpServer.FileName`.
func (httpServer *Server) Run(timeout int) {
	if len(httpServer.HostedFiles) == 0 {
		output.PrintFrameworkError("No files provided via httpServeFile.FilesToServe or programmatically")

		return
	}

	// set up handlers for each of the files
	for _, hosted := range httpServer.HostedFiles {
		http.HandleFunc("/"+hosted.RandomName, func(writer http.ResponseWriter, req *http.Request) {
			output.PrintfFrameworkStatus("Connection from %s requested %s", req.RemoteAddr, req.URL.Path)
			httpServer.Channel().AddSession(nil, req.RemoteAddr)

			writer.Header().Set("Server", httpServer.ServerField)

			name := path.Base(req.URL.Path)
			// cannot used hosted as the values move on with the loop
			for _, selected := range httpServer.HostedFiles {
				if selected.RandomName == name {
					http.ServeContent(writer, req, selected.RandomName, time.Time{}, bytes.NewReader(selected.FileData))

					return
				}
			}

			writer.WriteHeader(http.StatusNotFound)
		})
	}

	var wg sync.WaitGroup
	connectionString := fmt.Sprintf("%s:%d", httpServer.HTTPAddr, httpServer.HTTPPort)
	wg.Add(1)
	go func() {
		if httpServer.TLS {
			output.PrintfFrameworkStatus("Starting an HTTPS server on %s", connectionString)
			tlsConfig := &tls.Config{
				Certificates: []tls.Certificate{httpServer.Certificate},
				// We have no control over the SSL versions supported on the remote target. Be permissive for more targets.
				MinVersion: tls.VersionSSL30,
			}
			server := http.Server{
				Addr:      connectionString,
				TLSConfig: tlsConfig,
			}
			defer server.Close()
			// Track if the server has signaled for shutdown and if so mark the waitgroup and trigger shutdown
			go func() {
				for {
					if httpServer.Channel().Shutdown.Load() {
						server.Close()
						httpServer.Shutdown()
						wg.Done()

						break
					}
					time.Sleep(10 * time.Millisecond)
				}
			}()
			// Handle timeouts
			go func() {
				time.Sleep(time.Duration(timeout) * time.Second)
				output.PrintFrameworkError("Timeout met. Shutting down shell listener.")
				// We do not care about sessions with file
				httpServer.channel.Shutdown.Store(true)
			}()

			_ = server.ListenAndServeTLS("", "")
		} else {
			output.PrintfFrameworkStatus("Starting an HTTP server on %s", connectionString)
			server := http.Server{
				Addr: connectionString,
			}
			defer server.Close()
			// Track if the server has signaled for shutdown and if so mark the waitgroup and trigger shutdown
			go func() {
				for {
					if httpServer.Channel().Shutdown.Load() {
						server.Close()
						httpServer.Shutdown()
						wg.Done()

						break
					}
					time.Sleep(10 * time.Millisecond)
				}
			}()
			// Handle timeouts
			go func() {
				time.Sleep(time.Duration(timeout) * time.Second)
				output.PrintFrameworkError("Timeout met. Shutting down shell listener.")
				// We do not care about sessions with file
				httpServer.channel.Shutdown.Store(true)
			}()
			_ = http.ListenAndServe(connectionString, nil)
		}
	}()

	wg.Wait()
	httpServer.Channel().Shutdown.Store(true)
}

// Returns the random name of the provided filename. If filename is empty, return the first entry.
func (httpServer *Server) GetRandomName(filename string) string {
	if len(filename) == 0 {
		for _, hosted := range httpServer.HostedFiles {
			return hosted.RandomName
		}
	}

	hosted, found := httpServer.HostedFiles[filename]
	if !found {
		output.PrintfFrameworkError("Requested a file that doesn't exist: %s", filename)

		return ""
	}

	return hosted.RandomName
}