diff --git a/go.mod b/go.mod index e69b2d6..9e8d3b9 100644 --- a/go.mod +++ b/go.mod @@ -9,4 +9,5 @@ require ( github.com/go-telegram-bot-api/telegram-bot-api v4.6.4+incompatible github.com/jinzhu/gorm v1.9.12 github.com/technoweenie/multipartstreamer v1.0.1 // indirect + github.com/sahilm/fuzzy v0.1.0 ) diff --git a/go.sum b/go.sum index e793f73..0bbb28f 100644 --- a/go.sum +++ b/go.sum @@ -36,6 +36,8 @@ github.com/mattn/go-sqlite3 v2.0.1+incompatible h1:xQ15muvnzGBHpIpdrNi1DA5x0+TcB github.com/mattn/go-sqlite3 v2.0.1+incompatible/go.mod h1:FPy6KqzDD04eiIsT53CuJW3U88zkxoIYsOqkbpncsNc= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/sahilm/fuzzy v0.1.0 h1:FzWGaw2Opqyu+794ZQ9SYifWv2EIXpwP4q8dY1kDAwI= +github.com/sahilm/fuzzy v0.1.0/go.mod h1:VFvziUEIMCrT6A6tw2RFIXPXXmzXbOsSHF0DOI8ZK9Y= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/testify v1.3.0 h1:TivCn/peBQ7UY8ooIcPgZFpTNSz0Q2U6UrFlUfqbe0Q= github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= diff --git a/internal/bot/bot.go b/internal/bot/bot.go index 627eee2..78b2036 100644 --- a/internal/bot/bot.go +++ b/internal/bot/bot.go @@ -9,6 +9,7 @@ import ( "git.1750studios.com/ToddShepard/DB640/internal/telegram" "git.1750studios.com/ToddShepard/DB640/internal/twitter" "github.com/jinzhu/gorm" + "github.com/sahilm/fuzzy" ) var stream *twitter.Stream @@ -32,7 +33,11 @@ func Init() { go twitter.StreamDemux(stream, handleHashtagTweet) log.Printf("[TWITTER] Connection established\n") - updates, err = telegram.Init() + err = telegram.Init() + if err != nil { + log.Fatalf("Could not establish telegram connection: %+v", err) + } + updates, err = telegram.GetChan() if err != nil { log.Fatalf("Could not establish telegram connection: %+v", err) } @@ -43,6 +48,7 @@ func Init() { // DeInit stops the stream and deinitzializes the bot func DeInit() { stream.Stop() + telegram.DeInit() database.Close() } @@ -98,11 +104,41 @@ func handleTelegram() { // ignore any non-Message Updates continue } - log.Printf("[TELEGRAM] %s: %s\n", update.Message.From.UserName, update.Message.Text) - if update.Message.Text == "/start" || update.Message.Text == "/help" { - reply := "Willkommen beim DB 640 Telegram Bot!\n\nEinfach den gewünschten DB 640 Betriebsstellencode schicken, und der Bot antwortet mit der zugehörigen Betriebsstelle!\n\nZum Beispiel: Nb -> Wiener Neustadt Hbf (in Nb)" - telegram.SendReply(reply, update) - continue + log.Printf("[TELEGRAM] %s: %s\n", update.Message.From.String(), update.Message.Text) + + if update.Message.IsCommand() { + switch update.Message.Command() { + case "start", "help": + reply := "Willkommen beim DB 640 Telegram Bot!\n\nEinfach den gewünschten DB 640 Betriebsstellencode schicken, und der Bot antwortet mit der zugehörigen Betriebsstelle!\n\nZum Beispiel: Nb -> Wiener Neustadt Hbf (in Nb)" + telegram.SendReply(reply, update) + continue + case "find": + if update.Message.CommandArguments() == "" { + reply := "Benutze /find um einen Code für eine Betriebsstelle zu finden!" + telegram.SendReply(reply, update) + continue + } + var bs database.Betriebsstellen + var reply string + if database.Db.Find(&bs, "name LIKE ?", "%"+update.Message.CommandArguments()+"%").Error != gorm.ErrRecordNotFound { + + results := fuzzy.FindFrom(update.Message.CommandArguments(), bs) + for i, r := range results { + reply = reply + bs[r.Index].Code + ": " + bs[r.Index].Name + "\n" + if i == 49 { + break + } + } + reply = reply[0 : len(reply)-1] + } else { + reply = "Keine Betriebsstelle mit Namen '" + update.Message.CommandArguments() + "' gefunden!" + } + err := telegram.SendReply(reply, update) + if err != nil { + log.Printf("Cannot send reply, error: %+v", err) + } + continue + } } var bs database.Betriebsstelle diff --git a/internal/database/database.go b/internal/database/database.go index 723a58b..ae99943 100644 --- a/internal/database/database.go +++ b/internal/database/database.go @@ -1,6 +1,8 @@ package database import ( + "time" + "github.com/jinzhu/gorm" // SQLite dialect for gorm @@ -9,11 +11,40 @@ import ( // Betriebsstelle defines the database model of the Betriebsstelle table type Betriebsstelle struct { - gorm.Model + ID uint `gorm:"primary_key"` + CreatedAt time.Time + UpdatedAt time.Time + DeletedAt *time.Time `sql:"index"` + Code string `gorm:"unique_index"` Name string } +// Betriebsstellen is an slice of Betriebsstelle +type Betriebsstellen []Betriebsstelle + +// String outputs the Name of the Betriebssstelle +func (b Betriebsstellen) String(i int) string { + return b[i].Name +} + +// Len outputs the length of the Betriebsstellen slice +func (b Betriebsstellen) Len() int { + return len(b) +} + +// TGChat represents a telegram chat the bot is part of +type TGChat struct { + ID int64 `gorm:"primary_key"` + CreatedAt time.Time + UpdatedAt time.Time + DeletedAt *time.Time `sql:"index"` + + FirstName string + LastName string + UserName string +} + // Db is the GORM database handle var Db *gorm.DB @@ -23,7 +54,7 @@ func Open(dialect string, connection string) (err error) { if err != nil { return err } - Db.AutoMigrate(&Betriebsstelle{}) + Db.AutoMigrate(&Betriebsstelle{}, &TGChat{}) return nil } diff --git a/internal/telegram/telegram.go b/internal/telegram/telegram.go index f1a688f..ab40e18 100644 --- a/internal/telegram/telegram.go +++ b/internal/telegram/telegram.go @@ -4,7 +4,9 @@ import ( "log" "git.1750studios.com/ToddShepard/DB640/internal/config" + "git.1750studios.com/ToddShepard/DB640/internal/database" tgbotapi "github.com/go-telegram-bot-api/telegram-bot-api" + "github.com/jinzhu/gorm" ) // UpdateChan is telegram UpdatesChannel @@ -13,28 +15,62 @@ type UpdateChan = tgbotapi.UpdatesChannel var bot *tgbotapi.BotAPI // Init initzializes the telegram bot -func Init() (UpdateChan, error) { - var err error +func Init() (err error) { bot, err = tgbotapi.NewBotAPI(config.C.Telegram.APIKey) if err != nil { - return nil, err + return } log.Printf("Authorized on account %s", bot.Self.UserName) + return +} +// GetChan returns the Telegram update channel +func GetChan() (UpdateChan, error) { u := tgbotapi.NewUpdate(0) u.Timeout = 60 return bot.GetUpdatesChan(u) } +// DeInit deinitzialisizes the bot +func DeInit() { + bot.StopReceivingUpdates() +} + // SendReply sends a reply to the given update func SendReply(reply string, update tgbotapi.Update) error { + var chat database.TGChat + chat.ID = update.Message.Chat.ID + + if database.Db.First(&chat).Error != gorm.ErrRecordNotFound { + chat.FirstName = update.Message.Chat.FirstName + chat.LastName = update.Message.Chat.LastName + chat.UserName = update.Message.Chat.UserName + database.Db.Save(&chat) + } else { + chat.FirstName = update.Message.Chat.FirstName + chat.LastName = update.Message.Chat.LastName + chat.UserName = update.Message.Chat.UserName + database.Db.Create(&chat) + } msg := tgbotapi.NewMessage(update.Message.Chat.ID, reply) - //msg.ReplyToMessageID = update.Message.MessageID if _, err := bot.Send(msg); err != nil { return err } return nil } + +// SendAll sends a message to all active users +func SendAll(msg string) { + log.Printf("[TELEGRAM] Sending to all chats: %s", msg) + var chats []database.TGChat + database.Db.Find(&chats) + for _, chat := range chats { + tgmsg := tgbotapi.NewMessage(chat.ID, msg) + if _, err := bot.Send(tgmsg); err != nil && err.Error() == "Forbidden: bot was blocked by the user" { + database.Db.Unscoped().Delete(&chat) + } + } +} diff --git a/tools/sendmsg/main.go b/tools/sendmsg/main.go new file mode 100644 index 0000000..535be12 --- /dev/null +++ b/tools/sendmsg/main.go @@ -0,0 +1,66 @@ +package main + +import ( + "bufio" + "flag" + "io/ioutil" + "log" + "os" + + "git.1750studios.com/ToddShepard/DB640/internal/config" + "git.1750studios.com/ToddShepard/DB640/internal/database" + "git.1750studios.com/ToddShepard/DB640/internal/telegram" +) + +func main() { + msg := flag.String("m", "", "Specifies the message to send") + file := flag.String("f", "", "Specifies the file of the message to send") + cfg := flag.String("c", "", "Config file") + flag.Parse() + if *cfg == "" { + log.Fatalf("Config file must not be empty!") + } + if *msg == "" && *file == "" { + log.Fatalf("Message (or file) must not be empty!") + } + if *msg != "" && *file != "" { + log.Fatalf("Only use message or file!") + } + + err := config.LoadConfig(*cfg) + if err != nil { + log.Fatalf("Config error: %+v", err) + } + + err = database.Open(config.C.Database.Dialect, config.C.Database.Connection) + if err != nil { + log.Fatalf("Could not establish database connection: %+v", err) + } + log.Printf("[DATABASE] Connection established\n") + + err = telegram.Init() + if err != nil { + log.Fatalf("Could not establish telegram connection: %+v", err) + } + log.Printf("[TELEGRAM] Connection established\n") + + if *msg != "" { + telegram.SendAll(*msg) + } else { + var reader *bufio.Reader + if *file == "-" { + reader = bufio.NewReader(os.Stdin) + } else { + file, err := os.Open(*file) + if err != nil { + log.Fatalf("File error: %+v", err) + } + reader = bufio.NewReader(file) + } + bytes, err := ioutil.ReadAll(reader) + if err != nil { + log.Fatalf("ReadAll error: %+v", err) + } + telegram.SendAll(string(bytes)) + } +}