bot.py 5.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181
  1. import os
  2. import logging
  3. from telegram import *
  4. from telegram.ext import *
  5. from ctftime_client import CTFTimeClient
  6. from dbs import CTFDb, GroupDb
  7. db = CTFDb()
  8. groups = GroupDb()
  9. TOKEN = os.environ['BOT_TOKEN']
  10. # Enable logging
  11. logging.basicConfig(format='%(asctime)s - %(name)s - \
  12. %(levelname)s - %(message)s',
  13. level=logging.INFO)
  14. logger = logging.getLogger(__name__)
  15. def error(bot, update, error):
  16. logger.warning('Update "%s" caused error "%s"' % (update, error))
  17. def start(bot, update):
  18. """Start the bot reminder about starting and announced CTF"""
  19. groups.add(update.message.chat_id)
  20. update.message.reply_text('Hi! I will notify you when new CTFs \
  21. are announced and remind them!')
  22. def stop(bot, update):
  23. """Start the bot reminder about starting and announced CTF"""
  24. groups.remove(update.message.chat_id)
  25. update.message.reply_text('Stopped')
  26. def fetch(bot, job):
  27. """Job function to check for newly announced CTFs"""
  28. to_print = db.add_events(CTFTimeClient.fetch_data())
  29. if to_print is []:
  30. return
  31. message = ""
  32. for ctf in to_print:
  33. message += db.starting_message(ctf)
  34. if message is not "":
  35. message = "*New CTF announced:*\n" + message
  36. for element in groups.groups:
  37. bot.sendMessage(element,
  38. text=message,
  39. parse_mode='MARKDOWN',
  40. disable_web_page_preview=True)
  41. def starting(bot, job):
  42. """Job function to check for newly announced CTFs"""
  43. to_print = db.starting_ctf()
  44. if to_print is []:
  45. return
  46. message = ""
  47. for ctf in to_print:
  48. message += db.starting_message(ctf)
  49. if message is not "":
  50. message = "*The following CTF are starting in less than 24 hours:*\n" + message
  51. for element in groups.groups:
  52. bot.sendMessage(element,
  53. text=message,
  54. parse_mode='MARKDOWN',
  55. disable_web_page_preview=True)
  56. def upcoming(bot, update):
  57. """List 5 upcoming CTFs in database"""
  58. message = ""
  59. for ctf in db.upcoming():
  60. message += db.starting_message(ctf)
  61. if message is not "":
  62. message = "*Upcoming events:*\n" + message
  63. bot.sendMessage(update.message.chat_id,
  64. text=message,
  65. parse_mode='MARKDOWN',
  66. disable_web_page_preview=True)
  67. else:
  68. bot.sendMessage(update.message.chat_id,
  69. text="No upcoming CTF")
  70. def current(bot, update):
  71. """List current running CTFs in database"""
  72. message = ""
  73. for ctf in db.running():
  74. message += db.finishing_message(ctf)
  75. if message is not "":
  76. message = "*Now running events:*\n" + message
  77. bot.sendMessage(update.message.chat_id,
  78. text=message,
  79. parse_mode='MARKDOWN',
  80. disable_web_page_preview=True)
  81. else:
  82. bot.sendMessage(update.message.chat_id,
  83. text="No CTF is currently running")
  84. def info(bot, update, args):
  85. """Get info form a given CTFid"""
  86. if len(args) != 1:
  87. update.message.reply_text('Usage: `/info <ctf_id>`',
  88. parse_mode='MARKDOWN')
  89. return
  90. update.message.reply_text(db.get_info_message(args[0]),
  91. parse_mode='MARKDOWN',
  92. disable_web_page_preview=True)
  93. def ping(bot, update):
  94. """Classic ping command"""
  95. message = "Pong. I'm not running in this group! Start me with `/start`"
  96. if update.message.chat_id in groups.groups:
  97. message = "Pong."
  98. update.message.reply_text(message, parse_mode='MARKDOWN')
  99. def usage(bot, update):
  100. message = "*CTF Reminder* will remind your CTF as they start!\n"
  101. message += "`/start` to start the reminder\n"
  102. message += "`/stop` to stop the reminder\n"
  103. message += "`/upcoming` to list all the upcoming CTFs\n"
  104. message += "`/now` or `/current` to list all the currently running CTFs\n"
  105. message += "`/info <ctf_id>` to get info for specific CTF\n"
  106. message += "`/ping` to check if the reminder is started\n"
  107. message += "\nGet the bot [Source Code]"
  108. message += "(https://github.com/TheZ3ro/ctf_reminder)"
  109. update.message.reply_text(message, parse_mode='MARKDOWN',
  110. disable_web_page_preview=True)
  111. def main():
  112. db.add_events(CTFTimeClient.fetch_data())
  113. updater, bot = Updater(TOKEN), Bot(TOKEN)
  114. # Get the dispatcher to register handlers
  115. dp = updater.dispatcher
  116. # on different commands - answer in Telegram
  117. dp.add_handler(CommandHandler("start", start))
  118. dp.add_handler(CommandHandler("stop", stop))
  119. dp.add_handler(CommandHandler("upcoming", upcoming))
  120. dp.add_handler(CommandHandler("now", current))
  121. dp.add_handler(CommandHandler("current", current))
  122. dp.add_handler(CommandHandler("info", info, pass_args=True))
  123. dp.add_handler(CommandHandler("ping", ping))
  124. dp.add_handler(CommandHandler("help", usage))
  125. # log all errors
  126. dp.add_error_handler(error)
  127. # Start the Bot
  128. updater.start_polling(clean=True)
  129. job_queue = updater.job_queue
  130. job_queue.run_repeating(fetch, 6 * 60 * 60)
  131. job_queue.run_repeating(starting, 1 * 60 * 60)
  132. # Block until you press Ctrl-C or the process receives SIGINT, SIGTERM or
  133. # SIGABRT. This should be used most of the time, since start_polling() is
  134. # non-blocking and will stop the bot gracefully.
  135. updater.idle()
  136. if __name__ == '__main__':
  137. main()