// Package exploit is the entrypoint for exploits developed with the go-exploit framework.
//
// The exploit package invokes command line parsing, handles c2, and calls into the three stages of
// exploitation (as defined by go-exploit). In order to use this framework, implementing exploits
// should follow this general template:
//
//	package main
//
//	import (
//		"github.com/vulncheck-oss/go-exploit"
//		"github.com/vulncheck-oss/go-exploit/c2"
//		"github.com/vulncheck-oss/go-exploit/config"
//		"github.com/vulncheck-oss/go-exploit/output"
//	)
//
//	type MyExploit struct{}
//
//	func generatePayload(conf *config.Config) (string, bool) {
//	        generated := ""
//
//	        switch conf.C2Type {
//	        case c2.SimpleShellServer:
//	                // generated = reverse.Bash.TCPRedirection(conf.Lhost, conf.Lport) // Example
//	        default:
//	                output.PrintError("Invalid payload")
//
//	                return generated, false
//	        }
//
//	        return generated, true
//	}
//
//	func (sploit MyExploit) ValidateTarget(conf *config.Config) bool {
//		return false
//	}
//
//	func (sploit MyExploit) CheckVersion(conf *config.Config) exploit.VersionCheckType {
//		return exploit.NotImplemented
//	}
//
//	func (sploit MyExploit) RunExploit(conf *config.Config) bool {
//		generated, ok := generatePayload(conf)
//	 	if !ok {
//	     		return false
//	 	}
//		return true
//	}
//
//	func main() {
//		supportedC2 := []c2.Impl{
//			c2.SimpleShellServer,
//			c2.SimpleShellClient,
//		}
//		conf := config.NewRemoteExploit(
//			config.ImplementedFeatures{AssetDetection: false, VersionScanning: false, Exploitation: false},
//			config.CodeExecution, supportedC2, "Vendor", []string{"Product"},
//			[]string{"cpe:2.3:a:vendor:product"}, "CVE-2024-1270", "HTTP", 8080)
//
//		sploit := MyExploit{}
//		exploit.RunProgram(sploit, conf)
//	}
package exploit

import (
	"crypto/tls"
	"fmt"
	"os"
	"os/signal"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/Masterminds/semver"
	"github.com/vulncheck-oss/go-exploit/c2"
	"github.com/vulncheck-oss/go-exploit/c2/channel"
	"github.com/vulncheck-oss/go-exploit/cli"
	"github.com/vulncheck-oss/go-exploit/config"
	"github.com/vulncheck-oss/go-exploit/db"
	"github.com/vulncheck-oss/go-exploit/output"
	"github.com/vulncheck-oss/go-exploit/protocol"
)

// The return type for CheckVersion().
type VersionCheckType int

const (
	// The target is not vulnerable.
	NotVulnerable VersionCheckType = 0
	// The target is vulnerable.
	Vulnerable VersionCheckType = 1
	// Based on incomplete information, the target might be vulnerable.
	PossiblyVulnerable VersionCheckType = 2
	// Something went wrong during CheckVersion().
	Unknown VersionCheckType = 3
	// CheckVersion() is not implemented.
	NotImplemented VersionCheckType = 4
)

