package simpleshell

import (
	"net"
	"sync/atomic"
	"time"

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

type Client struct {
	ConnectAddr string
	ConnectPort int
	channel     *channel.Channel
}

var clientSingleton *Client

func GetClientInstance() *Client {
	if clientSingleton == nil {
		clientSingleton = new(Client)
	}

	return clientSingleton
}

func (shellClient *Client) CreateFlags() {
}

func (shellClient *Client) Shutdown() bool {
	// Account for non-running case
	if shellClient.Channel() == nil {
		return true
	}
	ok := shellClient.Channel().RemoveSessions()

	// we done here
	output.PrintfFrameworkStatus("Connection closed")

	return ok
}

func (shellClient *Client) Channel() *channel.Channel {
	return shellClient.channel
}

func (shellClient *Client) 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
	}

	shellClient.ConnectAddr = channel.IPAddr
	shellClient.ConnectPort = channel.Port
	shellClient.channel = channel

	if !channel.IsClient {
		output.PrintFrameworkError("Called SimpleShellClient as a server. Use bport.")

		return false
	}

	if channel.Port == 0 {
		output.PrintFrameworkError("Provided an invalid bind port.")

		return false
	}

	return true
}

func (shellClient *Client) Run(timeout int) {
	conn, ok := connect(shellClient.ConnectAddr, shellClient.ConnectPort, timeout)
	if !ok {
		return
	}
	// Track if the C2 is indicated to shutdown for any reason.
	go func() {
		for {
			if shellClient.Channel().Shutdown.Load() {
				shellClient.Shutdown()

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

	output.PrintfFrameworkStatus("Active shell on %s:%d", shellClient.ConnectAddr, shellClient.ConnectPort)
	shellClient.Channel().AddSession(&conn, conn.RemoteAddr().String())

	cli.Basic(conn, shellClient.channel)

	shellClient.Channel().Shutdown.Store(true)
}

func connect(host string, port int, timeout int) (net.Conn, bool) {
	// loop every three seconds until timeout
	for i := 0; i < timeout; i += 3 {
		// TCPConnect is proxy aware, so it will use the proxy if configured
		conn, ok := protocol.TCPConnect(host, port)
		if ok {
			output.PrintfFrameworkSuccess("Connected to %s:%d!", host, port)

			return conn, true
		}
		// this is technically not correct. TCPConnect itself has a 10 second timeout. If the server doesn't
		// send a RST when we try to connect then we'll wait the full 10 + this 3 for a full 13... which means
		// we won't honor the timeout correctly.
		time.Sleep(3 * time.Second)
	}

	return nil, false
}