package config

import (
	"bytes"
	"flag"
	"fmt"
	"strings"
	"text/template"

	"github.com/vulncheck-oss/go-exploit/c2"
	"github.com/vulncheck-oss/go-exploit/c2/shelltunnel"
	"github.com/vulncheck-oss/go-exploit/output"
	"github.com/vulncheck-oss/go-exploit/protocol"
)

type ExploitType int

const (
	CodeExecution         ExploitType = 0
	InformationDisclosure ExploitType = 1
	Webshell              ExploitType = 2
	FileFormat            ExploitType = 3
	Local                 ExploitType = 4
)

type ImplementedFeatures struct {
	AssetDetection  bool
	VersionScanning bool
	Exploitation    bool
}

type SSLSupport int

const (
	SSLDisabled     SSLSupport = 0
	SSLEnabled      SSLSupport = 1
	SSLAutodiscover SSLSupport = 2
)

type RhostTriplet struct {
	Rhost string
	Rport int
	SSL   SSLSupport
}

// The config struct contains a mix of module specified configurations
// and user specified configurations. The Config struct is first generated
// by the exploit implementation and then modified by option parsing.
type Config struct {
	// the following are values configured by the exploit module

	// implemented features describes which three stages the exploit implements
	Impl ImplementedFeatures
	// the vendor of the targeted product
	Vendor string
	// the targeted products
	Products []string
	// A combination of the Vendor and Products strings
	Product string
	// the CPE for the targeted product
	CPE []string
	// the CVE being tested
	CVE string
	// the protocol being targeted
	Protocol string
	// the type of exploit being executed
	ExType ExploitType
	// the c2 supported by the exploit
	SupportedC2 []c2.Impl

	// Some exploits need to define custom flags. Use the Create*Flag functions
	// to store them in the following data structures. They can then be fetched
	// using the Get*Flag functions

	StringFlagsMap map[string]*string
	IntFlagsMap    map[string]*int
	UintFlagsMap   map[string]*uint
	BoolFlagsMap   map[string]*bool

	// the following are values configured by the user

	// target host, the target address/name the exploit will work on
	Rhost string
	// target port, the target port the exploit will work on
	Rport int
	// a list of specific targets
	RhostsNTuple []RhostTriplet
	// local host for remote exploits
	Lhost string
	// local port
	Lport int
	// bind port
	Bport int
	// indicates if the framework should autodetect ssl/plain
	DetermineSSL bool
	// indicates if ssl is used in comms
	SSL bool
	// indicates if we run the target verify
	DoVerify bool
	// indicates if we run the version check
	DoVersionCheck bool
	// indicates if we run the exploit
	DoExploit bool
	// automatically start the c2 or not
	C2AutoStart bool
	// the user requested c2 to use
	C2Type c2.Impl
	// C2 server timeout
	C2Timeout int
	// Indicates if the c2 server will be handled elsewhere
	ThirdPartyC2Server bool
	// The database we are working with
	DBName string
	// File format template
	FileTemplateData string
	// File format exploit output
	FileFormatFilePath string
}

// Convert ExploitType to String.
func (eType ExploitType) String() string {
	switch eType {
	case CodeExecution:
		return "CodeExecution"
	case InformationDisclosure:
		return "InformationDisclosure"
	case Webshell:
		return "Webshell"
	case FileFormat:
		return "FileFormat"
	case Local:
		return "Local"
	default:
		return "Invalid exploit type"
	}
}

// Deprecated: New does not affectively describe the affected/targeted product. Use NewRemoteExploit.
func New(extype ExploitType, supportedC2 []c2.Impl, product string, cve string, defaultPort int) *Config {
	returnVal := new(Config)
	returnVal.ExType = extype
	returnVal.SupportedC2 = supportedC2
	returnVal.Product = product
	returnVal.CVE = cve
	returnVal.Rport = defaultPort

	return returnVal
}

// Deprecated: NewLocal does not affectively describe the affected/targeted product. Use NewLocalExploit.
func NewLocal(extype ExploitType, supportedC2 []c2.Impl, product string, cve string) *Config {
	returnVal := new(Config)
	returnVal.ExType = extype
	returnVal.SupportedC2 = supportedC2
	returnVal.Product = product
	returnVal.CVE = cve

	return returnVal
}