// Exploit is the implementing interface for go-exploit exploits. The functions
// are called in order: ValidateTarget, CheckVersion, RunExploit.
//
// # ValidateTarget
//
// ValidateTarget is for determining if the target is the type of software the
// implemented exploit would like to exploit. This is to avoid throwing an exploit
// at target would never be vulnerable. For example, if an exploit is targeting
// Confluence, then ValidateTarget might look like the following
//
//	func (sploit ConfluenceExploit) ValidateTarget(conf *config.Config) bool {
//		url := protocol.GenerateURL(conf.Rhost, conf.Rport, conf.SSL, "/")
//		resp, _, ok := protocol.HTTPSendAndRecv("GET", url, "")
//		if !ok {
//			return false
//		}
//
//		if resp.StatusCode != 200 {
//			output.PrintfError("Received an unexpected HTTP status code: %d", resp.StatusCode)
//
//			return false
//		}
//		_, ok = resp.Header["X-Confluence-Request-Time"]
//
//		return ok
//	}
//
// Above you can see ValidateTarget returns true *only* if it finds the X-Confluence-Request-Time
// HTTP header. The exploit will not continue on if false is returned. If true is returned then
// it will move on to the next stage (CheckVersion).
//
// # CheckVersion
//
// CheckVersion is for determning if the target is an affected version or not. Again, to avoid
// throwing an exploit at a target that is not vulnerable. CheckVersion is intended to be a
// non-intrusive version check. That generally means doing things like:
//
//   - Extracting the version number from a login page
//   - Examining the HTTP Last-Modified header
//   - Looking for new functionality introduce in the patch
//
// For example, to check for CVE-2022-30525, you could do something like this.
//
//	func (sploit ZyxelExploit) CheckVersion(conf *config.Config) exploit.VersionCheckType {
//		url := protocol.GenerateURL(conf.Rhost, conf.Rport, conf.SSL, "/")
//		resp, bodyString, ok := protocol.HTTPSendAndRecv("GET", url, "")
//		if !ok {
//			return exploit.Unknown
//		}
//
//		if resp.StatusCode != 200 {
//			output.PrintfError("Received an unexpected HTTP status code: %d", resp.StatusCode)
//
//			return exploit.Unknown
//		}
//
//		if !strings.Contains(bodyString, "zyFunction.js") {
//			output.PrintError("The HTTP response did not contain an expected JavaScript include")
//
//			return exploit.Unknown
//		}
//
//		re := regexp.MustCompile(`src="/ext-js/app/common/zyFunction.js\?v=([0-9]+)"></script>`)
//		res := re.FindAllStringSubmatch(bodyString, -1)
//		if len(res) == 0 {
//			output.PrintError("Could not extract the build date from the target")
//
//			return exploit.Unknown
//		}
//
//		output.PrintfStatus("The device has a self-reported firmware publication date of %s", res[0][1])
//		date64, _ := strconv.ParseInt(res[0][1], 10, 64)
//		if date64 < 220415000000 {
//			return exploit.Vulnerable
//		}
//
//		return exploit.NotVulnerable
//	}
//
// Regardless, the goal is to avoid throwing the exploit until you are somewhat sure that it should
// land. This cannot always be accomplished so the return of exploit.NotImplemented is always on offer,
// and the attacker can skip this step via configuration if they please.
//
// # RunExploit
//
// RunExploit should contain the logic for exploiting the target. There is almost no requirement on this
// function other than the attacker do their thing. The on thing the implementation should do is return
// false if believe their attack has failed.
type Exploit interface {
	ValidateTarget(conf *config.Config) bool
	CheckVersion(conf *config.Config) VersionCheckType
	RunExploit(conf *config.Config) bool
}

// For syncing c2 and exploit threads.
var globalWG sync.WaitGroup

// doVerify is a wrapper around the implemented exploit's ValidateTarget() function. The results will
// be logged in a parsable fashion (e.g. verified=false) and stored in the sqlite db if provided.
func doVerify(sploit Exploit, conf *config.Config) bool {
	output.PrintFrameworkStatus(fmt.Sprintf("Validating %s target", conf.Product), "host", conf.Rhost, "port", conf.Rport)

	// first check to see if we already verified this target as specific software
	result, ok := db.GetVerified(conf.Product, conf.Rhost, conf.Rport)
	if !ok {
		// if we couldn't query the DB (or there was an error) call down into the exploit
		result = sploit.ValidateTarget(conf)
		// update the database with the result so we can skip it next time
		db.UpdateVerified(conf.Product, result, "", conf.Rhost, conf.Rport)
	} else {
		output.PrintFrameworkTrace("Verified software cache hit", "result", result)
	}

	if result {
		output.PrintFrameworkSuccess("Target verification succeeded!", "host", conf.Rhost, "port", conf.Rport, "verified", true)
	} else {
		output.PrintFrameworkStatus("The target isn't recognized as "+conf.Product, "host", conf.Rhost, "port", conf.Rport, "verified", false)
	}

	return result
}

// doVersionCheck is a wrapper around the implemented exploit's CheckVersion() function.
func doVersionCheck(sploit Exploit, conf *config.Config) bool {
	output.PrintFrameworkStatus("Running a version check on the remote target", "host", conf.Rhost, "port", conf.Rport)
	result := sploit.CheckVersion(conf)
	switch result {
	case NotVulnerable:
		output.PrintFrameworkStatus("The target appears to be a patched version.", "host", conf.Rhost, "port", conf.Rport, "vulnerable", "no")

		return false
	case Vulnerable:
		output.PrintFrameworkSuccess("The target appears to be a vulnerable version!", "host", conf.Rhost, "port", conf.Rport, "vulnerable", "yes")
	case PossiblyVulnerable:
		output.PrintFrameworkSuccess("The target *might* be a vulnerable version. Continuing.", "host", conf.Rhost, "port", conf.Rport, "vulnerable", "possibly")
	case Unknown:
		output.PrintFrameworkStatus("The result of the version check returned an unknown state.", "host", conf.Rhost, "port", conf.Rport, "vulnerable", "unknown")

		return false
	case NotImplemented:
		output.PrintFrameworkStatus("This exploit has not implemented a version check", "host", conf.Rhost, "port", conf.Rport, "vulnerable", "notimplemented")
	}

	return true
}

