137 lines
4.1 KiB
Python
Executable File
137 lines
4.1 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
import logging
|
|
import os
|
|
import shutil
|
|
import sys
|
|
from glob import glob
|
|
|
|
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):
|
|
if 'api.telegram.org' in url \
|
|
and ('animation' in url or 'video' in url) \
|
|
and "mp4" not in url:
|
|
ext = '.mp4'
|
|
else:
|
|
ext = ''
|
|
local_filename = OUT_DIR + '/' + datestr(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)
|
|
|
|
|
|
# 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()
|