mirror of
https://github.com/c9s/bbgo.git
synced 2024-11-13 02:23:51 +00:00
interact: improve authentication process
This commit is contained in:
parent
62e5706657
commit
17322cbc09
|
@ -131,6 +131,7 @@ func main() {
|
||||||
|
|
||||||
globalInteraction := interact.New()
|
globalInteraction := interact.New()
|
||||||
globalInteraction.SetMessenger(&interact.Telegram{
|
globalInteraction.SetMessenger(&interact.Telegram{
|
||||||
|
Private: true,
|
||||||
Bot: b,
|
Bot: b,
|
||||||
})
|
})
|
||||||
|
|
||||||
|
|
|
@ -20,6 +20,7 @@ const (
|
||||||
var ErrAuthenticationFailed = errors.New("authentication failed")
|
var ErrAuthenticationFailed = errors.New("authentication failed")
|
||||||
|
|
||||||
type Authorizer interface {
|
type Authorizer interface {
|
||||||
|
StartAuthorizing()
|
||||||
Authorize() error
|
Authorize() error
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@ -50,8 +51,9 @@ func (it *AuthInteract) Commands(interact *Interact) {
|
||||||
|
|
||||||
it.OneTimePasswordKey = key
|
it.OneTimePasswordKey = key
|
||||||
}
|
}
|
||||||
interact.Command("/auth", func(reply Reply) error {
|
interact.Command("/auth", func(reply Reply, authorizer Authorizer) error {
|
||||||
reply.Message("Enter your authentication token")
|
reply.Message("Enter your authentication token")
|
||||||
|
authorizer.StartAuthorizing()
|
||||||
return nil
|
return nil
|
||||||
}).Next(func(token string, reply Reply) error {
|
}).Next(func(token string, reply Reply) error {
|
||||||
if token == it.Token {
|
if token == it.Token {
|
||||||
|
@ -72,6 +74,7 @@ func (it *AuthInteract) Commands(interact *Interact) {
|
||||||
}).NamedNext(StateAuthenticated, func(code string, reply Reply, authorizer Authorizer) error {
|
}).NamedNext(StateAuthenticated, func(code string, reply Reply, authorizer Authorizer) error {
|
||||||
if totp.Validate(code, it.OneTimePasswordKey.Secret()) {
|
if totp.Validate(code, it.OneTimePasswordKey.Secret()) {
|
||||||
reply.Message("Great! You're authenticated!")
|
reply.Message("Great! You're authenticated!")
|
||||||
|
interact.SetOriginState(StateAuthenticated)
|
||||||
return authorizer.Authorize()
|
return authorizer.Authorize()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@ -87,12 +90,14 @@ func (it *AuthInteract) Commands(interact *Interact) {
|
||||||
case AuthModeToken:
|
case AuthModeToken:
|
||||||
if code == it.Token {
|
if code == it.Token {
|
||||||
reply.Message("Great! You're authenticated!")
|
reply.Message("Great! You're authenticated!")
|
||||||
|
interact.SetOriginState(StateAuthenticated)
|
||||||
return authorizer.Authorize()
|
return authorizer.Authorize()
|
||||||
}
|
}
|
||||||
|
|
||||||
case AuthModeOTP:
|
case AuthModeOTP:
|
||||||
if totp.Validate(code, it.OneTimePasswordKey.Secret()) {
|
if totp.Validate(code, it.OneTimePasswordKey.Secret()) {
|
||||||
reply.Message("Great! You're authenticated!")
|
reply.Message("Great! You're authenticated!")
|
||||||
|
interact.SetOriginState(StateAuthenticated)
|
||||||
return authorizer.Authorize()
|
return authorizer.Authorize()
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
|
@ -7,6 +7,7 @@ import (
|
||||||
"strconv"
|
"strconv"
|
||||||
"strings"
|
"strings"
|
||||||
"text/scanner"
|
"text/scanner"
|
||||||
|
"time"
|
||||||
|
|
||||||
log "github.com/sirupsen/logrus"
|
log "github.com/sirupsen/logrus"
|
||||||
)
|
)
|
||||||
|
@ -47,6 +48,8 @@ type Messenger interface {
|
||||||
|
|
||||||
// Interact implements the interaction between bot and message software.
|
// Interact implements the interaction between bot and message software.
|
||||||
type Interact struct {
|
type Interact struct {
|
||||||
|
startTime time.Time
|
||||||
|
|
||||||
// commands is the default public command map
|
// commands is the default public command map
|
||||||
commands map[string]*Command
|
commands map[string]*Command
|
||||||
|
|
||||||
|
@ -63,6 +66,7 @@ type Interact struct {
|
||||||
|
|
||||||
func New() *Interact {
|
func New() *Interact {
|
||||||
return &Interact{
|
return &Interact{
|
||||||
|
startTime: time.Now(),
|
||||||
commands: make(map[string]*Command),
|
commands: make(map[string]*Command),
|
||||||
originState: StatePublic,
|
originState: StatePublic,
|
||||||
currentState: StatePublic,
|
currentState: StatePublic,
|
||||||
|
@ -114,6 +118,13 @@ func (it *Interact) setState(s State) {
|
||||||
}
|
}
|
||||||
|
|
||||||
func (it *Interact) handleResponse(text string, ctxObjects ...interface{}) error {
|
func (it *Interact) handleResponse(text string, ctxObjects ...interface{}) error {
|
||||||
|
// we only need response when executing a command
|
||||||
|
switch it.currentState {
|
||||||
|
case StatePublic, StateAuthenticated:
|
||||||
|
return nil
|
||||||
|
|
||||||
|
}
|
||||||
|
|
||||||
args := parseCommand(text)
|
args := parseCommand(text)
|
||||||
|
|
||||||
f, ok := it.statesFunc[it.currentState]
|
f, ok := it.statesFunc[it.currentState]
|
||||||
|
@ -190,7 +201,8 @@ func (it *Interact) SetMessenger(messenger Messenger) {
|
||||||
// builtin initializes the built-in commands
|
// builtin initializes the built-in commands
|
||||||
func (it *Interact) builtin() error {
|
func (it *Interact) builtin() error {
|
||||||
it.Command("/uptime", func(reply Reply) error {
|
it.Command("/uptime", func(reply Reply) error {
|
||||||
reply.Message("uptime")
|
uptime := time.Since(it.startTime)
|
||||||
|
reply.Message(fmt.Sprintf("uptime %s", uptime))
|
||||||
return nil
|
return nil
|
||||||
})
|
})
|
||||||
|
|
||||||
|
|
|
@ -13,14 +13,17 @@ type TelegramReply struct {
|
||||||
message string
|
message string
|
||||||
menu *telebot.ReplyMarkup
|
menu *telebot.ReplyMarkup
|
||||||
buttons [][]telebot.Btn
|
buttons [][]telebot.Btn
|
||||||
|
set bool
|
||||||
}
|
}
|
||||||
|
|
||||||
func (r *TelegramReply) Message(message string) {
|
func (r *TelegramReply) Message(message string) {
|
||||||
r.message = message
|
r.message = message
|
||||||
|
r.set = true
|
||||||
}
|
}
|
||||||
|
|
||||||
func (r *TelegramReply) RemoveKeyboard() {
|
func (r *TelegramReply) RemoveKeyboard() {
|
||||||
r.menu.ReplyKeyboardRemove = true
|
r.menu.ReplyKeyboardRemove = true
|
||||||
|
r.set = true
|
||||||
}
|
}
|
||||||
|
|
||||||
func (r *TelegramReply) AddButton(text string) {
|
func (r *TelegramReply) AddButton(text string) {
|
||||||
|
@ -29,6 +32,7 @@ func (r *TelegramReply) AddButton(text string) {
|
||||||
r.buttons = append(r.buttons, []telebot.Btn{})
|
r.buttons = append(r.buttons, []telebot.Btn{})
|
||||||
}
|
}
|
||||||
r.buttons[len(r.buttons)-1] = append(r.buttons[len(r.buttons)-1], button)
|
r.buttons[len(r.buttons)-1] = append(r.buttons[len(r.buttons)-1], button)
|
||||||
|
r.set = true
|
||||||
}
|
}
|
||||||
|
|
||||||
func (r *TelegramReply) build() {
|
func (r *TelegramReply) build() {
|
||||||
|
@ -47,21 +51,30 @@ type TelegramAuthorizer struct {
|
||||||
func (a *TelegramAuthorizer) Authorize() error {
|
func (a *TelegramAuthorizer) Authorize() error {
|
||||||
a.Telegram.Owner = a.Message.Sender
|
a.Telegram.Owner = a.Message.Sender
|
||||||
a.Telegram.OwnerChat = a.Message.Chat
|
a.Telegram.OwnerChat = a.Message.Chat
|
||||||
|
a.Telegram.authorizing = false
|
||||||
log.Infof("[interact][telegram] authorized owner %+v and chat %+v", a.Message.Sender, a.Message.Chat)
|
log.Infof("[interact][telegram] authorized owner %+v and chat %+v", a.Message.Sender, a.Message.Chat)
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (a *TelegramAuthorizer) StartAuthorizing() {
|
||||||
|
a.Telegram.authorizing = true
|
||||||
|
}
|
||||||
|
|
||||||
type Telegram struct {
|
type Telegram struct {
|
||||||
Bot *telebot.Bot `json:"-"`
|
Bot *telebot.Bot `json:"-"`
|
||||||
|
|
||||||
|
// Private is used to protect the telegram bot, users not authenticated can not see messages or sending commands
|
||||||
|
Private bool `json:"private,omitempty"`
|
||||||
|
|
||||||
|
authorizing bool
|
||||||
|
|
||||||
// Owner is the authorized bot owner
|
// Owner is the authorized bot owner
|
||||||
// This field is exported in order to be stored in file
|
// This field is exported in order to be stored in file
|
||||||
Owner *telebot.User `json:"owner"`
|
Owner *telebot.User `json:"owner,omitempty"`
|
||||||
|
|
||||||
// OwnerChat is the chat of the authorized bot owner
|
// OwnerChat is the chat of the authorized bot owner
|
||||||
// This field is exported in order to be stored in file
|
// This field is exported in order to be stored in file
|
||||||
OwnerChat *telebot.Chat `json:"chat"`
|
OwnerChat *telebot.Chat `json:"chat,omitempty"`
|
||||||
|
|
||||||
// textMessageResponder is used for interact to register its message handler
|
// textMessageResponder is used for interact to register its message handler
|
||||||
textMessageResponder Responder
|
textMessageResponder Responder
|
||||||
|
@ -80,19 +93,32 @@ func (tm *Telegram) SetTextMessageResponder(textMessageResponder Responder) {
|
||||||
|
|
||||||
func (tm *Telegram) Start(context.Context) {
|
func (tm *Telegram) Start(context.Context) {
|
||||||
tm.Bot.Handle(telebot.OnText, func(m *telebot.Message) {
|
tm.Bot.Handle(telebot.OnText, func(m *telebot.Message) {
|
||||||
log.Infof("[interact][telegram] onText: %+v", m)
|
log.Infof("[telegram] onText: %+v", m)
|
||||||
|
|
||||||
|
if tm.Private && !tm.authorizing {
|
||||||
|
// ignore the message directly if it's not authorized yet
|
||||||
|
if tm.Owner == nil {
|
||||||
|
log.Warn("[telegram] telegram is set to private mode, skipping")
|
||||||
|
return
|
||||||
|
} else if tm.Owner != nil && tm.Owner.ID != m.Sender.ID {
|
||||||
|
log.Warnf("[telegram] telegram is set to private mode, owner does not match: %d != %d", tm.Owner.ID, m.Sender.ID)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
authorizer := tm.newAuthorizer(m)
|
authorizer := tm.newAuthorizer(m)
|
||||||
reply := tm.newReply()
|
reply := tm.newReply()
|
||||||
if tm.textMessageResponder != nil {
|
if tm.textMessageResponder != nil {
|
||||||
if err := tm.textMessageResponder(m.Text, reply, authorizer); err != nil {
|
if err := tm.textMessageResponder(m.Text, reply, authorizer); err != nil {
|
||||||
log.WithError(err).Errorf("[interact][telegram] response handling error")
|
log.WithError(err).Errorf("[telegram] response handling error")
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if reply.set {
|
||||||
reply.build()
|
reply.build()
|
||||||
if _, err := tm.Bot.Send(m.Sender, reply.message, reply.menu); err != nil {
|
if _, err := tm.Bot.Send(m.Sender, reply.message, reply.menu); err != nil {
|
||||||
log.WithError(err).Errorf("[interact][telegram] message send error")
|
log.WithError(err).Errorf("[telegram] message send error")
|
||||||
|
}
|
||||||
}
|
}
|
||||||
})
|
})
|
||||||
go tm.Bot.Start()
|
go tm.Bot.Start()
|
||||||
|
@ -100,18 +126,21 @@ func (tm *Telegram) Start(context.Context) {
|
||||||
|
|
||||||
func (tm *Telegram) AddCommand(command string, responder Responder) {
|
func (tm *Telegram) AddCommand(command string, responder Responder) {
|
||||||
tm.Bot.Handle(command, func(m *telebot.Message) {
|
tm.Bot.Handle(command, func(m *telebot.Message) {
|
||||||
|
authorizer := tm.newAuthorizer(m)
|
||||||
reply := tm.newReply()
|
reply := tm.newReply()
|
||||||
if err := responder(m.Payload, reply); err != nil {
|
if err := responder(m.Payload, reply, authorizer); err != nil {
|
||||||
log.WithError(err).Errorf("[interact][telegram] responder error")
|
log.WithError(err).Errorf("[interact][telegram] responder error")
|
||||||
tm.Bot.Send(m.Sender, fmt.Sprintf("error: %v", err))
|
tm.Bot.Send(m.Sender, fmt.Sprintf("error: %v", err))
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
// build up the response objects
|
// build up the response objects
|
||||||
|
if reply.set {
|
||||||
reply.build()
|
reply.build()
|
||||||
if _, err := tm.Bot.Send(m.Sender, reply.message, reply.menu); err != nil {
|
if _, err := tm.Bot.Send(m.Sender, reply.message, reply.menu); err != nil {
|
||||||
log.WithError(err).Errorf("[interact][telegram] message send error")
|
log.WithError(err).Errorf("[interact][telegram] message send error")
|
||||||
}
|
}
|
||||||
|
}
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue
Block a user