User info saved in db.

Signed-off-by: Ettore Dreucci <ettore.dreucci@gmail.com>
This commit is contained in:
2018-06-25 12:36:54 +02:00
parent 5f5dfad1ed
commit b42027527a
3 changed files with 153 additions and 21 deletions

View File

@@ -2,11 +2,13 @@ package main
import ( import (
"bufio" "bufio"
"encoding/json"
"errors" "errors"
"fmt" "fmt"
"log" "log"
"os" "os"
"regexp" "regexp"
"strconv"
"strings" "strings"
"github.com/dixonwille/wmenu" "github.com/dixonwille/wmenu"
@@ -15,8 +17,11 @@ import (
) )
const ( const (
tkSet = "botTokens" tkSet = "botToken"
botHash = "botInfos" botHash = "botInfo"
userSet = "userID"
userHash = "userInfo"
authUserSet = "authUser"
) )
var redisClient *redis.Client var redisClient *redis.Client
@@ -26,10 +31,16 @@ var (
ErrRedisConnection = errors.New("redis: couldn't connect to remote instance") 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 is thrown when it's not possible to add a key in a set
ErrRedisAddSet = errors.New("redis: couldn't add key in 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 is thrown when it's not possible to retrive keys from a set
ErrRedisRetriveSet = errors.New("redis: couldn't retrive keys from 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 is thrown when it's not possible to add a key in a hash
ErrRedisAddHash = errors.New("redis: couldn't add key in 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 is thrown when it's not possible to parse the bot token
ErrTokenParsing = errors.New("botToken: cannot parse token") ErrTokenParsing = errors.New("botToken: cannot parse token")
//ErrTokenInvalid is thrown when the string parsed isn't a valid telegram bot 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") ErrAddToken = errors.New("couldn't add one or more tokens")
//ErrRemoveToken is thrown when one or more bot tokens hasn't been removed //ErrRemoveToken is thrown when one or more bot tokens hasn't been removed
ErrRemoveToken = errors.New("couldn't remove one or more tokens") 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 { func redisInit(addr string, pwd string, db int) error {
@@ -131,21 +144,29 @@ func removeBotTokens() error {
menu.AllowMultiple() menu.AllowMultiple()
menu.LoopOnInvalid() menu.LoopOnInvalid()
menu.Action(func(opts []wmenu.Opt) error { menu.Action(func(opts []wmenu.Opt) error {
var returnErr error
for _, opt := range opts { for _, opt := range opts {
if opt.Value == nil { if opt.Value == nil {
log.Println("Couldn't remove bot: nil token") log.Println("Couldn't remove bot: nil token")
return ErrNilPointer returnErr = ErrNilPointer
} } else {
err := removeBotToken(opt.Value.(string)) err := removeBotToken(opt.Value.(string))
if err != nil { if err != nil {
log.Printf("Couldn't remove bot: %v", err) 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 := range tokens {
for token, botInfo := range botsInfo { for token, jsonBotInfo := range botsInfo {
menu.Option(botInfo, token, false, nil) botInfo := &tb.Bot{}
json.Unmarshal([]byte(jsonBotInfo), &botInfo)
menu.Option(botInfo.Me.Username, token, false, nil)
} }
err = menu.Run() err = menu.Run()
if err != nil { if err != nil {
@@ -182,11 +203,16 @@ func getBotTokens() ([]string, error) {
return tokens, nil return tokens, nil
} }
func addBotInfo(bot *tb.Bot, botToken string) error { func addBotInfo(botToken string, bot *tb.Bot) error {
if redisClient == nil { if redisClient == nil {
return ErrNilPointer 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 { if err != nil {
log.Printf("Error in adding bot info: %v", err) log.Printf("Error in adding bot info: %v", err)
return ErrRedisAddHash return ErrRedisAddHash
@@ -194,3 +220,70 @@ func addBotInfo(bot *tb.Bot, botToken string) error {
return nil 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
}

View File

@@ -23,6 +23,7 @@ type flags struct {
var cmdFlags flags var cmdFlags flags
var ( 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 is thrown when it's not possible to read from the standard input
ErrStdRead = errors.New("stdin: couldn't read string from stdin") ErrStdRead = errors.New("stdin: couldn't read string from stdin")
//ErrMainMenu is thrown when a menu couldn't be started //ErrMainMenu is thrown when a menu couldn't be started
@@ -71,7 +72,7 @@ func getFlags() error {
} }
func mainMenu() 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 := wmenu.NewMenu("What do you want to do?")
menu.LoopOnInvalid() menu.LoopOnInvalid()
menu.Option("Start Bot(s)", nil, true, func(opt wmenu.Opt) error { menu.Option("Start Bot(s)", nil, true, func(opt wmenu.Opt) error {
@@ -84,12 +85,14 @@ func mainMenu() error {
return removeBotTokens() return removeBotTokens()
}) })
var returnErr error
for { for {
err := menu.Run() err := menu.Run()
if err != nil { if err != nil {
log.Printf("Error in main menu: %v", err) log.Printf("Error in main menu: %v", err)
return ErrMainMenu returnErr = ErrMainMenu
} }
} }
return nil return returnErr
} }

View File

@@ -3,6 +3,7 @@ package main
import ( import (
"errors" "errors"
"log" "log"
"strconv"
"sync" "sync"
"time" "time"
@@ -14,6 +15,8 @@ var bots []*tb.Bot
var ( var (
//ErrNilPointer is thrown when a pointer is nil //ErrNilPointer is thrown when a pointer is nil
ErrNilPointer = errors.New("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 { func botsInit() error {
@@ -29,16 +32,44 @@ func botsInit() error {
} }
for _, token := range tokens { 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{ bot, err := tb.NewBot(tb.Settings{
Token: token, Token: token,
Poller: &tb.LongPoller{Timeout: 10 * time.Second}, Poller: middlePoller,
}) })
if err != nil { if err != nil {
log.Printf("Error in enstablishing connection for bot %s: %v", bot.Me.Username, err) log.Printf("Error in enstablishing connection for bot %s: %v", bot.Me.Username, err)
} else { } else {
bots = append(bots, bot) bots = append(bots, bot)
err = addBotInfo(bot, token) err = addBotInfo(token, bot)
if err != nil { if err != nil {
log.Printf("Error: bot %s info couldn't be added: %v", bot.Me.Username, err) 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 var wg sync.WaitGroup
for i := range bots { for i := range bots {
defer wg.Done()
if bots[i] != nil { if bots[i] != nil {
go botStart(bots[i]) wg.Add(1)
go botStart(bots[i], &wg)
} }
} }
wg.Wait()
return nil return nil
} }
func botStart(bot *tb.Bot) error { func botStart(bot *tb.Bot, wg *sync.WaitGroup) error {
defer wg.Done()
if bot == nil { if bot == nil {
return ErrNilPointer 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.Handle("/hello", func(m *tb.Message) {
bot.Send(m.Sender, "hello world") bot.Send(m.Sender, "hello world")
}) })
bot.Handle("/userID", func(m *tb.Message) {
bot.Send(m.Sender, strconv.Itoa(m.Sender.ID))
})
bot.Start() bot.Start()