Add checks

This commit is contained in:
Ceda EI 2018-10-03 19:56:07 +05:30
parent 71f798fe5e
commit b2b20e7f05
1 changed files with 13 additions and 3 deletions

16
bot.py
View File

@ -17,14 +17,20 @@ logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - \
%(message)s', level=logging.INFO) %(message)s', level=logging.INFO)
def start(bot, update): def check_user(bot, update):
chat_id = update.message.chat_id chat_id = update.message.chat_id
if chat_id != config.master_id: if chat_id != config.master_id:
text = "You need to host your own instance! \ Source: \ text = "You need to host your own instance! Source: \
https://gitlab.com/ceda_ei/pc-monitor-bot" https://gitlab.com/ceda_ei/pc-monitor-bot"
bot.send_message(chat_id=chat_id, text=text) bot.send_message(chat_id=chat_id, text=text)
return return False
return True
def start(bot, update):
if not check_user(bot, update):
return False
chat_id = update.message.chat_id
custom_keyboard = [['/state', '/screenshot'], ['/lock', '/video']] custom_keyboard = [['/state', '/screenshot'], ['/lock', '/video']]
reply_markup = telegram.ReplyKeyboardMarkup(custom_keyboard) reply_markup = telegram.ReplyKeyboardMarkup(custom_keyboard)
bot.send_message(chat_id=chat_id, text="What to do?", bot.send_message(chat_id=chat_id, text="What to do?",
@ -40,6 +46,8 @@ def get_screenshot():
def screenshot(bot, update): def screenshot(bot, update):
if not check_user(bot, update):
return False
screenshot = get_screenshot() screenshot = get_screenshot()
chat_id = update.message.chat_id chat_id = update.message.chat_id
if not screenshot: if not screenshot:
@ -50,6 +58,8 @@ def screenshot(bot, update):
def lock(bot, update): def lock(bot, update):
if not check_user(bot, update):
return False
command = config.lock.split() command = config.lock.split()
exitcode = subprocess.call(command) exitcode = subprocess.call(command)
chat_id = update.message.chat_id chat_id = update.message.chat_id