bbgo_origin/pkg/interact/interact.go

267 lines
6.2 KiB
Go
Raw Normal View History

2022-01-12 13:45:12 +00:00
package interact
2022-01-11 18:54:13 +00:00
import (
2022-01-13 14:15:05 +00:00
"context"
"fmt"
"time"
2022-01-11 18:54:13 +00:00
2022-01-13 14:15:05 +00:00
log "github.com/sirupsen/logrus"
2022-01-11 18:54:13 +00:00
)
2022-01-13 14:15:05 +00:00
type CustomInteraction interface {
Commands(interact *Interact)
}
2022-01-14 18:52:46 +00:00
type Initializer interface {
Initialize() error
}
2022-01-13 14:24:51 +00:00
type Messenger interface {
TextMessageResponder
CommandResponder
2022-01-13 17:58:04 +00:00
Start(ctx context.Context)
}
type Session interface {
ID() string
SetOriginState(state State)
GetOriginState() State
SetState(state State)
GetState() State
IsAuthorized() bool
SetAuthorized()
SetAuthorizing(b bool)
}
2022-01-11 18:54:13 +00:00
// Interact implements the interaction between bot and message software.
type Interact struct {
startTime time.Time
2022-01-13 15:41:22 +00:00
// commands is the default public command map
commands map[string]*Command
2022-01-13 15:41:22 +00:00
// privateCommands is the private command map, need auth
privateCommands map[string]*Command
2022-01-13 14:15:05 +00:00
states map[State]State
statesFunc map[State]interface{}
authenticatedSessions map[string]Session
2022-01-13 14:15:05 +00:00
2022-01-14 18:52:46 +00:00
customInteractions []CustomInteraction
2022-01-16 11:06:26 +00:00
messengers []Messenger
}
func New() *Interact {
return &Interact{
2022-01-14 18:52:46 +00:00
startTime: time.Now(),
commands: make(map[string]*Command),
privateCommands: make(map[string]*Command),
states: make(map[State]State),
statesFunc: make(map[State]interface{}),
}
}
func (it *Interact) AddCustomInteraction(custom CustomInteraction) {
custom.Commands(it)
2022-01-14 18:52:46 +00:00
it.customInteractions = append(it.customInteractions, custom)
2022-01-13 14:15:05 +00:00
}
2022-01-14 18:52:46 +00:00
func (it *Interact) PrivateCommand(command, desc string, f interface{}) *Command {
cmd := NewCommand(command, desc, f)
it.privateCommands[command] = cmd
2022-01-13 15:41:22 +00:00
return cmd
}
2022-01-13 18:57:39 +00:00
func (it *Interact) Command(command string, desc string, f interface{}) *Command {
cmd := NewCommand(command, desc, f)
it.commands[command] = cmd
return cmd
}
func (it *Interact) getNextState(session Session, currentState State) (nextState State, final bool) {
var ok bool
2022-01-13 14:15:05 +00:00
final = false
nextState, ok = it.states[currentState]
if ok {
2022-01-13 14:15:05 +00:00
// check if it's the final state
if _, hasTransition := it.statesFunc[nextState]; !hasTransition {
2022-01-13 14:15:05 +00:00
final = true
}
return nextState, final
}
2022-01-13 14:15:05 +00:00
// state not found, return to the origin state
return session.GetOriginState(), final
2022-01-11 18:54:13 +00:00
}
func (it *Interact) handleResponse(session Session, text string, ctxObjects ...interface{}) error {
// We only need response when executing a command
switch session.GetState() {
case StatePublic, StateAuthenticated:
return nil
}
args := parseCommand(text)
state := session.GetState()
f, ok := it.statesFunc[state]
if !ok {
return fmt.Errorf("state function of %s is not defined", state)
}
ctxObjects = append(ctxObjects, session)
_, err := parseFuncArgsAndCall(f, args, ctxObjects...)
if err != nil {
return err
}
nextState, end := it.getNextState(session, state)
if end {
session.SetState(session.GetOriginState())
return nil
}
session.SetState(nextState)
return nil
}
func (it *Interact) getCommand(session Session, command string) (*Command, error) {
if session.IsAuthorized() {
if cmd, ok := it.privateCommands[command]; ok {
2022-01-13 16:17:41 +00:00
return cmd, nil
}
} else {
if _, ok := it.privateCommands[command]; ok {
return nil, fmt.Errorf("private command can not be executed in the public mode, type /auth to get authorized")
2022-01-13 16:17:41 +00:00
}
}
// find any public command
if cmd, ok := it.commands[command]; ok {
2022-01-13 16:17:41 +00:00
return cmd, nil
}
return nil, fmt.Errorf("command %s not found", command)
}
func (it *Interact) runCommand(session Session, command string, args []string, ctxObjects ...interface{}) error {
cmd, err := it.getCommand(session, command)
2022-01-13 16:17:41 +00:00
if err != nil {
return err
}
ctxObjects = append(ctxObjects, session)
session.SetState(cmd.initState)
if _, err := parseFuncArgsAndCall(cmd.F, args, ctxObjects...); err != nil {
return err
}
// if we can successfully execute the command, then we can go to the next state.
state := session.GetState()
nextState, end := it.getNextState(session, state)
if end {
session.SetState(session.GetOriginState())
return nil
}
session.SetState(nextState)
return nil
2022-01-12 13:45:12 +00:00
}
2022-01-11 18:54:13 +00:00
2022-01-16 11:06:26 +00:00
func (it *Interact) AddMessenger(messenger Messenger) {
2022-01-13 17:58:04 +00:00
// pass Responder function
messenger.SetTextMessageResponder(func(session Session, message string, reply Reply, ctxObjects ...interface{}) error {
return it.handleResponse(session, message, append(ctxObjects, reply)...)
2022-01-13 14:24:51 +00:00
})
2022-01-16 11:06:26 +00:00
it.messengers = append(it.messengers, messenger)
2022-01-13 14:15:05 +00:00
}
2022-01-13 16:17:41 +00:00
// builtin initializes the built-in commands
func (it *Interact) builtin() error {
2022-01-13 18:57:39 +00:00
it.Command("/uptime", "show bot uptime", func(reply Reply) error {
uptime := time.Since(it.startTime)
reply.Message(fmt.Sprintf("uptime %s", uptime))
2022-01-13 16:17:41 +00:00
return nil
})
return nil
}
func (it *Interact) init() error {
if err := it.builtin(); err != nil {
2022-01-13 16:17:41 +00:00
return err
}
2022-01-14 18:52:46 +00:00
if err := it.registerCommands(it.commands); err != nil {
return err
}
if err := it.registerCommands(it.privateCommands); err != nil {
return err
}
return nil
}
func (it *Interact) registerCommands(commands map[string]*Command) error {
for n, cmd := range commands {
for s1, s2 := range cmd.states {
if _, exist := it.states[s1]; exist {
return fmt.Errorf("state %s already exists", s1)
}
it.states[s1] = s2
}
for s, f := range cmd.statesFunc {
it.statesFunc[s] = f
}
2022-01-13 14:15:05 +00:00
// register commands to the service
2022-01-16 11:06:26 +00:00
if len(it.messengers) == 0 {
2022-01-13 14:15:05 +00:00
return fmt.Errorf("messenger is not set")
}
2022-01-16 11:06:26 +00:00
// commandName is used in the closure, we need to copy the variable
2022-01-13 16:17:41 +00:00
commandName := n
2022-01-16 11:06:26 +00:00
for _, messenger := range it.messengers {
messenger.AddCommand(cmd, func(session Session, message string, reply Reply, ctxObjects ...interface{}) error {
args := parseCommand(message)
return it.runCommand(session, commandName, args, append(ctxObjects, reply)...)
})
}
}
return nil
2022-01-12 13:45:12 +00:00
}
func (it *Interact) Start(ctx context.Context) error {
2022-01-16 11:06:26 +00:00
if len(it.messengers) == 0 {
log.Warn("messenger is not set, skip initializing")
return nil
}
if err := it.init(); err != nil {
2022-01-13 14:15:05 +00:00
return err
}
2022-01-14 18:52:46 +00:00
for _, custom := range it.customInteractions {
log.Infof("checking %T custom interaction...", custom)
if initializer, ok := custom.(Initializer); ok {
log.Infof("initializing %T custom interaction...", custom)
if err := initializer.Initialize(); err != nil {
return err
}
}
}
2022-01-13 14:15:05 +00:00
// TODO: use go routine and context
2022-01-16 11:06:26 +00:00
for _, m := range it.messengers {
m.Start(ctx)
}
2022-01-13 14:15:05 +00:00
return nil
2022-01-11 18:54:13 +00:00
}