From 17322cbc090338b5e4cfed1517a0912795da22ca Mon Sep 17 00:00:00 2001 From: c9s Date: Fri, 14 Jan 2022 02:36:06 +0800 Subject: [PATCH] interact: improve authentication process --- examples/interact/main.go | 7 +++--- pkg/interact/auth.go | 7 +++++- pkg/interact/interact.go | 14 ++++++++++- pkg/interact/telegram.go | 53 ++++++++++++++++++++++++++++++--------- 4 files changed, 64 insertions(+), 17 deletions(-) diff --git a/examples/interact/main.go b/examples/interact/main.go index 2e5b049b4..8524b4887 100644 --- a/examples/interact/main.go +++ b/examples/interact/main.go @@ -131,13 +131,14 @@ func main() { globalInteraction := interact.New() globalInteraction.SetMessenger(&interact.Telegram{ - Bot: b, + Private: true, + Bot: b, }) globalInteraction.AddCustomInteraction(&interact.AuthInteract{ Strict: true, - Mode: interact.AuthModeToken, - Token: "123", + Mode: interact.AuthModeToken, + Token: "123", }) globalInteraction.AddCustomInteraction(&PositionInteraction{}) diff --git a/pkg/interact/auth.go b/pkg/interact/auth.go index 9d7ad19a8..c984247b8 100644 --- a/pkg/interact/auth.go +++ b/pkg/interact/auth.go @@ -20,6 +20,7 @@ const ( var ErrAuthenticationFailed = errors.New("authentication failed") type Authorizer interface { + StartAuthorizing() Authorize() error } @@ -50,8 +51,9 @@ func (it *AuthInteract) Commands(interact *Interact) { 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") + authorizer.StartAuthorizing() return nil }).Next(func(token string, reply Reply) error { if token == it.Token { @@ -72,6 +74,7 @@ func (it *AuthInteract) Commands(interact *Interact) { }).NamedNext(StateAuthenticated, func(code string, reply Reply, authorizer Authorizer) error { if totp.Validate(code, it.OneTimePasswordKey.Secret()) { reply.Message("Great! You're authenticated!") + interact.SetOriginState(StateAuthenticated) return authorizer.Authorize() } @@ -87,12 +90,14 @@ func (it *AuthInteract) Commands(interact *Interact) { case AuthModeToken: if code == it.Token { reply.Message("Great! You're authenticated!") + interact.SetOriginState(StateAuthenticated) return authorizer.Authorize() } case AuthModeOTP: if totp.Validate(code, it.OneTimePasswordKey.Secret()) { reply.Message("Great! You're authenticated!") + interact.SetOriginState(StateAuthenticated) return authorizer.Authorize() } } diff --git a/pkg/interact/interact.go b/pkg/interact/interact.go index abc2ddda9..707331616 100644 --- a/pkg/interact/interact.go +++ b/pkg/interact/interact.go @@ -7,6 +7,7 @@ import ( "strconv" "strings" "text/scanner" + "time" log "github.com/sirupsen/logrus" ) @@ -47,6 +48,8 @@ type Messenger interface { // Interact implements the interaction between bot and message software. type Interact struct { + startTime time.Time + // commands is the default public command map commands map[string]*Command @@ -63,6 +66,7 @@ type Interact struct { func New() *Interact { return &Interact{ + startTime: time.Now(), commands: make(map[string]*Command), originState: StatePublic, currentState: StatePublic, @@ -114,6 +118,13 @@ func (it *Interact) setState(s State) { } 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) f, ok := it.statesFunc[it.currentState] @@ -190,7 +201,8 @@ func (it *Interact) SetMessenger(messenger Messenger) { // builtin initializes the built-in commands func (it *Interact) builtin() 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 }) diff --git a/pkg/interact/telegram.go b/pkg/interact/telegram.go index f151af8f7..e38763cd1 100644 --- a/pkg/interact/telegram.go +++ b/pkg/interact/telegram.go @@ -13,14 +13,17 @@ type TelegramReply struct { message string menu *telebot.ReplyMarkup buttons [][]telebot.Btn + set bool } func (r *TelegramReply) Message(message string) { r.message = message + r.set = true } func (r *TelegramReply) RemoveKeyboard() { r.menu.ReplyKeyboardRemove = true + r.set = true } 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[len(r.buttons)-1] = append(r.buttons[len(r.buttons)-1], button) + r.set = true } func (r *TelegramReply) build() { @@ -47,21 +51,30 @@ type TelegramAuthorizer struct { func (a *TelegramAuthorizer) Authorize() error { a.Telegram.Owner = a.Message.Sender 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) return nil } +func (a *TelegramAuthorizer) StartAuthorizing() { + a.Telegram.authorizing = true +} + type Telegram struct { 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 // 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 // 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 Responder @@ -80,19 +93,32 @@ func (tm *Telegram) SetTextMessageResponder(textMessageResponder Responder) { func (tm *Telegram) Start(context.Context) { 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) reply := tm.newReply() if tm.textMessageResponder != 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") } } - reply.build() - if _, err := tm.Bot.Send(m.Sender, reply.message, reply.menu); err != nil { - log.WithError(err).Errorf("[interact][telegram] message send error") + if reply.set { + reply.build() + if _, err := tm.Bot.Send(m.Sender, reply.message, reply.menu); err != nil { + log.WithError(err).Errorf("[telegram] message send error") + } } }) go tm.Bot.Start() @@ -100,17 +126,20 @@ func (tm *Telegram) Start(context.Context) { func (tm *Telegram) AddCommand(command string, responder Responder) { tm.Bot.Handle(command, func(m *telebot.Message) { + authorizer := tm.newAuthorizer(m) 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") tm.Bot.Send(m.Sender, fmt.Sprintf("error: %v", err)) return } // build up the response objects - reply.build() - if _, err := tm.Bot.Send(m.Sender, reply.message, reply.menu); err != nil { - log.WithError(err).Errorf("[interact][telegram] message send error") + if reply.set { + reply.build() + if _, err := tm.Bot.Send(m.Sender, reply.message, reply.menu); err != nil { + log.WithError(err).Errorf("[interact][telegram] message send error") + } } }) }