From f9342b164d2df51024ece07b84b63777988c7c5f Mon Sep 17 00:00:00 2001
From: Aaron Johnson <ajohnson@skyfalltech.net>
Date: Fri, 28 Dec 2018 00:42:17 -0600
Subject: [PATCH] Chasing a bug (bot crashes after being added to group)

---
 bot_config.yaml.example |  4 ++++
 commands.rb             |  7 +++++++
 run.rb                  | 19 ++++++++++++++-----
 3 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/bot_config.yaml.example b/bot_config.yaml.example
index c43d20a..105d6cd 100644
--- a/bot_config.yaml.example
+++ b/bot_config.yaml.example
@@ -8,4 +8,8 @@ admin:
     - 11111111
 authorized_chats:
     - 11111111
+allowed_sources:
+    - 'private'
+    - 'group'
+    - 'supergroup'
 
diff --git a/commands.rb b/commands.rb
index ee7dfaf..764b47c 100644
--- a/commands.rb
+++ b/commands.rb
@@ -1,3 +1,10 @@
+def validate_incoming_data(message)
+    message = message.message if message.is_a? Telegram::Bot::Types::CallbackQuery
+    return "Received message is not from a valid source! Type: \"#{message.chat.type}\". Ignoring." if ! @allowed_sources.include?(message.chat.type)
+    puts message.chat.type
+    return true
+end
+
 def message_from_admin?(message, adm)
     #return @admin_userids.include? message.from.id.to_s
     #puts adm.class
diff --git a/run.rb b/run.rb
index dddfe52..1bb873f 100755
--- a/run.rb
+++ b/run.rb
@@ -13,6 +13,7 @@ token = @conf['token']
 telnet = @conf['telnet']
 admin = @conf['admin']
 @auth_chat = @conf['authorized_chats']
+@allowed_sources = @conf['allowed_sources']
 
 # Sanity check
 errcount = 0
@@ -72,12 +73,12 @@ def handle_message(message)
 
     #Format command
     command = message.text.split(" ")[0].split("@")[0].downcase #Strip command from arguments and @tags
+    #command = message.text
 
     reply = 'Empty String'
     #conf = YAML.load_file("bot_config.yaml")
     telnet = @conf['telnet']
     adm = @conf['admin']
-    #case message.text.split(" ")[0].split("@")[0].downcase #Strip command from arguments and @tags
     puts "Received command from #{message.from.username}: #{command}"
     case command
         when '/start'
@@ -142,11 +143,19 @@ end
 
 Telegram::Bot::Client.run(token) do |bot|
     bot.listen do |message|
-        reply = handle_message(message)
-        if reply == "dance" || reply == "flex"
-            bot.api.sendVideo(chat_id: message.chat.id, video: "https://img.skyfalltech.net/togra/#{reply}.gif")
+        #pp message
+        validation = validate_incoming_data(message)
+        #puts "DEBUG: #{validation}"
+        if validation
+            reply = handle_message(message)
+            if reply == "dance" || reply == "flex"
+                bot.api.sendVideo(chat_id: message.chat.id, video: "https://img.skyfalltech.net/togra/#{reply}.gif")
+            else
+                bot.api.send_message(chat_id: message.chat.id, text: "#{reply}")
+            end
         else
-            bot.api.send_message(chat_id: message.chat.id, text: "#{reply}")
+            puts "Received bad data! [#{message.chat.type}]"
+            puts validation
         end
     end
 end