connect and run commands working

bot/bot.go

@@ -19,7 +19,6 @@ var (
 )
 
 func Run() {
-	pb.Init()
 	discord, err := discordgo.New("Bot " + BotToken)
 	if err != nil {
 		log.Fatal(err)
@@ -27,16 +26,23 @@ func Run() {
 
 	discord.AddHandler(newMessage)
 
+	log.Info("Opening Discord connection...")
 	discord.Open()
 	defer discord.Close()
 	log.Info("BitBot is running...")
 
+	// Try initializing PocketBase after Discord is connected
+	log.Info("Initializing PocketBase...")
+	pb.Init()
+	log.Info("PocketBase initialized successfully.")
+
 	c := make(chan os.Signal, 1)
 	signal.Notify(c, os.Interrupt)
 	<-c
 }
 
 var conversationHistoryMap = make(map[string][]openai.ChatCompletionMessage)
+var sshConnections = make(map[string]*SSHConnection)
 
 func newMessage(discord *discordgo.Session, message *discordgo.MessageCreate) {
 	if message.Author.ID == discord.State.User.ID || message.Content == "" {
@@ -105,16 +111,42 @@ func newMessage(discord *discordgo.Session, message *discordgo.MessageCreate) {
 			}
 
 			connectionDetails := commandParts[1]
-			err := SSHConnectToRemoteServer(connectionDetails)
+			sshConn, err := SSHConnectToRemoteServer(connectionDetails)
 			if err != nil {
 				discord.ChannelMessageSend(message.ChannelID, "Error connecting to remote server.")
 				return
 			}
 
+			// Store the SSH connection for later use
+			sshConnections[message.Author.ID] = sshConn
+
 			discord.ChannelMessageSend(message.ChannelID, "Connected to remote server!")
 		} else {
 			discord.ChannelMessageSend(message.ChannelID, "You are not authorized to use this command.")
 		}
+	} else if strings.HasPrefix(message.Content, "!exe") {
+		if message.Author.ID == AllowedUserID {
+			// Check if there is an active SSH connection for this user
+			sshConn, ok := sshConnections[message.Author.ID]
+			if !ok {
+				discord.ChannelMessageSend(message.ChannelID, "You are not connected to any remote server. Use !ssh first.")
+				return
+			}
+
+			// Extract the command after "!exe"
+			command := strings.TrimPrefix(message.Content, "!exe ")
+
+			// Execute the command on the remote server
+			response, err := sshConn.ExecuteCommand(command)
+			if err != nil {
+				discord.ChannelMessageSend(message.ChannelID, "Error executing command on remote server.")
+				return
+			}
+
+			discord.ChannelMessageSend(message.ChannelID, "Remote server response: "+response)
+		} else {
+			discord.ChannelMessageSend(message.ChannelID, "You are not authorized to use this command.")
+		}
 	}
 
 	conversationHistoryMap[userID] = conversationHistory

bot/sshclient.go

@@ -9,6 +9,7 @@ import (
 	"os"
 	"strings"
 
+	"github.com/charmbracelet/log"
 	"golang.org/x/crypto/ssh"
 )
 
@@ -18,6 +19,48 @@ const (
 	bits           = 2048
 )
 
+type SSHConnection struct {
+	client    *ssh.Client
+	commands  chan string
+	responses chan string
+}
+
+func NewSSHConnection(client *ssh.Client) *SSHConnection {
+	return &SSHConnection{
+		client:    client,
+		commands:  make(chan string),
+		responses: make(chan string),
+	}
+}
+
+func (conn *SSHConnection) startCommandExecution() {
+	for cmd := range conn.commands {
+		// Execute the command and send the response back
+		response, err := executeRemoteCommand(conn.client, cmd)
+		if err != nil {
+			log.Error(err)
+			conn.responses <- "Error executing command"
+		} else {
+			conn.responses <- response
+		}
+	}
+}
+
+func executeRemoteCommand(client *ssh.Client, command string) (string, error) {
+	session, err := client.NewSession()
+	if err != nil {
+		return "", err
+	}
+	defer session.Close()
+
+	output, err := session.CombinedOutput(command)
+	if err != nil {
+		return "", err
+	}
+
+	return string(output), nil
+}
+
 func GenerateAndSaveSSHKeyPairIfNotExist() error {
 	if KeyFilesExist(privateKeyPath, publicKeyPath) {
 		return nil
@@ -118,10 +161,10 @@ func SavePublicKeyToFile(filename string, publicKey ssh.PublicKey) error {
 	return nil
 }
 
-func SSHConnectToRemoteServer(connectionDetails string) error {
+func SSHConnectToRemoteServer(connectionDetails string) (*SSHConnection, error) {
 	privateKey, err := LoadPrivateKey(privateKeyPath)
 	if err != nil {
-		return err
+		return nil, err
 	}
 
 	config := &ssh.ClientConfig{
@@ -136,7 +179,7 @@ func SSHConnectToRemoteServer(connectionDetails string) error {
 	// Extract username, host, and port from connectionDetails
 	parts := strings.Split(connectionDetails, "@")
 	if len(parts) != 2 {
-		return errors.New("invalid connection format")
+		return nil, errors.New("invalid connection format")
 	}
 	username := parts[0]
 	hostPort := parts[1]
@@ -146,14 +189,26 @@ func SSHConnectToRemoteServer(connectionDetails string) error {
 	// Connect to the remote server
 	client, err := ssh.Dial("tcp", hostPort, config)
 	if err != nil {
-		return err
+		return nil, err
 	}
-	defer client.Close()
 
-	// Now you have a connected SSH client to the remote server.
-	// You can use this client to perform remote commands or file transfers.
+	// Create an SSH connection instance
+	conn := NewSSHConnection(client)
 
-	return nil
+	// Start the goroutine for command execution
+	go conn.startCommandExecution()
+
+	// Now you have an active SSH connection with command execution capabilities
+	return conn, nil
+}
+
+func (conn *SSHConnection) ExecuteCommand(command string) (string, error) {
+	// Send the command to the goroutine for execution
+	conn.commands <- command
+
+	// Receive the response
+	response := <-conn.responses
+	return response, nil
 }
 
 func LoadPrivateKey(path string) (ssh.Signer, error) {

go.mod

@@ -101,6 +101,6 @@ require (
 	github.com/gorilla/websocket v1.5.0 // indirect
 	github.com/joho/godotenv v1.5.1
 	github.com/sashabaranov/go-openai v1.14.1
-	golang.org/x/crypto v0.12.0 // indirect
+	golang.org/x/crypto v0.12.0
 	golang.org/x/sys v0.11.0 // indirect
 )