diff --git a/commands/discommand.go b/commands/discommand.go index fe30ed4..b550888 100644 --- a/commands/discommand.go +++ b/commands/discommand.go @@ -4,6 +4,7 @@ import ( "sync" "git.wh64.net/muffin/goMuffin/configs" + "git.wh64.net/muffin/goMuffin/databases" "git.wh64.net/muffin/goMuffin/utils" "github.com/bwmarrin/discordgo" ) @@ -17,6 +18,7 @@ type modalParse func(ctx *ModalContext) bool type componentParse func(ctx *ComponentContext) bool type Category string +type CommandFlags uint8 type DetailedDescription struct { Usage string @@ -30,6 +32,7 @@ type Command struct { Category Category RegisterApplicationCommand bool RegisterMessageCommand bool + Flags CommandFlags MessageRun messageRun ChatInputRun chatInputRun } @@ -78,6 +81,11 @@ const ( DeveloperOnly Category = "개발자 전용" ) +const ( + CommandFlagsIsDeveloper CommandFlags = 1 << iota + CommandFlagsIsRegistered +) + var ( commandMutex sync.Mutex componentMutex sync.Mutex @@ -118,10 +126,15 @@ func (d *DiscommandStruct) LoadModal(m *Modal) { d.Modals = append(d.Modals, m) } -func (d *DiscommandStruct) MessageRun(name string, s *discordgo.Session, m *discordgo.MessageCreate, args []string) { - if command, ok := d.Commands[name]; ok { - if command.Category == DeveloperOnly && m.Author.ID != configs.Config.Bot.OwnerId { - utils.NewMessageSender(&utils.MessageCreate{MessageCreate: m, Session: s}). +func (d *DiscommandStruct) MessageRun(name string, s *discordgo.Session, msg *discordgo.MessageCreate, args []string) { + m := &utils.MessageCreate{ + MessageCreate: msg, + Session: s, + } + + if command, ok := d.Commands[name]; ok && command.RegisterMessageCommand { + if command.Flags&CommandFlagsIsDeveloper != 0 && m.Author.ID != configs.Config.Bot.OwnerId { + utils.NewMessageSender(m). AddComponents(utils.GetErrorContainer(discordgo.TextDisplay{Content: "해당 명령어는 개발자만 사용 가능해요."})). SetComponentsV2(true). SetReply(true). @@ -129,24 +142,48 @@ func (d *DiscommandStruct) MessageRun(name string, s *discordgo.Session, m *disc return } - if !command.RegisterMessageCommand { + if command.Flags&CommandFlagsIsRegistered != 0 && !databases.Database.IsUser(m.Author.ID) { + utils.NewMessageSender(m). + AddComponents(utils.GetUserIsNotRegisteredErrContainer(configs.Config.Bot.Prefix)). + SetComponentsV2(true). + SetReply(true). + Send() return } - command.MessageRun(&MsgContext{&utils.MessageCreate{ - MessageCreate: m, - Session: s, - }, &args, command}) + command.MessageRun(&MsgContext{m, &args, command}) } } -func (d *DiscommandStruct) ChatInputRun(name string, s *discordgo.Session, i *discordgo.InteractionCreate) { +func (d *DiscommandStruct) ChatInputRun(name string, s *discordgo.Session, inter *discordgo.InteractionCreate) { + i := &utils.InteractionCreate{ + InteractionCreate: inter, + Session: s, + Options: utils.GetInteractionOptions(inter), + } + + i.InteractionCreate.User = utils.GetInteractionUser(inter) + if command, ok := d.Commands[name]; ok && command.RegisterApplicationCommand { - command.ChatInputRun(&ChatInputContext{&utils.InteractionCreate{ - InteractionCreate: i, - Session: s, - Options: utils.GetInteractionOptions(i), - }, command}) + if command.Flags&CommandFlagsIsDeveloper != 0 && i.User.ID != configs.Config.Bot.OwnerId { + utils.NewMessageSender(i). + AddComponents(utils.GetErrorContainer(discordgo.TextDisplay{Content: "해당 명령어는 개발자만 사용 가능해요."})). + SetComponentsV2(true). + SetReply(true). + Send() + return + } + + if command.Flags&CommandFlagsIsRegistered != 0 && !databases.Database.IsUser(i.User.ID) { + utils.NewMessageSender(i). + AddComponents(utils.GetUserIsNotRegisteredErrContainer(configs.Config.Bot.Prefix)). + SetComponentsV2(true). + SetReply(true). + Send() + return + } + + command.ChatInputRun(&ChatInputContext{i, command}) } } diff --git a/commands/reloadPrompt.go b/commands/reloadPrompt.go index 3c4a033..0681b95 100644 --- a/commands/reloadPrompt.go +++ b/commands/reloadPrompt.go @@ -21,6 +21,7 @@ var ReloadPromptCommand *Command = &Command{ Category: DeveloperOnly, RegisterApplicationCommand: false, RegisterMessageCommand: true, + Flags: CommandFlagsIsDeveloper, MessageRun: func(ctx *MsgContext) { err := chatbot.ChatBot.ReloadPrompt() if err != nil { diff --git a/commands/switchMode.go b/commands/switchMode.go index 673668e..4888b56 100644 --- a/commands/switchMode.go +++ b/commands/switchMode.go @@ -20,6 +20,7 @@ var SwitchModeCommand *Command = &Command{ Category: DeveloperOnly, RegisterApplicationCommand: false, RegisterMessageCommand: true, + Flags: CommandFlagsIsDeveloper, MessageRun: func(ctx *MsgContext) { chatbot.ChatBot.SwitchMode() diff --git a/databases/User.go b/databases/User.go new file mode 100644 index 0000000..101dee9 --- /dev/null +++ b/databases/User.go @@ -0,0 +1,29 @@ +package databases + +import ( + "context" + "time" + + "go.mongodb.org/mongo-driver/v2/bson" +) + +type InsertUser struct { + UserId string `bson:"user_id"` + CreatedAt time.Time `bson:"created_at"` +} + +type User struct { + Id bson.ObjectID `bson:"_id"` + UserId string `bson:"user_id"` + CreatedAt time.Time `bson:"created_at"` +} + +func (d *MuffinDatabase) IsUser(userId string) bool { + var user *User + d.Users.FindOne(context.TODO(), bson.D{{Key: "user_id", Value: userId}}).Decode(&user) + + if user != nil { + return true + } + return false +} diff --git a/databases/database.go b/databases/database.go index 75a9cf8..fe7da03 100644 --- a/databases/database.go +++ b/databases/database.go @@ -13,6 +13,7 @@ type MuffinDatabase struct { Learns *mongo.Collection Texts *mongo.Collection Memory *mongo.Collection + Users *mongo.Collection } var Database *MuffinDatabase @@ -36,5 +37,6 @@ func Connect() (*MuffinDatabase, error) { Learns: client.Database(configs.Config.Database.Name).Collection("learn"), Texts: client.Database(configs.Config.Database.Name).Collection("text"), Memory: client.Database(configs.Config.Database.Name).Collection("memory"), + Users: client.Database(configs.Config.Database.Name).Collection("user"), }, nil } diff --git a/utils/embed.go b/utils/embed.go index 4093445..861d260 100644 --- a/utils/embed.go +++ b/utils/embed.go @@ -1,6 +1,10 @@ package utils -import "github.com/bwmarrin/discordgo" +import ( + "fmt" + + "github.com/bwmarrin/discordgo" +) const ( EmbedDefault int = 0xaddb87 @@ -33,3 +37,9 @@ func GetSuccessContainer(components ...discordgo.MessageComponent) *discordgo.Co c.Components = append(c.Components, components...) return c } + +func GetUserIsNotRegisteredErrContainer(prefix string) *discordgo.Container { + return GetErrorContainer(discordgo.TextDisplay{ + Content: fmt.Sprintf("해당 기능은 등록된 사용자만 쓸 수 있어요. `%s가입`으로 가입해주새요.", prefix), + }) +} diff --git a/utils/interactions.go b/utils/interactions.go index 0cc4461..3f0dda7 100644 --- a/utils/interactions.go +++ b/utils/interactions.go @@ -53,6 +53,17 @@ func GetInteractionOptions(i *discordgo.InteractionCreate) map[string]*discordgo return optsMap } +func GetInteractionUser(i *discordgo.InteractionCreate) *discordgo.User { + if i.Member != nil { + return i.Member.User + } + + if i.User != nil { + return i.User + } + return nil +} + // DeferReply to this interaction. func (i *InteractionCreate) DeferReply(data *discordgo.InteractionResponseData) error { err := i.Session.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{ diff --git a/utils/messageBuilder.go b/utils/messageBuilder.go index 612a1ce..bd92070 100644 --- a/utils/messageBuilder.go +++ b/utils/messageBuilder.go @@ -63,7 +63,7 @@ func (s *MessageSender) Send() error { var flags discordgo.MessageFlags if s.ComponentsV2 { - flags = flags | discordgo.MessageFlagsIsComponentsV2 + flags |= discordgo.MessageFlagsIsComponentsV2 } switch m := s.m.(type) { @@ -85,7 +85,7 @@ func (s *MessageSender) Send() error { return err case *InteractionCreate: if s.Ephemeral { - flags = flags | discordgo.MessageFlagsEphemeral + flags |= discordgo.MessageFlagsEphemeral } if m.Replied || m.Deferred {