// Defines a new remote exploit and associates with CVE/Product/Protocol metadata. Usage example:
//
//	conf := config.NewRemoteExploit(
//	  config.ImplementedFeatures{AssetDetection: true, VersionScanning: true, Exploitation: true},
//	  config.CodeExecution, []c2.Impl{c2.SimpleShellServer},
//	  "Atlassian", []string{"Confluence"}, []string{"cpe:2.3:a:atlassian:confluence"},
//	  "CVE-2023-22527", "HTTP", 8090)
func NewRemoteExploit(implemented ImplementedFeatures, extype ExploitType, supportedC2 []c2.Impl, vendor string,
	product []string, cpe []string, cve string, protocol string, defaultPort int,
) *Config {
	newConf := new(Config)
	newConf.Product = deDupeProductName(product, vendor)
	newConf.InitFlagsStructs()
	newConf.Impl = implemented
	newConf.ExType = extype
	newConf.SupportedC2 = supportedC2
	newConf.Vendor = vendor
	newConf.Products = product
	newConf.C2AutoStart = true
	newConf.CPE = cpe
	newConf.CVE = cve
	newConf.Protocol = protocol
	newConf.Rport = defaultPort

	return newConf
}

func deDupeProductName(product []string, vendor string) string {
	joinedProducts := strings.Join(product, "/")
	if joinedProducts == vendor {
		return vendor
	}

	return fmt.Sprintf("%s %s", vendor, joinedProducts)
}

// Defines a new remote exploit and associates with CVE/Product/Protocol metadata. Usage example:.
func NewLocalExploit(implemented ImplementedFeatures, extype ExploitType, supportedC2 []c2.Impl, vendor string,
	product []string, cpe []string, cve string,
) *Config {
	newConf := new(Config)
	newConf.Product = deDupeProductName(product, vendor)
	newConf.InitFlagsStructs()
	newConf.Impl = implemented
	newConf.ExType = extype
	newConf.SupportedC2 = supportedC2
	newConf.Vendor = vendor
	newConf.Products = product
	newConf.C2AutoStart = true
	newConf.CPE = cpe
	newConf.CVE = cve

	return newConf
}

func (conf *Config) InitFlagsStructs() {
	conf.StringFlagsMap = make(map[string]*string)
	conf.IntFlagsMap = make(map[string]*int)
	conf.UintFlagsMap = make(map[string]*uint)
	conf.BoolFlagsMap = make(map[string]*bool)
}

// Create a command line flag for the string var "name" with the default value of "value" and
// store the result locally.
func (conf *Config) CreateStringFlag(name string, value string, usage string) {
	valuePtr := &value
	conf.StringFlagsMap[name] = valuePtr
	flag.StringVar(conf.StringFlagsMap[name], name, value, usage)
}

// Create a command line flag for the string var "name" with the default value of "value" and
// store the result locally *using an external "param" pointer*.
func (conf *Config) CreateStringVarFlag(param *string, name string, value string, usage string) {
	conf.StringFlagsMap[name] = param
	flag.StringVar(param, name, value, usage)
}

// Create a command line flag for the uint var "name" with the default value of "value" and
// store the result locally.
func (conf *Config) CreateUintFlag(name string, value uint, usage string) {
	valuePtr := &value
	conf.UintFlagsMap[name] = valuePtr
	flag.UintVar(conf.UintFlagsMap[name], name, value, usage)
}

// Create a command line flag for the uint var "name" with the default value of "value" and
// store the result locally *using an external "param" pointer*.
func (conf *Config) CreateUintVarFlag(param *uint, name string, value uint, usage string) {
	conf.UintFlagsMap[name] = param
	flag.UintVar(param, name, value, usage)
}

// Create a command line flag for the int var "name" with the default value of "value" and
// store the result locally.
func (conf *Config) CreateIntFlag(name string, value int, usage string) {
	valuePtr := &value
	conf.IntFlagsMap[name] = valuePtr
	flag.IntVar(conf.IntFlagsMap[name], name, value, usage)
}

// Create a command line flag for the int var "name" with the default value of "value" and
// store the result locally *using an external "param" pointer*.
func (conf *Config) CreateIntVarFlag(param *int, name string, value int, usage string) {
	conf.IntFlagsMap[name] = param
	flag.IntVar(param, name, value, usage)
}

