package c2

import (
	"github.com/vulncheck-oss/go-exploit/c2/channel"
	"github.com/vulncheck-oss/go-exploit/c2/external"
	"github.com/vulncheck-oss/go-exploit/c2/httpservefile"
	"github.com/vulncheck-oss/go-exploit/c2/httpserveshell"
	"github.com/vulncheck-oss/go-exploit/c2/httpshellserver"
	"github.com/vulncheck-oss/go-exploit/c2/shelltunnel"
	"github.com/vulncheck-oss/go-exploit/c2/simpleshell"
	"github.com/vulncheck-oss/go-exploit/c2/sslshell"
	"github.com/vulncheck-oss/go-exploit/output"
)

// An interface used by both reverse shells, bind shells, and stagers.
type Interface interface {
	CreateFlags()
	Init(channel *channel.Channel) bool
	Run(timeout int)
	Shutdown() bool
	Channel() *channel.Channel
}

// Internal representation of a C2 implementation. Each C2 is managed by
// the framework by the Impl struct type and should be unique.
type Impl struct {
	Name     string
	Category category
}

// Categories allow for type hints and invalid checks.
type category int

const (
	InvalidCategory           category = -1
	SimpleShellServerCategory category = 0
	SimpleShellClientCategory category = 1
	SSLShellServerCategory    category = 2
	HTTPServeFileCategory     category = 3
	HTTPServeShellCategory    category = 4
	ExternalCategory          category = 5
	ShellTunnelCategory       category = 6
	HTTPShellServerCategory   category = 7
)

// Simplified names in order to keep the old calling convention and allow
// for quick references in supported C2 functions.
var (
	SimpleShellServer = internalSupported["SimpleShellServer"]
	SimpleShellClient = internalSupported["SimpleShellClient"]
	SSLShellServer    = internalSupported["SSLShellServer"]
	HTTPServeFile     = internalSupported["HTTPServeFile"]
	HTTPServeShell    = internalSupported["HTTPServeShell"]
	ShellTunnel       = internalSupported["ShellTunnel"]
	HTTPShellServer   = internalSupported["HTTPShellServer"]
	// We do not want external to be called directly because external
	// internally is not useful.
)

// The internal representation and model for keeping track of C2s. This is
// modified by external modules, but first-party framework supported
// channels are defined here.
var internalSupported = map[string]Impl{
	"SimpleShellServer": {Name: "SimpleShellServer", Category: SimpleShellServerCategory},
	"SimpleShellClient": {Name: "SimpleShellClient", Category: SimpleShellClientCategory},
	"SSLShellServer":    {Name: "SSLShellServer", Category: SSLShellServerCategory},
	"HTTPServeFile":     {Name: "HTTPServeFile", Category: HTTPServeFileCategory},
	"HTTPServeShell":    {Name: "HTTPServeShell", Category: HTTPServeShellCategory},
	"HTTPShellServer":   {Name: "HTTPShellServer", Category: HTTPShellServerCategory},
	// Ensure the internal supported External module name is an error if used
	// directly.
	"External":    {Name: "", Category: InvalidCategory},
	"ShellTunnel": {Name: "ShellTunnel", Category: ShellTunnelCategory},
}

// Add an external C2 to the supported list. Use this to integrate a new C2
// into the framework. This is expected to be called early in the
// configuration process in order to expose the name of the C2 to the rest
// of the framework.
func AddC2(name string) Impl {
	_, exists := internalSupported[name]
	if exists {
		// We might not want to panic, but it does simplify the call
		panic("C2 type already exists")
	}

	i := Impl{Name: name, Category: ExternalCategory}
	internalSupported[name] = i

	return i
}

// Factory pattern for creating c2 interfaces. Note that this is
// returning an interface, which is a bit anti-Go but it's more or less
// exactly what we want so.
func GetInstance(implementation Impl) (Interface, bool) {
	switch implementation.Category {
	case SimpleShellServerCategory:
		return simpleshell.GetServerInstance(), true
	case SimpleShellClientCategory:
		return simpleshell.GetClientInstance(), true
	case SSLShellServerCategory:
		return sslshell.GetInstance(), true
	case HTTPServeFileCategory:
		return httpservefile.GetInstance(), true
	case HTTPServeShellCategory:
		return httpserveshell.GetInstance(), true
	case ExternalCategory:
		if implementation.Name != "" {
			return external.GetInstance(implementation.Name), true
		}
	case HTTPShellServerCategory:
		return httpshellserver.GetInstance(), true
	case ShellTunnelCategory:
		return shelltunnel.GetInstance(), true
	case InvalidCategory:
		// Calling your external C2 as explicitly invalid is odd.
		output.PrintFrameworkError("Invalid C2 Server")
	default:
		output.PrintFrameworkError("Invalid C2 Server")
	}

	return nil, false
}

// Call into the c2 impl so that it can create command line flags.
func CreateFlags(implementation Impl) {
	switch implementation.Category {
	case SimpleShellServerCategory:
		simpleshell.GetServerInstance().CreateFlags()
	case SimpleShellClientCategory:
		simpleshell.GetClientInstance().CreateFlags()
	case SSLShellServerCategory:
		sslshell.GetInstance().CreateFlags()
	case HTTPServeFileCategory:
		httpservefile.GetInstance().CreateFlags()
	case HTTPServeShellCategory:
		httpserveshell.GetInstance().CreateFlags()
	case ExternalCategory:
		if implementation.Name != "" {
			external.GetInstance(implementation.Name).CreateFlags()
		}
	case HTTPShellServerCategory:
		httpshellserver.GetInstance().CreateFlags()
	case ShellTunnelCategory:
		shelltunnel.GetInstance().CreateFlags()
	case InvalidCategory:
		// Calling your external C2 as explicitly invalid is odd.
		output.PrintFrameworkError("Invalid C2 Server")
	default:
		output.PrintFrameworkError("Invalid C2 Server")
	}
}

// HasSessions returns if the underlying channel has active sessions. This is useful for code that
// needs to validate if callbacks have occurred and is a helper wrapper around the channel package
// function of the same name.
func HasSessions(implementation Impl) bool {
	switch implementation.Category {
	case SimpleShellServerCategory:
		return simpleshell.GetServerInstance().Channel().HasSessions()
	case SimpleShellClientCategory:
		return simpleshell.GetClientInstance().Channel().HasSessions()
	case SSLShellServerCategory:
		return sslshell.GetInstance().Channel().HasSessions()
	case HTTPServeFileCategory:
		return httpservefile.GetInstance().Channel().HasSessions()
	case HTTPServeShellCategory:
		return httpserveshell.GetInstance().Channel().HasSessions()
	case ExternalCategory:
		if implementation.Name != "" {
			return external.GetInstance(implementation.Name).Channel().HasSessions()
		}
	case HTTPShellServerCategory:
		return httpshellserver.GetInstance().Channel().HasSessions()
	case ShellTunnelCategory:
		return shelltunnel.GetInstance().Channel().HasSessions()
	case InvalidCategory:
	default:
	}
	output.PrintFrameworkError("Invalid C2 Server")

	return false
}

// Return the internal representation of a C2 from a string.
func StringToImpl(c2Name string) (Impl, bool) {
	for _, value := range internalSupported {
		if value.Name == c2Name {
			return value, true
		}
	}

	return Impl{Name: "", Category: InvalidCategory}, false
}