public/bitbot
1
1
Fork 0
Code Issues 2 Pull Requests Packages Projects Releases Wiki Activity

use discord role system to role check

Browse Source
This commit is contained in:
Kristian 2023-10-06 21:57:30 +02:00
parent a983bcac7e
commit 34f3bcac3d
1 changed files with 16 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

23
bot/bot.go
View File

@ -44,6 +44,15 @@ func Run() {
var conversationHistoryMap = make(map[string][]openai.ChatCompletionMessage) var conversationHistoryMap = make(map[string][]openai.ChatCompletionMessage)
var sshConnections = make(map[string]*SSHConnection) var sshConnections = make(map[string]*SSHConnection)
func hasAdminRole(roles []string) bool {
for _, role := range roles {
if role == AllowedUserID {
return true
}
}
return false
}
func newMessage(discord *discordgo.Session, message *discordgo.MessageCreate) { func newMessage(discord *discordgo.Session, message *discordgo.MessageCreate) {
if message.Author.ID == discord.State.User.ID || message.Content == "" { if message.Author.ID == discord.State.User.ID || message.Content == "" {
return return
@ -70,7 +79,7 @@ func newMessage(discord *discordgo.Session, message *discordgo.MessageCreate) {
} else if strings.HasPrefix(message.Content, "!bit") || isPrivateChannel { } else if strings.HasPrefix(message.Content, "!bit") || isPrivateChannel {
chatGPT(discord, message.ChannelID, message.Content, conversationHistory) chatGPT(discord, message.ChannelID, message.Content, conversationHistory)
} else if strings.HasPrefix(message.Content, "!genkey") { } else if strings.HasPrefix(message.Content, "!genkey") {
if message.Author.ID == AllowedUserID { if hasAdminRole(message.Member.Roles) {
err := GenerateAndSaveSSHKeyPairIfNotExist() err := GenerateAndSaveSSHKeyPairIfNotExist()
if err != nil { if err != nil {
discord.ChannelMessageSend(message.ChannelID, "Error generating or saving key pair.") discord.ChannelMessageSend(message.ChannelID, "Error generating or saving key pair.")
@ -81,7 +90,7 @@ func newMessage(discord *discordgo.Session, message *discordgo.MessageCreate) {
discord.ChannelMessageSend(message.ChannelID, "You are not authorized to use this command.") discord.ChannelMessageSend(message.ChannelID, "You are not authorized to use this command.")
} }
} else if strings.HasPrefix(message.Content, "!showkey") { } else if strings.HasPrefix(message.Content, "!showkey") {
if message.Author.ID == AllowedUserID { if hasAdminRole(message.Member.Roles) {
publicKey, err := GetPublicKey() publicKey, err := GetPublicKey()
if err != nil { if err != nil {
discord.ChannelMessageSend(message.ChannelID, "Error fetching public key.") discord.ChannelMessageSend(message.ChannelID, "Error fetching public key.")
@ -92,7 +101,7 @@ func newMessage(discord *discordgo.Session, message *discordgo.MessageCreate) {
discord.ChannelMessageSend(message.ChannelID, "You are not authorized to use this command.") discord.ChannelMessageSend(message.ChannelID, "You are not authorized to use this command.")
} }
} else if strings.HasPrefix(message.Content, "!regenkey") { } else if strings.HasPrefix(message.Content, "!regenkey") {
if message.Author.ID == AllowedUserID { if hasAdminRole(message.Member.Roles) {
err := GenerateAndSaveSSHKeyPair() err := GenerateAndSaveSSHKeyPair()
if err != nil { if err != nil {
discord.ChannelMessageSend(message.ChannelID, "Error regenerating and saving key pair.") discord.ChannelMessageSend(message.ChannelID, "Error regenerating and saving key pair.")
@ -103,7 +112,7 @@ func newMessage(discord *discordgo.Session, message *discordgo.MessageCreate) {
discord.ChannelMessageSend(message.ChannelID, "You are not authorized to use this command.") discord.ChannelMessageSend(message.ChannelID, "You are not authorized to use this command.")
} }
} else if strings.HasPrefix(message.Content, "!ssh") { } else if strings.HasPrefix(message.Content, "!ssh") {
if message.Author.ID == AllowedUserID { if hasAdminRole(message.Member.Roles) {
commandParts := strings.Fields(message.Content) commandParts := strings.Fields(message.Content)
if len(commandParts) != 2 { if len(commandParts) != 2 {
discord.ChannelMessageSend(message.ChannelID, "Invalid command format. Use !ssh username@remote-host:port") discord.ChannelMessageSend(message.ChannelID, "Invalid command format. Use !ssh username@remote-host:port")
@ -125,7 +134,7 @@ func newMessage(discord *discordgo.Session, message *discordgo.MessageCreate) {
discord.ChannelMessageSend(message.ChannelID, "You are not authorized to use this command.") discord.ChannelMessageSend(message.ChannelID, "You are not authorized to use this command.")
} }
} else if strings.HasPrefix(message.Content, "!exe") { } else if strings.HasPrefix(message.Content, "!exe") {
if message.Author.ID == AllowedUserID { if hasAdminRole(message.Member.Roles) {
// Check if there is an active SSH connection for this user // Check if there is an active SSH connection for this user
sshConn, ok := sshConnections[message.Author.ID] sshConn, ok := sshConnections[message.Author.ID]
if !ok { if !ok {
@ -148,7 +157,7 @@ func newMessage(discord *discordgo.Session, message *discordgo.MessageCreate) {
discord.ChannelMessageSend(message.ChannelID, "You are not authorized to use this command.") discord.ChannelMessageSend(message.ChannelID, "You are not authorized to use this command.")
} }
} else if strings.HasPrefix(message.Content, "!exit") { } else if strings.HasPrefix(message.Content, "!exit") {
if message.Author.ID == AllowedUserID { if hasAdminRole(message.Member.Roles) {
// Check if there is an active SSH connection for this user // Check if there is an active SSH connection for this user
sshConn, ok := sshConnections[message.Author.ID] sshConn, ok := sshConnections[message.Author.ID]
if !ok { if !ok {
@ -174,7 +183,7 @@ func newMessage(discord *discordgo.Session, message *discordgo.MessageCreate) {
"!regenkey - Regenerate and save SSH key pair.\n" + "!regenkey - Regenerate and save SSH key pair.\n" +
"!ssh username@remote-host:port - Connect to a remote server via SSH.\n" + "!ssh username@remote-host:port - Connect to a remote server via SSH.\n" +
"!exe command - Execute a command on the remote server (after !ssh).\n" + "!exe command - Execute a command on the remote server (after !ssh).\n" +
"!sshexit - Close the SSH connection (after !ssh).\n" "!exit - Close the SSH connection (after !ssh).\n"
discord.ChannelMessageSend(message.ChannelID, adminHelpMessage) discord.ChannelMessageSend(message.ChannelID, adminHelpMessage)
} else { } else {