delojza/robot.py

import logging
import os
from datetime import datetime

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__))
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 date():
    return datetime.now().strftime("%Y-%m-%d@%H:%M")


def download_ydl(urls):
    ydl_opts = {
        'noplaylist': True,
        'download_archive': DIR + '/downloaded.lst',
        'outtmpl': OUT_DIR + '/' + date() + '__%(title)s__%(id)s.%(ext)s'
    }
    with youtube_dl.YoutubeDL(ydl_opts) as ydl:
        ydl.download(urls)


def download_raw(url):
    ext = ''
    if 'api.telegram.org' in url and 'animation' in url:
        ext = '.mp4'
    local_filename = OUT_DIR + '/' + date() + '__' + url.split('/')[-1] + ext
    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)


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:
        logger.info("Downloading %s" % ytdl_urls)
        update.message.reply_text('Downloading now...')
        download_ydl(ytdl_urls)


def handle_photo(bot, update):
    photo = max(update.message.photo, key=lambda p: p.width)
    url = bot.getFile(photo.file_id).file_path
    update.message.reply_text('Downloading now...')
    logger.info("Downloading '%s'" % url)
    download_raw(url)


def handle_audio(bot, update):
    url = bot.getFile(update.message.audio.file_id).file_path
    update.message.reply_text('Downloading now...')
    logger.info("Downloading '%s'" % url)
    download_raw(url)


def handle_video(bot, update):
    url = bot.getFile(update.message.video.file_id).file_path
    update.message.reply_text('Downloading now...')
    logger.info("Downloading '%s'" % url)
    download_raw(url)


def handle_document(bot, update):
    url = bot.getFile(update.message.document.file_id).file_path
    update.message.reply_text('Downloading now...')
    logger.info("Downloading '%s'" % url)
    download_raw(url)


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, handle_photo))
    dp.add_handler(MessageHandler(Filters.audio, handle_audio))
    dp.add_handler(MessageHandler(Filters.video, handle_video))
    dp.add_handler(MessageHandler(Filters.document, handle_document))

    updater.start_polling()

    logger.info("Started Telegram bot...")

    updater.idle()


if __name__ == '__main__':
    main()