package cli

import (
	"bufio"
	"net"
	"os"
	"sync"
	"testing"
	"time"

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

// backgroundResponse handles the network connection reading for response data and contains a
// trigger to the shutdown of the channel to ensure cleanup happens on socket close.
func backgroundResponse(ch *channel.Channel, wg *sync.WaitGroup, conn net.Conn, responseCh chan string) {
	defer wg.Done()
	defer func(channel *channel.Channel) {
		// Signals for both routines to stop, this should get triggered when socket is closed
		// and causes it to fail the read
		channel.Shutdown.Store(true)
	}(ch)
	responseBuffer := make([]byte, 1024)
	for {
		if ch.Shutdown.Load() {
			return
		}

		err := conn.SetReadDeadline(time.Now().Add(1 * time.Second))
		if err != nil {
			output.PrintfFrameworkError("Error setting read deadline: %s, exiting.", err)

			return
		}

		bytesRead, err := conn.Read(responseBuffer)
		if err != nil && !os.IsTimeout(err) {
			// things have gone sideways, but the command line won't know that
			// until they attempt to execute a command and the socket fails.
			// i think that's largely okay.
			return
		}

		if bytesRead > 0 {
			// I think there is technically a race condition here where the socket
			// could have move data to write, but the user has already called exit
			// below. I that that's tolerable for now.
			responseCh <- string(responseBuffer[:bytesRead])
			// Update "Last Seen"
			ok := ch.UpdateLastSeenByConn(conn, time.Now())
			if !ok {
				output.PrintFrameworkError("Failed to update LastSeen value for connection")

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

// A very basic reverse/bind shell handler.
func Basic(conn net.Conn, ch *channel.Channel) {
	// Create channels for communication between goroutines.
	responseCh := make(chan string)

	// Use a WaitGroup to wait for goroutines to finish.
	var wg sync.WaitGroup

	// Goroutine to read responses from the server.
	wg.Add(1)

	// If running in the test context inherit the channel input setting, this will let us control the
	// input of the shell programmatically.
	if !testing.Testing() {
		ch.Input = os.Stdin
	}
	go backgroundResponse(ch, &wg, conn, responseCh)

	// Goroutine to handle responses and print them.
	wg.Add(1)
	go func(channel *channel.Channel) {
		defer wg.Done()
		for {
			if channel.Shutdown.Load() {
				return
			}
			select {
			case response := <-responseCh:
				output.PrintShell(response)
			default:
			}
			time.Sleep(10 * time.Millisecond)
		}
	}(ch)

	go func(channel *channel.Channel) {
		// no waitgroup for this one because blocking IO, but this should not matter
		// since we are intentionally not trying to be a multi-implant C2 framework.
		// There still remains the issue that you would need to hit enter to find out
		// that the socket is dead but at least we can stop Basic() regardless of this fact.
		// This issue of unblocking stdin is discussed at length here https://github.com/golang/go/issues/24842
		for {
			reader := bufio.NewReader(ch.Input)
			command, _ := reader.ReadString('\n')
			if channel.Shutdown.Load() {
				break
			}
			if command == "exit\n" {
				channel.Shutdown.Store(true)

				break
			}
			ok := protocol.TCPWrite(conn, []byte(command))
			if !ok {
				channel.Shutdown.Store(true)

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

	// wait until the go routines are clean up
	wg.Wait()
	close(responseCh)
}