From b42027527af1a9ce5be139560b48ee7996c3d4d6 Mon Sep 17 00:00:00 2001 From: Ettore Dreucci Date: Mon, 25 Jun 2018 12:36:54 +0200 Subject: [PATCH] User info saved in db. Signed-off-by: Ettore Dreucci --- src/redisAPI.go | 117 ++++++++++++++++++++++++++++++++++++++++----- src/sys.go | 9 ++-- src/telegramAPI.go | 48 ++++++++++++++++--- 3 files changed, 153 insertions(+), 21 deletions(-) diff --git a/src/redisAPI.go b/src/redisAPI.go index 0092556..ff10c18 100644 --- a/src/redisAPI.go +++ b/src/redisAPI.go @@ -2,11 +2,13 @@ package main import ( "bufio" + "encoding/json" "errors" "fmt" "log" "os" "regexp" + "strconv" "strings" "github.com/dixonwille/wmenu" @@ -15,8 +17,11 @@ import ( ) const ( - tkSet = "botTokens" - botHash = "botInfos" + tkSet = "botToken" + botHash = "botInfo" + userSet = "userID" + userHash = "userInfo" + authUserSet = "authUser" ) var redisClient *redis.Client @@ -26,10 +31,16 @@ var ( ErrRedisConnection = errors.New("redis: couldn't connect to remote instance") //ErrRedisAddSet is thrown when it's not possible to add a key in a set ErrRedisAddSet = errors.New("redis: couldn't add key in set") + //ErrRedisRemSet is thrown when it's not possible to remove a key from a given set + ErrRedisRemSet = errors.New("redis: couldn't remove key from set") //ErrRedisRetriveSet is thrown when it's not possible to retrive keys from a set ErrRedisRetriveSet = errors.New("redis: couldn't retrive keys from set") + //ErrRedisCheckSet is thrown when it's not possible to check if a key is in a given set + ErrRedisCheckSet = errors.New("redis: couldn't check if key is in set") //ErrRedisAddHash is thrown when it's not possible to add a key in a hash ErrRedisAddHash = errors.New("redis: couldn't add key in hash") + //ErrRedisDelHash is thrown when it's not possible to remove a key from a hash + ErrRedisDelHash = errors.New("redis: couldn't remove key from hash") //ErrTokenParsing is thrown when it's not possible to parse the bot token ErrTokenParsing = errors.New("botToken: cannot parse token") //ErrTokenInvalid is thrown when the string parsed isn't a valid telegram bot token @@ -38,6 +49,8 @@ var ( ErrAddToken = errors.New("couldn't add one or more tokens") //ErrRemoveToken is thrown when one or more bot tokens hasn't been removed ErrRemoveToken = errors.New("couldn't remove one or more tokens") + //ErrJSONMarshall is thrown when it's impossible to marshall a given struct + ErrJSONMarshall = errors.New("json: couldn't marshall struct") ) func redisInit(addr string, pwd string, db int) error { @@ -131,21 +144,29 @@ func removeBotTokens() error { menu.AllowMultiple() menu.LoopOnInvalid() menu.Action(func(opts []wmenu.Opt) error { + var returnErr error for _, opt := range opts { if opt.Value == nil { log.Println("Couldn't remove bot: nil token") - return ErrNilPointer - } - err := removeBotToken(opt.Value.(string)) - if err != nil { - log.Printf("Couldn't remove bot: %v", err) + returnErr = ErrNilPointer + } else { + err := removeBotToken(opt.Value.(string)) + if err != nil { + log.Printf("Couldn't remove bot token: %v", err) + } + err = removeBotInfo(opt.Value.(string)) + if err != nil { + log.Printf("Couldn't remove bot info: %v", err) + } } } - return nil + return returnErr }) //for _, token := range tokens { - for token, botInfo := range botsInfo { - menu.Option(botInfo, token, false, nil) + for token, jsonBotInfo := range botsInfo { + botInfo := &tb.Bot{} + json.Unmarshal([]byte(jsonBotInfo), &botInfo) + menu.Option(botInfo.Me.Username, token, false, nil) } err = menu.Run() if err != nil { @@ -182,11 +203,16 @@ func getBotTokens() ([]string, error) { return tokens, nil } -func addBotInfo(bot *tb.Bot, botToken string) error { +func addBotInfo(botToken string, bot *tb.Bot) error { if redisClient == nil { return ErrNilPointer } - err := redisClient.HSet(botHash, botToken, bot.Me.Username).Err() + jsonBot, err := json.Marshal(&bot) + if err != nil { + log.Printf("Error marshalling bot info: %v", err) + return ErrJSONMarshall + } + err = redisClient.HSet(botHash, botToken, string(jsonBot)).Err() if err != nil { log.Printf("Error in adding bot info: %v", err) return ErrRedisAddHash @@ -194,3 +220,70 @@ func addBotInfo(bot *tb.Bot, botToken string) error { return nil } + +func removeBotInfo(botToken string) error { + if redisClient == nil { + return ErrNilPointer + } + err := redisClient.HDel(botHash, botToken).Err() + if err != nil { + log.Printf("Error in removing bot info: %v", err) + return ErrRedisDelHash + } + return nil +} + +func addUser(user *tb.User) error { + if redisClient == nil { + return ErrNilPointer + } + err := redisClient.SAdd(userSet, user.ID).Err() + if err != nil { + log.Printf("Error in adding user ID: %v", err) + return ErrRedisAddSet + } + jsonUser, err := json.Marshal(&user) + if err != nil { + log.Printf("Error in marshalling user to json: %v", err) + return ErrJSONMarshall + } + err = redisClient.HSet(userHash, strconv.Itoa(user.ID), jsonUser).Err() + if err != nil { + log.Printf("Error adding user info in hash: %v", err) + return ErrRedisAddHash + } + + return nil +} + +func isAuthrizedUser(userID int) (bool, error) { + if redisClient == nil { + return false, ErrNilPointer + } + auth, err := redisClient.SIsMember(authUserSet, strconv.Itoa(userID)).Result() + if err != nil { + log.Printf("Error checking if user is authorized: %v", err) + return false, ErrRedisCheckSet + } + return auth, nil +} + +func authorizeUser(userID int, authorized bool) error { + if redisClient == nil { + return ErrNilPointer + } + if authorized { + err := redisClient.SAdd(authUserSet, strconv.Itoa(userID)).Err() + if err != nil { + log.Printf("Error adding token to set: %v", err) + return ErrRedisAddSet + } + } else { + err := redisClient.SRem(authUserSet, strconv.Itoa(userID)).Err() + if err != nil { + log.Printf("Error removing token from set: %v", err) + return ErrRedisRemSet + } + } + return nil +} diff --git a/src/sys.go b/src/sys.go index 97aecc7..e086ba7 100644 --- a/src/sys.go +++ b/src/sys.go @@ -23,6 +23,7 @@ type flags struct { var cmdFlags flags var ( + welcomeMessage = "Welcome in barandaBot! Here you can control the bot(s) options and configurations." //ErrStdRead is thrown when it's not possible to read from the standard input ErrStdRead = errors.New("stdin: couldn't read string from stdin") //ErrMainMenu is thrown when a menu couldn't be started @@ -71,7 +72,7 @@ func getFlags() error { } func mainMenu() error { - fmt.Println("Welcome in barandaBot! Here you can control the bot(s) options and configurations.") + fmt.Println(welcomeMessage) menu := wmenu.NewMenu("What do you want to do?") menu.LoopOnInvalid() menu.Option("Start Bot(s)", nil, true, func(opt wmenu.Opt) error { @@ -84,12 +85,14 @@ func mainMenu() error { return removeBotTokens() }) + var returnErr error + for { err := menu.Run() if err != nil { log.Printf("Error in main menu: %v", err) - return ErrMainMenu + returnErr = ErrMainMenu } } - return nil + return returnErr } diff --git a/src/telegramAPI.go b/src/telegramAPI.go index 1911cf6..17c3ca4 100644 --- a/src/telegramAPI.go +++ b/src/telegramAPI.go @@ -3,6 +3,7 @@ package main import ( "errors" "log" + "strconv" "sync" "time" @@ -14,6 +15,8 @@ var bots []*tb.Bot var ( //ErrNilPointer is thrown when a pointer is nil ErrNilPointer = errors.New("pointer is nil") + //ErrIDFromMsg is thrown when the message doesn't contain user infos + ErrIDFromMsg = errors.New("telegram: couldn't retrive user ID from message") ) func botsInit() error { @@ -29,16 +32,44 @@ func botsInit() error { } for _, token := range tokens { + poller := &tb.LongPoller{Timeout: 15 * time.Second} + middlePoller := tb.NewMiddlewarePoller(poller, func(upd *tb.Update) bool { + if upd.Message == nil { + return true + } + if upd.Message.Sender != nil { + err := addUser(upd.Message.Sender) + if err != nil { + log.Printf("Error in adding user info: %v", err) + } + err = authorizeUser(upd.Message.Sender.ID, true) + if err != nil { + log.Printf("Error in authorizing user: %v", err) + } + } else { + log.Printf("%v", ErrIDFromMsg) + } + auth, err := isAuthrizedUser(upd.Message.Sender.ID) + if err != nil { + log.Printf("Error checking if user is authorized: %v", err) + } + if !auth { + return false + } + + return true + }) + bot, err := tb.NewBot(tb.Settings{ Token: token, - Poller: &tb.LongPoller{Timeout: 10 * time.Second}, + Poller: middlePoller, }) if err != nil { log.Printf("Error in enstablishing connection for bot %s: %v", bot.Me.Username, err) } else { bots = append(bots, bot) - err = addBotInfo(bot, token) + err = addBotInfo(token, bot) if err != nil { log.Printf("Error: bot %s info couldn't be added: %v", bot.Me.Username, err) } @@ -59,23 +90,28 @@ func botsStart() error { var wg sync.WaitGroup for i := range bots { - defer wg.Done() if bots[i] != nil { - go botStart(bots[i]) + wg.Add(1) + go botStart(bots[i], &wg) } } + wg.Wait() return nil } -func botStart(bot *tb.Bot) error { +func botStart(bot *tb.Bot, wg *sync.WaitGroup) error { + defer wg.Done() if bot == nil { return ErrNilPointer } - log.Printf("Started bot %s", bot.Me.Username) + log.Printf("Started %s", bot.Me.Username) bot.Handle("/hello", func(m *tb.Message) { bot.Send(m.Sender, "hello world") }) + bot.Handle("/userID", func(m *tb.Message) { + bot.Send(m.Sender, strconv.Itoa(m.Sender.ID)) + }) bot.Start()