-
Notifications
You must be signed in to change notification settings - Fork 911
/
system.go
353 lines (283 loc) · 9.06 KB
/
system.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
package dcmd
import (
"fmt"
"log"
"runtime/debug"
"strings"
"github.com/botlabs-gg/yagpdb/v2/lib/discordgo"
"github.com/botlabs-gg/yagpdb/v2/lib/dstate"
"github.com/pkg/errors"
)
type System struct {
Root *Container
Prefix PrefixProvider
ResponseSender ResponseSender
State dstate.StateTracker
}
func NewStandardSystem(staticPrefix string) (system *System) {
sys := &System{
Root: &Container{HelpTitleEmoji: "ℹ️", HelpColor: 0xbeff7a},
ResponseSender: &StdResponseSender{LogErrors: true},
}
if staticPrefix != "" {
sys.Prefix = NewSimplePrefixProvider(staticPrefix)
}
sys.Root.AddMidlewares(ArgParserMW)
return sys
}
// You can add this as a handler directly to discordgo, it will recover from any panics that occured in commands
// and log errors using the standard logger
func (sys *System) HandleMessageCreate(s *discordgo.Session, m *discordgo.MessageCreate) {
// Set up handler to recover from panics
defer func() {
if r := recover(); r != nil {
sys.handlePanic(s, r, false)
}
}()
err := sys.CheckMessage(s, m)
if err != nil {
log.Println("[DCMD ERROR]: Failed checking message:", err)
}
}
// CheckMessage checks the message for commands, and triggers any command that the message should trigger
// you should not add this as an discord handler directly, if you want to do that you should add "system.HandleMessageCreate" instead.
func (sys *System) CheckMessage(s *discordgo.Session, m *discordgo.MessageCreate) error {
data, err := sys.FillDataLegacyMessage(s, m.Message)
if err != nil {
return err
}
if !sys.FindPrefix(data) {
// No prefix found in the message for a command to be triggered
return nil
}
response, err := sys.Root.Run(data)
return sys.ResponseSender.SendResponse(data, response, err)
}
// CheckInteraction checks a interaction and runs a command if found
func (sys *System) CheckInteraction(s *discordgo.Session, interaction *discordgo.Interaction) error {
data, err := sys.FillDataInteraction(s, interaction)
if err != nil {
return err
}
response, err := sys.Root.Run(data)
return sys.ResponseSender.SendResponse(data, response, err)
}
// CheckMessageWtihPrefetchedPrefix is the same as CheckMessage but you pass in a prefetched command prefix
func (sys *System) CheckMessageWtihPrefetchedPrefix(s *discordgo.Session, m *discordgo.MessageCreate, prefetchedPrefix string) error {
data, err := sys.FillDataLegacyMessage(s, m.Message)
if err != nil {
return err
}
if !sys.FindPrefixWithPrefetched(data, prefetchedPrefix) {
// No prefix found in the message for a command to be triggered
return nil
}
response, err := sys.Root.Run(data)
return sys.ResponseSender.SendResponse(data, response, err)
}
// FindPrefix checks if the message has a proper command prefix (either from the PrefixProvider or a direction mention to the bot)
// It sets the source field, and MsgStripped in data if found
func (sys *System) FindPrefix(data *Data) (found bool) {
if data.Source == TriggerSourceDM {
data.TraditionalTriggerData.MessageStrippedPrefix = data.TraditionalTriggerData.Message.Content
data.TriggerType = TriggerTypeDirect
return true
}
if sys.FindMentionPrefix(data) {
return true
}
// Check for custom prefix
if sys.Prefix == nil {
return false
}
prefix := sys.Prefix.Prefix(data)
if prefix == "" {
return false
}
data.TraditionalTriggerData.PrefixUsed = prefix
if strings.HasPrefix(data.TraditionalTriggerData.Message.Content, prefix) {
data.TriggerType = TriggerTypePrefix
data.TraditionalTriggerData.MessageStrippedPrefix = strings.TrimSpace(strings.Replace(data.TraditionalTriggerData.Message.Content, prefix, "", 1))
found = true
}
return
}
// FindPrefixWithPrefetched is the same as FindPrefix but you pass in a prefetched command prefix
func (sys *System) FindPrefixWithPrefetched(data *Data, commandPrefix string) (found bool) {
msg := data.TraditionalTriggerData.Message
if data.Source == TriggerSourceDM {
data.TraditionalTriggerData.MessageStrippedPrefix = msg.Content
data.TriggerType = TriggerTypeDirect
return true
}
if sys.FindMentionPrefix(data) {
return true
}
data.TraditionalTriggerData.PrefixUsed = commandPrefix
if strings.HasPrefix(msg.Content, commandPrefix) {
data.TriggerType = TriggerTypePrefix
data.TraditionalTriggerData.MessageStrippedPrefix = strings.TrimSpace(strings.Replace(msg.Content, commandPrefix, "", 1))
found = true
}
return
}
func (sys *System) FindMentionPrefix(data *Data) (found bool) {
if data.Session.State.User == nil {
return false
}
ok := false
stripped := ""
msg := data.TraditionalTriggerData.Message
// Check for mention
id := discordgo.StrID(data.Session.State.User.ID)
if strings.Index(msg.Content, "<@"+id+">") == 0 { // Normal mention
ok = true
stripped = strings.Replace(msg.Content, "<@"+id+">", "", 1)
data.TraditionalTriggerData.PrefixUsed = "<@" + id + ">"
} else if strings.Index(msg.Content, "<@!"+id+">") == 0 { // Nickname mention
ok = true
data.TraditionalTriggerData.PrefixUsed = "<@!" + id + ">"
stripped = strings.Replace(msg.Content, "<@!"+id+">", "", 1)
}
if ok {
data.TraditionalTriggerData.MessageStrippedPrefix = strings.TrimSpace(stripped)
data.TriggerType = TriggerTypeMention
return true
}
return false
}
var (
ErrChannelNotFound = errors.New("Channel not found")
ErrGuildNotFound = errors.New("Guild not found")
ErrMemberNotAvailable = errors.New("Member not provided in message")
ErrMemberNotAvailableInteraction = errors.New("Member not provided in interaction")
)
func (sys *System) FillDataLegacyMessage(s *discordgo.Session, m *discordgo.Message) (*Data, error) {
var gs *dstate.GuildSet
var cs *dstate.ChannelState
if m.GuildID != 0 {
gs = sys.State.GetGuild(m.GuildID)
if gs == nil {
return nil, ErrGuildNotFound
}
cs = gs.GetChannelOrThread(m.ChannelID)
if cs == nil {
return nil, ErrChannelNotFound
}
}
data := &Data{
Session: s,
System: sys,
ChannelID: m.ChannelID,
Author: m.Author,
TraditionalTriggerData: &TraditionalTriggerData{
Message: m,
},
}
if m.GuildID == 0 {
data.Source = TriggerSourceDM
} else {
data.Source = TriggerSourceGuild
if m.Member == nil || m.Author == nil {
return nil, ErrMemberNotAvailable
}
member := *m.Member
member.User = m.Author // user field is not provided in Message.Member, its weird but *shrug*
member.GuildID = m.GuildID
data.GuildData = &GuildContextData{
CS: cs,
GS: gs,
MS: dstate.MemberStateFromMember(&member),
}
}
return data, nil
}
func (sys *System) FillDataInteraction(s *discordgo.Session, interaction *discordgo.Interaction) (*Data, error) {
var gs *dstate.GuildSet
var cs *dstate.ChannelState
if interaction.GuildID != 0 {
gs = sys.State.GetGuild(interaction.GuildID)
if gs == nil {
return nil, ErrGuildNotFound
}
cs = gs.GetChannelOrThread(interaction.ChannelID)
if cs == nil {
return nil, ErrChannelNotFound
}
}
user := interaction.User
if interaction.Member != nil {
user = interaction.Member.User
}
data := &Data{
Session: s,
System: sys,
ChannelID: interaction.ChannelID,
Author: user,
TriggerType: TriggerTypeSlashCommands,
SlashCommandTriggerData: &SlashCommandTriggerData{
Interaction: interaction,
},
}
if interaction.GuildID == 0 {
data.Source = TriggerSourceDM
} else {
data.Source = TriggerSourceGuild
// were working off the assumption that member is always provided when in a guild
if interaction.Member == nil {
return nil, ErrMemberNotAvailableInteraction
}
member := *interaction.Member
member.GuildID = interaction.GuildID
data.GuildData = &GuildContextData{
CS: cs,
GS: gs,
MS: dstate.MemberStateFromMember(&member),
}
}
return data, nil
}
func (sys *System) handlePanic(s *discordgo.Session, r interface{}, sendChatNotice bool) {
// TODO
stack := debug.Stack()
log.Printf("[DCMD PANIC] %v\n%s", r, string(stack))
}
// Retrieves the prefix that might be different on a per server basis
type PrefixProvider interface {
Prefix(data *Data) string
}
// Simple Prefix provider for global fixed prefixes
type SimplePrefixProvider struct {
prefix string
}
func NewSimplePrefixProvider(prefix string) PrefixProvider {
return &SimplePrefixProvider{prefix: prefix}
}
func (pp *SimplePrefixProvider) Prefix(d *Data) string {
return pp.prefix
}
func Indent(depth int) string {
indent := ""
for i := 0; i < depth; i++ {
indent += "-"
}
return indent
}
type ResponseSender interface {
SendResponse(cmdData *Data, resp interface{}, err error) error
}
type StdResponseSender struct {
LogErrors bool
}
func (s *StdResponseSender) SendResponse(cmdData *Data, resp interface{}, err error) error {
if err != nil && s.LogErrors {
log.Printf("[DCMD]: Command %q returned an error: %s", cmdData.Cmd.FormatNames(false, "/"), err)
}
var errR error
if resp == nil && err != nil {
_, errR = SendResponseInterface(cmdData, fmt.Sprintf("%q command returned an error: %s", cmdData.Cmd.FormatNames(false, "/"), err), true)
} else if resp != nil {
_, errR = SendResponseInterface(cmdData, resp, false)
}
return errR
}