Skip to content

Commit

Permalink
Split main into smaller functions and add TrollShieldBot interface
Browse files Browse the repository at this point in the history 
That way we can mock on tests easier.
  • Loading branch information
@ryukinix
ryukinix committed Jul 24, 2020
1 parent 52cee34 commit 92bd96d
Showing 1 changed file with 106 additions and 81 deletions.
187 changes: 106 additions & 81 deletions troll_shield.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,9 +12,16 @@ import (
"strings"
"sync"

tgbotapi "github.com/go-telegram-bot-api/telegram-bot-api"
telegram "github.com/go-telegram-bot-api/telegram-bot-api"
)

// TrollShieldBot aggregate the methods used by my bot to keep mocking easier
type TrollShieldBot interface {
GetChatMember(telegram.ChatConfigWithUser) (telegram.ChatMember, error)
KickChatMember(telegram.KickChatMemberConfig) (telegram.APIResponse, error)
Send(telegram.Chattable) (telegram.Message, error)
}

// blacklist groups, member from that groups will be kicked automatically
var trollGroups = []string{
"@ccppbrasil",
Expand All @@ -31,17 +38,78 @@ const logfile = "troll-shield.log"

var log = logger.New(os.Stderr, "", logger.LstdFlags)

// messageEvent return true if is a message event
func messageEvent(update *telegram.Update) bool {
return update.Message != nil
}

// newChatMemberEvent return true if a new member joined to chat
func newChatMemberEvent(update *telegram.Update) bool {
return messageEvent(update) && update.Message.NewChatMembers != nil
}

// fromChatEvent return true if the message is from a specific chat
func fromChatEvent(update *telegram.Update, username string) bool {
chat := update.Message.Chat
return messageEvent(update) && chat != nil && (chat.UserName == username || chat.Title == username)
}

// getUserName return the most meaningful name available from a telegram user
// if it has a @username, return it
// if not, try to return FirstName + LastName
// otherwise, only return the FirstName
func getUserName(user telegram.User) string {
username := user.FirstName
if user.UserName != "" {
username = fmt.Sprintf("@%v", user.UserName)
} else if user.LastName != "" {
username = fmt.Sprintf("%s %s", user.FirstName, user.LastName)
}
return username
}

func getUpdates(bot *telegram.BotAPI) telegram.UpdatesChannel {
u := telegram.NewUpdate(0)
u.Timeout = 60
updates, err := bot.GetUpdatesChan(u)
if err != nil {
log.Printf("getUpdates error: %v", err)
}

return updates
}

func reply(bot TrollShieldBot, update *telegram.Update, text string) {
msg := telegram.NewMessage(update.Message.Chat.ID, text)
msg.ReplyToMessageID = update.Message.MessageID

_, err := bot.Send(msg)
if err != nil {
log.Printf("[!] Send msg failed: %v", err)
}
}

func welcomeMessage(bot *telegram.BotAPI, update *telegram.Update, member telegram.User) {
username := getUserName(member)
text := fmt.Sprintf(
`Olá %s! Seja bem-vindo ao grupo oficial de Common Lisp do Brasil.
Leia as regras em: https://lisp.com.br/rules.html.`,
username,
)
reply(bot, update, text)
}

// findTrollHouses return a string with groups separeted by comma,
// that groups are well-known to being troll houses.
// otherwise, if nothing is found returns a empty string
func findTrollHouses(bot *tgbotapi.BotAPI, userID int) string {
func findTrollHouses(bot TrollShieldBot, userID int) string {
ch := make(chan string, len(trollGroups))
var wait sync.WaitGroup
for _, trollGroup := range trollGroups {
wait.Add(1)
go func(group string) {
defer wait.Done()
c, _ := bot.GetChatMember(tgbotapi.ChatConfigWithUser{
c, _ := bot.GetChatMember(telegram.ChatConfigWithUser{
SuperGroupUsername: group,
UserID: userID,
})
Expand All @@ -67,20 +135,29 @@ func findTrollHouses(bot *tgbotapi.BotAPI, userID int) string {
return strings.Join(houses, ", ")
}

// messageEvent: return true if is a message event
func messageEvent(update *tgbotapi.Update) bool {
return update.Message != nil
}

// newChatMemberEvent: return true if a new member joined to chat
func newChatMemberEvent(update *tgbotapi.Update) bool {
return messageEvent(update) && update.Message.NewChatMembers != nil
}

// fromChatEvent: return true if the message is from a specific chat
func fromChatEvent(update *tgbotapi.Update, username string) bool {
chat := update.Message.Chat
return messageEvent(update) && chat != nil && (chat.UserName == username || chat.Title == username)
// kickTroll ban the troll and send a message about where we can found the trolls
func kickTroll(bot *telegram.BotAPI, update *telegram.Update, user telegram.User, trollHouse string) {
chatMember := telegram.ChatMemberConfig{
ChatID: update.Message.Chat.ID,
UserID: user.ID,
}
resp, err := bot.KickChatMember(
telegram.KickChatMemberConfig{ChatMemberConfig: chatMember},
)

if resp.Ok == false || err != nil {
log.Printf(
"[!] Kicking %q did not work, error code %v: %v",
user.FirstName, resp.ErrorCode, resp.Description,
)
} else {
username := getUserName(user)
text := fmt.Sprintf(
"%v foi banido porque é membro do grupo: %v. Adeus.",
username, trollHouse,
)
reply(bot, update, text)
}
}

func setupLogging() {
Expand All @@ -93,43 +170,18 @@ func setupLogging() {

log.SetOutput(wrt)
// register log to BotLoggt
err = tgbotapi.SetLogger(log)
err = telegram.SetLogger(log)
if err != nil {
log.Printf("Set Telegram Bot Logging error: %v", err)
}
}

// getUserName: return the most meaningful name available from a telegram user
// if it has a @username, return it
// if not, try to return FirstName + LastName
// otherwise, only return the FirstName
func getUserName(user tgbotapi.User) string {
username := user.FirstName
if user.UserName != "" {
username = fmt.Sprintf("@%v", user.UserName)
} else if user.LastName != "" {
username = fmt.Sprintf("%s %s", user.FirstName, user.LastName)
}
return username
}

func reply(bot *tgbotapi.BotAPI, update *tgbotapi.Update, text string) {
msg := tgbotapi.NewMessage(update.Message.Chat.ID, text)
msg.ReplyToMessageID = update.Message.MessageID

_, err := bot.Send(msg)
if err != nil {
log.Printf("[!] Send msg failed: %v", err)
}
}

func main() {
setupLogging()
func setupBot() *telegram.BotAPI {
token, exists := os.LookupEnv("TELEGRAM_BOT_TOKEN")
if !exists {
log.Fatal("TELEGRAM_BOT_TOKEN env should be defined.")
}
bot, err := tgbotapi.NewBotAPI(token)
bot, err := telegram.NewBotAPI(token)

if err != nil {
log.Panic(err)
Expand All @@ -139,12 +191,13 @@ func main() {

log.Printf("Authorized on account @%s", bot.Self.UserName)

u := tgbotapi.NewUpdate(0)
u.Timeout = 60

updates, err := bot.GetUpdatesChan(u)
return bot
}

for update := range updates {
func main() {
setupLogging()
bot := setupBot()
for update := range getUpdates(bot) {
if messageEvent(&update) {
if update.Message.Text == "/lelerax" {
reply(bot, &update, "Estou vivo.")
Expand All @@ -154,37 +207,9 @@ func main() {
if newChatMemberEvent(&update) {
for _, member := range *update.Message.NewChatMembers {
if trollHouse := findTrollHouses(bot, member.ID); trollHouse != "" {
chatMember := tgbotapi.ChatMemberConfig{
ChatID: update.Message.Chat.ID,
UserID: member.ID,
}
resp, err := bot.KickChatMember(
tgbotapi.KickChatMemberConfig{ChatMemberConfig: chatMember},
)

if resp.Ok == false || err != nil {
log.Printf(
"[!] Kicking %q did not work, error code %v: %v",
member.FirstName, resp.ErrorCode, resp.Description,
)
} else {
username := getUserName(member)
text := fmt.Sprintf(
"%v foi banido porque é membro do grupo: %v. Adeus.",
username, trollHouse,
)
reply(bot, &update, text)
}
} else {
if fromChatEvent(&update, "commonlispbr") && !member.IsBot {
username := getUserName(member)
text := fmt.Sprintf(
`Olá %s! Seja bem-vindo ao grupo oficial de Common Lisp do Brasil.
Leia as regras em: https://lisp.com.br/rules.html.`,
username,
)
reply(bot, &update, text)
}
kickTroll(bot, &update, member, trollHouse)
} else if fromChatEvent(&update, "commonlispbr") && !member.IsBot {
welcomeMessage(bot, &update, member)
}
}
}
Expand Down

0 comments on commit 92bd96d

Please sign in to comment.