// Automatically determine if the remote target is using SSL or not. This *does* work
// even if a proxy is configured. This can be slow when dealing with non-SSL
// targets that don't respond to the handshake attempt, but it seems a reasonable trade-off.
// return bool (connected), bool (ssl).
func determineServerSSL(rhost string, rport int) (bool, bool) {
	conn, ok := protocol.TCPConnect(rhost, rport)
	if !ok {
		return false, false
	}
	defer conn.Close()

	tlsConn := tls.Client(conn, &tls.Config{InsecureSkipVerify: true})
	_ = tlsConn.SetReadDeadline(time.Now().Add(10 * time.Second))
	err := tlsConn.Handshake()

	return true, err == nil
}

// Invokes command line parsing based on the type of exploit that was implemented.
func parseCommandLine(conf *config.Config) bool {
	switch conf.ExType {
	case config.CodeExecution:
		return cli.CodeExecutionCmdLineParse(conf)
	case config.InformationDisclosure:
		return cli.InformationDisclosureCmdLineParse(conf)
	case config.Webshell:
		return cli.WebShellCmdLineParse(conf)
	case config.FileFormat:
		return cli.FormatFileCmdLineParse(conf)
	case config.Local:
		return cli.LocalCmdLineParse(conf)
	default:
		output.PrintFrameworkError("Invalid exploit type provided.")

		return false
	}
}

// Manually start the C2 server. This is used when Config.C2AutoStart is
// disabled and for when you may not want to start the server until
// another action is complete.
func StartC2(conf *config.Config) bool {
	return startC2Server(conf)
}

func startC2Server(conf *config.Config) bool {
	if conf.DoExploit && !conf.ThirdPartyC2Server && conf.Bport == 0 &&
		(conf.ExType != config.InformationDisclosure && conf.ExType != config.Webshell) {
		c2Impl, success := c2.GetInstance(conf.C2Type)
		if !success || c2Impl == nil {
			return false
		}

		sigint := make(chan os.Signal, 1)
		signal.Notify(sigint, os.Interrupt)

		var shutdown atomic.Bool
		shutdown.Store(false)
		c2channel := &channel.Channel{
			IPAddr:   conf.Lhost,
			Port:     conf.Lport,
			IsClient: false,
			Shutdown: &shutdown,
		}
		// Handle the signal interrupt channel. If the signal is triggered, then trigger the done
		// channel which will clean up the server and close cleanly.
		go func(sigint <-chan os.Signal, channel *channel.Channel) {
			<-sigint
			output.PrintfFrameworkStatus("Interrupt signal received")
			channel.Shutdown.Store(true)
		}(sigint, c2channel)

		success = c2Impl.Init(c2channel)
		if !success {
			return false
		}

		globalWG.Add(1)
		go func() {
			defer globalWG.Done()
			c2Impl.Run(conf.C2Timeout)
			output.PrintFrameworkStatus("C2 server exited")
		}()
	}

	return true
}

// execute verify, version check, and exploit. Return false if an unrecoverable error occurred.
//
//nolint:gocognit
func doScan(sploit Exploit, conf *config.Config) bool {
	// autodetect if the target is using SSL or not
	if conf.DetermineSSL {
		connected, sslEnabled := determineServerSSL(conf.Rhost, conf.Rport)
		if !connected {
			return true
		}
		conf.SSL = sslEnabled
	}

	if conf.DoVerify {
		if !doVerify(sploit, conf) {
			// C2 cleanup is meaningless with third party C2s
			if !conf.ThirdPartyC2Server {
				// Shuts down the C2 if verification fails
				c, ok := c2.GetInstance(conf.C2Type)
				if !ok {
					output.PrintFrameworkError("Could not get C2 configuration")

					return false
				}
				c.Shutdown()
			}

			return true
		}
	}

	if conf.DoVersionCheck {
		if !doVersionCheck(sploit, conf) {
			// C2 cleanup is meaningless with third party C2s
			if !conf.ThirdPartyC2Server {
				// Shuts down the C2 if version check fails
				c, ok := c2.GetInstance(conf.C2Type)
				if !ok {
					output.PrintFrameworkError("Could not get C2 configuration")

					return false
				}
				c.Shutdown()
			}

			return true
		}
	}

	if conf.DoExploit {
		// execute exploit attempts on a new thread
		globalWG.Add(1)
		go func() {
			defer globalWG.Done()
			ok := sploit.RunExploit(conf)
			if ok {
				output.PrintFrameworkSuccess("Exploit successfully completed", "exploited", true)
			} else {
				output.PrintFrameworkStatus("Exploit exited with an error", "exploited", false)
			}
		}()

		// if the "c2" connects to a bindshell, call init to update the rhost/bport
		// and then attempt to connect
		if !conf.ThirdPartyC2Server && conf.Bport != 0 {
			c2Impl, success := c2.GetInstance(conf.C2Type)
			if !success || c2Impl == nil {
				return false
			}

			success = c2Impl.Init(&channel.Channel{
				IPAddr:   conf.Rhost,
				Port:     conf.Bport,
				IsClient: true,
			})
			if !success {
				return false
			}

			globalWG.Add(1)
			go func() {
				defer globalWG.Done()
				c2Impl.Run(conf.C2Timeout)
				output.PrintFrameworkStatus("C2 client exited")
			}()
		}
	}

	return true
}

