#!/usr/bin/env python3 import logging import os import re import shutil import sys from glob import glob import filetype import requests import youtube_dl from telegram import MessageEntity from telegram.ext import Updater, CommandHandler, MessageHandler, Filters logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') logger = logging.getLogger("kunsax") DIR = os.path.dirname(os.path.realpath(__file__)) TMP_DIR = '/tmp' OUT_DIR = DIR + '/out' def ytdl_has(url): ies = youtube_dl.extractor.gen_extractors() for ie in ies: if ie.suitable(url) and ie.IE_NAME != 'generic' \ and '/channel/' not in url: # Site has dedicated extractor return True return False def datestr(date): return date.strftime("%Y-%m-%d@%H%M") def download_ydl(urls, date): ydl_opts = { 'noplaylist': True, 'restrictfilenames': True, 'download_archive': DIR + '/downloaded.lst', 'outtmpl': TMP_DIR + '/' + datestr(date) + '__%(title)s__%(id)s.%(ext)s' } with youtube_dl.YoutubeDL(ydl_opts) as ydl: ydl.download(urls) for filename in map(ydl.prepare_filename, map(ydl.extract_info, urls)): globbeds = glob(os.path.splitext(filename)[0] + '.*') for globbed in globbeds: logger.info("Moving %s to %s..." % (globbed, OUT_DIR)) shutil.move(globbed, OUT_DIR) def download_raw(url, date): local_filename = OUT_DIR + '/' + datestr(date) + '__' + url.split('/')[-1] r = requests.get(url, stream=True) with open(local_filename, 'wb') as f: for chunk in r.iter_content(chunk_size=1024): if chunk: f.write(chunk) if not re.match(r'.*\..{3,5}$', os.path.split(local_filename)[-1]): kind = filetype.guess(local_filename) if kind is None: logger.error("File has no extension and could not be determined!") else: logger.info('Moving file without extension... %s?' % kind.extension) shutil.move(local_filename, local_filename + '.' + kind.extension) # noinspection PyBroadException def handle_url(bot, update): ytdl_urls = list(filter(ytdl_has, map(lambda e: update.message.text[e.offset:e.offset + e.length], filter(lambda e: e.type == 'url', update.message.entities)))) if len(ytdl_urls) > 0: try: logger.info("Downloading %s" % ytdl_urls) update.message.reply_text('Downloading now...') download_ydl(ytdl_urls, update.message.date) except: type, value, _ = sys.exc_info() update.message.reply_text("Something is FUCKED: %s, %s" % (type, value)) # noinspection PyBroadException def handle_rest(bot, update): file = None if len(update.message.photo) > 0: photo = max(update.message.photo, key=lambda p: p.width) file = photo.file_id elif update.message.document is not None: file = update.message.document.file_id elif update.message.audio is not None: file = update.message.audio.file_id elif update.message.video is not None: file = update.message.video.file_id if file is not None: try: url = bot.getFile(file).file_path update.message.reply_text('Downloading now...') logger.info("Downloading '%s'" % url) download_raw(url, update.message.date) except: type, value, _ = sys.exc_info() update.message.reply_text("Something is FUCKED: %s, %s" % (type, value)) def start(bot, update): update.message.reply_text('WOOP WOOP') def error(bot, update, error): logger.error(error) if update is not None: update.message.reply_text("Something is fucked: %s" % error) def main(): updater = Updater("***REMOVED***") dp = updater.dispatcher dp.add_handler(CommandHandler("start", start)) dp.add_error_handler(error) dp.add_handler(MessageHandler(Filters.entity(MessageEntity.URL), handle_url)) dp.add_handler(MessageHandler(Filters.photo | Filters.video | Filters.audio | Filters.document, handle_rest)) updater.start_polling() logger.info("Started Telegram bot...") updater.idle() if __name__ == '__main__': main()