// Create a command line flag for the bool var "name" with the default value of "value" and
// store the result locally.
func (conf *Config) CreateBoolFlag(name string, value bool, usage string) {
	valuePtr := &value
	conf.BoolFlagsMap[name] = valuePtr
	flag.BoolVar(conf.BoolFlagsMap[name], name, value, usage)
}

// Create a command line flag for the bool var "name" with the default value of "value" and
// store the result locally *using an external "param" pointer*.
func (conf *Config) CreateBoolVarFlag(param *bool, name string, value bool, usage string) {
	conf.BoolFlagsMap[name] = param
	flag.BoolVar(param, name, value, usage)
}

// Fetch the configured string value for "name".
func (conf *Config) GetStringFlag(name string) string {
	value, ok := conf.StringFlagsMap[name]
	if !ok {
		output.PrintfFrameworkError("Requested invalid flag: %s", name)

		return ""
	}

	return *value
}

// Fetch the configured uint value for "name".
func (conf *Config) GetUintFlag(name string) uint {
	value, ok := conf.UintFlagsMap[name]
	if !ok {
		output.PrintfFrameworkError("Requested invalid flag: %s", name)

		return 0
	}

	return *value
}

// Fetch the configured uint value for "name".
func (conf *Config) GetIntFlag(name string) int {
	value, ok := conf.IntFlagsMap[name]
	if !ok {
		output.PrintfFrameworkError("Requested invalid flag: %s", name)

		return 0
	}

	return *value
}

// Fetch the configured uint value for "name".
func (conf *Config) GetBoolFlag(name string) bool {
	value, ok := conf.BoolFlagsMap[name]
	if !ok {
		output.PrintfFrameworkError("Requested invalid flag: %s", name)

		return false
	}

	return *value
}

// Apply the configuration settings to a Go text template. This will take
// the `Config` struct and apply it to a `text/template`, allowing for
// strings to be built directly from the already set configuration
// variables.
//
//	s := conf.ApplyTemplate(`CVE: {{.CVE}} - {{.Product}}`)
//	output.PrintStatus(s) // Output: CVE: CVE-2024-1337 - OFBiz
//
// Flags that are user defined with CreateStringFlag and other types are
// directly accessible from their map values, for example if a command line
// argument is added with conf.CreateStringFlag("output", "do output",
// "instructions") it will be accessible via the following ApplyTemplate
// call:
//
//	conf.ApplyTemplate(`Output flag {{.StringFlagsMap.output}}`)
//
// This function only returns the processed string and if a templating
// error occurs the function emits a framework error and sets the string to
// an empty string. This makes it harder to process any dynamic content and
// properly catch errors, but simplifies the return value to only provide a
// string.
//
// This should not be used with potentially attacker controlled input.
//
// Some Config types might be complex and will require usage of range
// components of text/template, follow the package docs if necessary.
func (conf *Config) ApplyTemplate(name string) string {
	t, err := template.New("config-string-template").Parse(name)
	if err != nil {
		output.PrintfFrameworkError("Could not create template: %s", err.Error())

		return ""
	}
	var buf bytes.Buffer
	if err := t.Execute(&buf, conf); err != nil {
		output.PrintfFrameworkError("Could not apply template: %s", err.Error())

		return ""
	}

	return buf.String()
}

// Generate a URL from a path from the current configuration. This is a
// way of invoking protocol.GenerateURL for developer ergonomics during
// exploit development.
func (conf *Config) GenerateURL(path string) string {
	return protocol.GenerateURL(conf.Rhost, conf.Rport, conf.SSL, path)
}

// Disable automatic start of c2 servers. Manually starting is required after
// this function is called. This is useful when you have an exploit that
// may have multiple stages and you are guaranteed to not need the C2
// setup. An example is an exploit that needs to retrieve a CAPTCHA may not
// want to start up the C2 until the first stage is retrieved and the
// CAPTCHA is solved.
func (conf *Config) DisableC2Start() {
	conf.C2AutoStart = false
}

// Some C2 (ShellTunnel) don't actually care how the payload is generated, but
// the underlying C2 might be implied depending on how the individual exploit
// has been developed. It is certainly not a requirement to call this function
// but it can help simplify the handling of secure shell vs insecure.
func (conf *Config) ResolveC2Payload() c2.Impl {
	if conf.C2Type != c2.ShellTunnel {
		return conf.C2Type
	}

	if shelltunnel.GetInstance().SSLShellServer {
		return c2.SSLShellServer
	}

	return c2.SimpleShellServer
}