package simpleshell

import (
	"net"
	"strconv"
	"strings"
	"sync"
	"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"
)

// The SimpleShellServer implements a basic reverse shell catcher. The server listens on a user provided
// port and catches incoming unencrypted connections. The c2 implements a basic command line that lacks
// any type of useful bash-like features (history, interactive behavior, displaying of stderr, etc).
// The server can accept multiple connections, but the user has no way of swapping between them unless
// the terminate the connection.
type Server struct {
	Listener net.Listener
	channel  *channel.Channel
}

var serverSingleton *Server

// A basic singleton interface for the c2.
func GetServerInstance() *Server {
	if serverSingleton == nil {
		serverSingleton = new(Server)
	}

	return serverSingleton
}

// User options for the simple shell server (currently empty).
func (shellServer *Server) CreateFlags() {
}

// Get the underlying C2 channel with metadata and session information.
func (shellServer *Server) Channel() *channel.Channel {
	return shellServer.channel
}

// Shutdown the C2 and cleanup any active connections.
func (shellServer *Server) Shutdown() bool {
	// Account for non-running case
	if shellServer.Channel() == nil {
		return true
	}
	output.PrintFrameworkStatus("C2 received shutdown, killing server and client sockets for shell server")
	if len(shellServer.Channel().Sessions) > 0 {
		for k, session := range shellServer.Channel().Sessions {
			output.PrintfFrameworkStatus("Connection closed for shell server: %s", session.RemoteAddr)
			shellServer.Channel().RemoveSession(k)
		}
	}
	shellServer.Listener.Close()

	return true
}

// Validate configuration and create the listening socket.
func (shellServer *Server) 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
	}
	shellServer.channel = channel
	if channel.IsClient {
		output.PrintFrameworkError("Called SimpleShellServer as a client. Use lhost and lport.")

		return false
	}

	output.PrintfFrameworkStatus("Starting listener on %s:%d", channel.IPAddr, channel.Port)

	var err error
	shellServer.Listener, err = net.Listen("tcp", channel.IPAddr+":"+strconv.Itoa(channel.Port))
	if err != nil {
		output.PrintFrameworkError("Couldn't create the server: " + err.Error())

		return false
	}

	return true
}

// Listen for incoming.
func (shellServer *Server) Run(timeout int) {
	// mutex for user input
	var cliLock sync.Mutex

	// terminate the server if no shells come in within timeout seconds
	go func() {
		time.Sleep(time.Duration(timeout) * time.Second)
		if !shellServer.Channel().HasSessions() {
			output.PrintFrameworkError("Timeout met. Shutting down shell listener.")
			shellServer.Channel().Shutdown.Store(true)
		}
	}()
	// Track if the shutdown is signaled for any reason.
	go func() {
		for {
			if shellServer.Channel().Shutdown.Load() {
				shellServer.Shutdown()

				break
			}
			time.Sleep(10 * time.Millisecond)
		}
	}()
	// Accept arbitrary connections. In the future we need something for the
	// user to select which connection to make active
	for {
		client, err := shellServer.Listener.Accept()
		if err != nil {
			if !strings.Contains(err.Error(), "use of closed network connection") {
				output.PrintFrameworkError(err.Error())
			}

			break
		}
		output.PrintfFrameworkSuccess("Caught new shell from %v", client.RemoteAddr())
		// Add the session for tracking first, so it can be cleaned up.
		shellServer.Channel().AddSession(&client, client.RemoteAddr().String())

		go handleSimpleConn(client, &cliLock, client.RemoteAddr(), shellServer.channel)
	}
}

func handleSimpleConn(conn net.Conn, cliLock *sync.Mutex, remoteAddr net.Addr, ch *channel.Channel) {
	// connections will stack up here. Currently that will mean a race
	// to the next connection but we can add in attacker handling of
	// connections latter
	cliLock.Lock()
	defer cliLock.Unlock()

	// close the connection when the shell is complete
	defer conn.Close()

	// Only complete the full session handshake once
	if !ch.Shutdown.Load() {
		output.PrintfFrameworkStatus("Active shell from %v", remoteAddr)
		cli.Basic(conn, ch)
	}
}