From fbdba2b92d2ff27d7718ef9dbfdfe0a0bbf4bb45 Mon Sep 17 00:00:00 2001 From: Kristian _server Date: Fri, 6 Oct 2023 20:32:18 +0200 Subject: [PATCH] connect and run commands working --- bot/bot.go | 36 ++++++++++++++++++++++-- bot/sshclient.go | 71 +++++++++++++++++++++++++++++++++++++++++------ go.mod | 2 +- pb_data/logs.db | Bin 53248 -> 53248 bytes 4 files changed, 98 insertions(+), 11 deletions(-) diff --git a/bot/bot.go b/bot/bot.go index 44be63c..f10a8b2 100644 --- a/bot/bot.go +++ b/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 diff --git a/bot/sshclient.go b/bot/sshclient.go index 61f0c10..aa3d98e 100644 --- a/bot/sshclient.go +++ b/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) { diff --git a/go.mod b/go.mod index 77de328..421810c 100644 --- a/go.mod +++ b/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 ) diff --git a/pb_data/logs.db b/pb_data/logs.db index 81d15245f18ee9f75ced869241a68299b6711526..6f42c57f4f14b89b6f76ea9638a5d63a202b62bd 100644 GIT binary patch delta 1162 zcmZozz}&Eac>|+?D0c`03zq-`mjK@$-U=QQE&&kSEM{<>nT3IYO?mT2V-7|pR>{r3 zO+;B0n;f0_wHs}PnZ)I78^Z-{ZRJxeb5fJcay7KR4824)I|rdB4FR;Ctu#zw|bMg~U4x>&>wjm-57fY6YEf#JdA2Su|-v8eoC%+PL`QTRcbcEA8yG0FtpS&H8RBX z2UHwxJ{K^=Sok+E@ZaFy!2gE-2>+4ILJt4oCP)X1OxvGpqroYUjx!WxLK^g ziC+k4SOGkYK}IJ{&Z>{|<7Jj-gbO>P@Z3;%t|&aHRSOvS7x2H}zreo(NCV+!v4&H8 zGC;!*fb)fc1u$zU7?@cZ7yuE-(iM~c^ew3uXEtU;%{X8M&N#%~aEQC&5XTV?{5DXy US-!v$SAbZcc$0(xS+u|a0EO)xR{#J2 delta 1071 zcmbW0%WD%+6vk&VP1MxRvCk9{+J#CXB%M2x$4uJM%q(0~umKB}LT&RhO(xBx>C`?z z`&f0=qF2O?C5qjsi@+@O4{%Xj34(UvLc~Qu1r_Q8A9yE$v8}Gk8See=H+O#L+;f?Q zGP6)-Dm2$x_my%vwkhU=W1HPpzbh!##LJNS&?GIPw2kh;L_EQlj8plxxSR zrU_@}a>=ANBh96>d@hJlC6%iGhAs2fbm1#1*Mgf&=)@9-Jk5w$bc!D)xt9lHYcF2+Gq8&**wel+6ah&V=V z1Z%cho0Qr_5?m(<38S5`clp6>IeDqkM6{<=-uHe*nf86(0Zq