tried putting in a class

This commit is contained in:
Stan James 2018-02-01 19:35:56 -07:00
parent 0079ae2aec
commit 20a70ff55f

86
bot.py
View File

@ -18,34 +18,47 @@ from time import strftime
import re import re
import unidecode import unidecode
# Define a few command handlers. These usually take the two arguments bot and
# update. Error handlers also receive the raised TelegramError object in error.
class TelegramMonitorBot:
def security_check_message(bot, update): def __init__(self):
# '[0-9a-fA-F]{40,40}',
self.message_ban_patterns = os.environ['MESSAGE_BAN_PATTERNS']
self.message_ban_re = re.compile(self.message_ban_patterns, re.IGNORECASE)
self.message_hide_patterns = os.environ['MESSAGE_HIDE_PATTERNS']
self.message_hide_re = re.compile(self.message_hide_patterns, re.IGNORECASE)
# # Define a few command handlers. These usually take the two arguments bot and
# # update. Error handlers also receive the raised TelegramError object in error.
# def security_check_username(bot, update):
# """ Test username for security violations """
# # ban_patterns = [
# # 'origin',
# # 'admin',
# # 'official',
# # ]
# pass
def security_check_message(self, bot, update):
""" Test message for security violations """ """ Test message for security violations """
ban_patterns = [ # Remove accents from letters (é->e, ñ->n, etc...)
'[0-9a-fA-F]{40,40}',
'Fart',
]
# Remove accents from letters
message = unidecode.unidecode(update.message.text) message = unidecode.unidecode(update.message.text)
regexp_pattern = "|".join(ban_patterns) if self.message_hide_re.search(message):
r = re.compile(regexp_pattern, re.IGNORECASE)
if r.search(message):
print("Bannable match: {}".format(update.message.text.encode('utf-8')))
# Delete the message # Delete the message
print("Hide match: {}".format(update.message.text.encode('utf-8')))
update.message.delete() update.message.delete()
# # Ban the user if self.message_ban_re.search(message):
# kick_success = update.message.chat.kick_member(update.message.from_user.id) # Ban the user
# print(kick_success) kick_success = update.message.chat.kick_member(update.message.from_user.id)
print(kick_success)
print("Ban match: {}".format(update.message.text.encode('utf-8')))
def logger(bot, update): def logger(self, bot, update):
"""Primary Logger. Handles incoming bot messages and saves them to DB""" """Primary Logger. Handles incoming bot messages and saves them to DB"""
user = update.message.from_user user = update.message.from_user
@ -70,13 +83,14 @@ def logger(bot, update):
) )
try: try:
security_check_message(bot, update) self.security_check_username(bot, update)
self.security_check_message(bot, update)
except Exception as e: except Exception as e:
print(e) print(e)
# DB queries # DB queries
def id_exists(id_value): def id_exists(self, id_value):
s = session() s = session()
bool_set = False bool_set = False
for id1 in s.query(User.id).filter_by(id=id_value): for id1 in s.query(User.id).filter_by(id=id_value):
@ -88,7 +102,7 @@ def id_exists(id_value):
return bool_set return bool_set
def log_message(user_id, user_message): def log_message(self, user_id, user_message):
try: try:
s = session() s = session()
@ -101,31 +115,31 @@ def log_message(user_id, user_message):
print(e) print(e)
def add_user(user_id, first_name, last_name, username): def add_user(self, user_id, first_name, last_name, username):
try: try:
s = session() s = session()
bool_set = False user = User(
user = User(id=user_id, first_name = first_name, last_name = last_name, username = username) id=user_id,
first_name=first_name,
last_name=last_name,
username=username)
s.add(user) s.add(user)
s.commit() s.commit()
s.close() s.close()
return id_exists(user_id)
if id_exists(user_id) == True:
bool_set = True
return bool_set
except Exception as e: except Exception as e:
print(e) print(e)
def error(bot, update, error): def error(self, bot, update, error):
"""Log Errors caused by Updates.""" """Log Errors caused by Updates."""
print("Update '{}' caused error '{}'".format(update, error), file=sys.stderr) print("Update '{}' caused error '{}'".format(update, error),
file=sys.stderr)
def main(): def start(self):
"""Start the bot.""" """Start the bot."""
# Create the EventHandler and pass it your bot's token. # Create the EventHandler and pass it your bot's token.
updater = Updater(os.environ["TELEGRAM_BOT_TOKEN"]) updater = Updater(os.environ["TELEGRAM_BOT_TOKEN"])
@ -135,12 +149,12 @@ def main():
# on different commands - answer in Telegram # on different commands - answer in Telegram
# on noncommand i.e message - echo the message on Telegram # on noncommand i.e message - echo the message on Telegram
dp.add_handler(MessageHandler(Filters.text, logger)) dp.add_handler(MessageHandler(Filters.text, self.logger))
# dp.add_handler(MessageHandler(Filters.status_update, status)) # dp.add_handler(MessageHandler(Filters.status_update, status))
# log all errors # log all errors
dp.add_error_handler(error) dp.add_error_handler(self.error)
# Start the Bot # Start the Bot
updater.start_polling() updater.start_polling()
@ -154,4 +168,6 @@ def main():
if __name__ == '__main__': if __name__ == '__main__':
main() c = TelegramMonitorBot()
c.start()