main.py 1.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657
  1. import logging
  2. import os
  3. from telegram import InlineKeyboardMarkup
  4. from telegram.ext import Updater, CallbackQueryHandler, MessageHandler, Filters
  5. from vid_utils import Video, BadLink
  6. logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', level=logging.INFO)
  7. logger = logging.getLogger(__name__)
  8. def get_format(update, context):
  9. logger.info("from {}: {}".format(update.message.chat_id, update.message.text)) # "history"
  10. try:
  11. video = Video(link=update.message.text, init_keyboard=True)
  12. except BadLink as e:
  13. update.message.reply_text("Bad link: {}".format(e))
  14. except Exception as e:
  15. update.message.reply_text("Error: {}".format(e))
  16. else:
  17. reply_markup = InlineKeyboardMarkup(video.keyboard)
  18. update.message.reply_text('Choose format:', reply_markup=reply_markup)
  19. def download_choosen_format(update, context):
  20. query = update.callback_query
  21. video_index = query.data
  22. context.bot.edit_message_text(text="Downloading...",
  23. chat_id=query.message.chat_id,
  24. message_id=query.message.message_id)
  25. video = Video(vid=video_index)
  26. video.download()
  27. with video.send() as files:
  28. for f in files:
  29. logger.info("Sending... {} ".format(f))
  30. context.bot.send_document(chat_id=query.message.chat_id, document=open(f, 'rb'))
  31. context.bot.delete_message(chat_id=query.message.chat_id,
  32. message_id=query.message.message_id)
  33. TOKEN = None
  34. with open(os.path.join(os.environ['CONF_FOLDER'], 'bot.token')) as f:
  35. TOKEN = f.read().strip()
  36. updater = Updater(token=TOKEN)
  37. updater.dispatcher.add_handler(MessageHandler(Filters.text, get_format))
  38. updater.dispatcher.add_handler(CallbackQueryHandler(download_choosen_format))
  39. updater.start_polling()
  40. updater.idle()