From b7929bf6a7dbfc3665a9e0225ac5fede3accecfd Mon Sep 17 00:00:00 2001
From: bakatrouble <bakatrouble@gmail.com>
Date: Mon, 12 Apr 2021 14:47:33 +0300
Subject: [PATCH] add qr recognition to robot

---
 Pipfile               |  1 +
 bots/modules/robot.py | 52 ++++++++++++++++++++++++++++++++++++++++---
 2 files changed, 50 insertions(+), 3 deletions(-)

diff --git a/Pipfile b/Pipfile
index 73d186f..5daec85 100644
--- a/Pipfile
+++ b/Pipfile
@@ -39,6 +39,7 @@ whitenoise = "*"
 django-picklefield = "*"
 boto3 = "*"
 filetype = "*"
+qrtools = "*"
 
 [dev-packages]
 
diff --git a/bots/modules/robot.py b/bots/modules/robot.py
index 7e06ba6..43fad28 100644
--- a/bots/modules/robot.py
+++ b/bots/modules/robot.py
@@ -1,12 +1,14 @@
-import logging
 import os
+import re
 import time
 from tempfile import TemporaryDirectory
 
 import boto3
 import filetype
+import twitter
 from django.db import models
-from telegram import Update, InlineKeyboardMarkup, InlineKeyboardButton
+from qrtools import QR
+from telegram import Update, InlineKeyboardMarkup, InlineKeyboardButton, MessageEntity
 from telegram.ext import CallbackContext, Dispatcher, MessageHandler, Filters
 from djconfig import config
 
@@ -21,6 +23,9 @@ class RobotBotModuleConfig(TelegramBotModuleConfig):
     def build_dispatcher(self, dispatcher: Dispatcher):
         dispatcher.add_handler(MessageHandler(Filters.audio | Filters.photo | Filters.video | Filters.document,
                                               self.save_file))
+        dispatcher.add_handler(MessageHandler(Filters.text & (Filters.entity(MessageEntity.URL) |
+                                                              Filters.entity(MessageEntity.TEXT_LINK)),
+                                              self.document_downloader))
         return dispatcher
 
     def save_file(self, update: Update, ctx: CallbackContext):
@@ -45,17 +50,58 @@ class RobotBotModuleConfig(TelegramBotModuleConfig):
 
         file = ctx.bot.get_file(file_id)
         filename = str(int(time.time())) + '-' + os.path.basename(file.file_path)
+        text = ''
         with TemporaryDirectory() as d:
             file_path = os.path.join(d, filename)
             file.download(file_path)
             tp = filetype.guess(file_path)
             s3 = boto3.client('s3', aws_access_key_id=config.aws_access_key_id, aws_secret_access_key=config.aws_secret_access_key)
             s3.upload_file(file_path, config.s3_bucket, filename, ExtraArgs={'ACL': 'public-read', 'ContentType': tp.mime})
+            if update.message.text == 'qr':
+                qr = QR(filename=file_path)
+                qr.decode()
+                text = f'QR:\n{qr.data}'
         file_url = f'https://{config.s3_bucket}.s3.amazonaws.com/{filename}'
+        text += f'\n\n{file_url}'
         if update.message.photo:
             markup.inline_keyboard = [[
                 InlineKeyboardButton('Google IS', f'https://www.google.com/searchbyimage?image_url={file_url}'),
                 InlineKeyboardButton('Yandex IS', f'https://yandex.ru/images/search?rpt=imageview&url={file_url}'),
                 InlineKeyboardButton('SauceNao', f'https://saucenao.com/search.php?url={file_url}')
             ]]
-        update.message.reply_text(file_url, disable_web_page_preview=True, reply_markup=markup)
+        update.message.reply_text(text, disable_web_page_preview=True, reply_markup=markup)
+
+    def document_downloader(self, update: Update, ctx: CallbackContext):
+        config._reload_maybe()
+
+        if self.users.count() and not self.users.filter(user_id=update.effective_user.id).count():
+            update.effective_message.reply_text('GTFO')
+            return
+
+        links = update.effective_message.parse_entities([MessageEntity.URL, MessageEntity.TEXT_LINK]).values()
+        for link in links:
+            try:
+                twitter_matches = re.match(r'https?://twitter.com/[A-Za-z0-9_]{1,15}/status/(\d+).*', link)
+                if link.endswith('.gif'):
+                    update.message.reply_document(link)
+                elif twitter_matches:
+                    api = twitter.Api(config.twitter_consumer_api_key, config.twitter_consumer_api_secret,
+                                      config.twitter_access_token, config.twitter_access_token_secret)
+                    tweet = api.GetStatus(twitter_matches.group(1))
+                    try:
+                        for media in tweet_info['extended_entities']['media']:
+                            if media['type'] == 'video':
+                                max_bitrate = -1
+                                max_bitrate_link = None
+                                for variant in media['video_info']['variants']:
+                                    if variant['content_type'] == 'video/mp4' and variant['bitrate'] > max_bitrate:
+                                        max_bitrate = variant['bitrate']
+                                        max_bitrate_link = variant['url']
+                                if max_bitrate_link is not None:
+                                    update.message.reply_video(max_bitrate_link)
+                    except KeyError:
+                        update.message.reply_text('Unable to retrieve video info from tweet')
+                else:
+                    update.message.reply_text('Can\'t detect link type')
+            except BadRequest:
+                update.message.reply_text('Unable to retrieve file')