// Prints the version to the log file using status VERSION and a parsable version string (version=).
// Additionally, updates the database if it's in use. Typically should be called from the exploit.
func StoreVersion(conf *config.Config, version string) {
	output.PrintVersion("The reported version is "+version, conf.Rhost, conf.Rport, version)
	db.UpdateVerified(conf.Product, true, version, conf.Rhost, conf.Rport)
}

// Compare a version to a semantic version constraint using the [Masterminds semver constraints](https://github.com/Masterminds/semver?tab=readme-ov-file#checking-version-constraints).
// Provide a version string and a constraint and if the semver is within the constraint a boolean
// response of whether the version is constrained or not will occur. Any errors from the constraint
// or version will propagate through the framework errors and the value will be false.
//
// Deprecated: The location of the version checking in this package made little sense, with the
// addition of the search package this function should be used from that package.
func CheckSemVer(version string, constraint string) bool {
	c, err := semver.NewConstraint(constraint)
	if err != nil {
		output.PrintfFrameworkError("Invalid constraint: %s", err.Error())

		return false
	}
	v, err := semver.NewVersion(version)
	if err != nil {
		output.PrintfFrameworkError("Invalid version: %s", err.Error())

		return false
	}

	return c.Check(v)
}

// modify godebug to re-enable old cipher suites that were removed in 1.22. This does have implications for our
// client fingerprint, and we should consider how to improve/fix that in the future. We also should be respectful
// of other disabling this feature, so we will check for it before re-enabling it.
func updateGoDebug() {
	currentGODEBUG := os.Getenv("GODEBUG")
	if strings.Contains(currentGODEBUG, "tlsrsakex") {
		// do nothing
		return
	}
	if len(currentGODEBUG) == 0 {
		os.Setenv("GODEBUG", "tlsrsakex=1")
	} else {
		// append our new setting to the end
		currentGODEBUG += ",tlsrsakex=1"
		os.Setenv("GODEBUG", currentGODEBUG)
	}
}

// Effectively the package main function. Parses configuration, starts command and control,
// controls which targets are scanned, initiates call down into the exploits implementation
// and is ultimately responsible for waiting for all c2 and attack threads to finish.
//
// This function also runs `flag.Parse()` so any defined flags will be parsed when RunProgram
// is called.
//
// This function should be called by the implementing exploit, likely in the main function.
func RunProgram(sploit Exploit, conf *config.Config) {
	updateGoDebug()
	if !parseCommandLine(conf) {
		return
	}

	// create and init the db if the user provided a database
	if !db.InitializeDB(conf.DBName) {
		return
	}

	// if the c2 server is meant to catch responses, initialize and start so it can bind
	if conf.C2AutoStart {
		if !startC2Server(conf) {
			return
		}
	}

	if conf.ExType == config.FileFormat || conf.ExType == config.Local {
		if !doScan(sploit, conf) {
			return
		}
		globalWG.Wait()
	} else {
		// loop over all the provided host / port combos
		for index, host := range conf.RhostsNTuple {
			// setup the conf for the downstream exploit
			conf.Rhost = host.Rhost
			conf.Rport = host.Rport
			switch host.SSL {
			case config.SSLDisabled:
				conf.SSL = false
				conf.DetermineSSL = false
			case config.SSLEnabled:
				conf.SSL = true
				conf.DetermineSSL = false
			case config.SSLAutodiscover:
				conf.SSL = false
				conf.DetermineSSL = true
			}

			output.PrintFrameworkStatus("Starting target", "index", index, "host", conf.Rhost,
				"port", conf.Rport, "ssl", conf.SSL, "ssl auto", conf.DetermineSSL)
			if !doScan(sploit, conf) {
				return
			}
			globalWG.Wait()
		}
	}
}