add qr recognition to robot

This commit is contained in:
bakatrouble 2021-04-12 14:47:33 +03:00
parent e72cb9a453
commit b7929bf6a7
2 changed files with 50 additions and 3 deletions

View File

@ -39,6 +39,7 @@ whitenoise = "*"
django-picklefield = "*" django-picklefield = "*"
boto3 = "*" boto3 = "*"
filetype = "*" filetype = "*"
qrtools = "*"
[dev-packages] [dev-packages]

View File

@ -1,12 +1,14 @@
import logging
import os import os
import re
import time import time
from tempfile import TemporaryDirectory from tempfile import TemporaryDirectory
import boto3 import boto3
import filetype import filetype
import twitter
from django.db import models 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 telegram.ext import CallbackContext, Dispatcher, MessageHandler, Filters
from djconfig import config from djconfig import config
@ -21,6 +23,9 @@ class RobotBotModuleConfig(TelegramBotModuleConfig):
def build_dispatcher(self, dispatcher: Dispatcher): def build_dispatcher(self, dispatcher: Dispatcher):
dispatcher.add_handler(MessageHandler(Filters.audio | Filters.photo | Filters.video | Filters.document, dispatcher.add_handler(MessageHandler(Filters.audio | Filters.photo | Filters.video | Filters.document,
self.save_file)) self.save_file))
dispatcher.add_handler(MessageHandler(Filters.text & (Filters.entity(MessageEntity.URL) |
Filters.entity(MessageEntity.TEXT_LINK)),
self.document_downloader))
return dispatcher return dispatcher
def save_file(self, update: Update, ctx: CallbackContext): def save_file(self, update: Update, ctx: CallbackContext):
@ -45,17 +50,58 @@ class RobotBotModuleConfig(TelegramBotModuleConfig):
file = ctx.bot.get_file(file_id) file = ctx.bot.get_file(file_id)
filename = str(int(time.time())) + '-' + os.path.basename(file.file_path) filename = str(int(time.time())) + '-' + os.path.basename(file.file_path)
text = ''
with TemporaryDirectory() as d: with TemporaryDirectory() as d:
file_path = os.path.join(d, filename) file_path = os.path.join(d, filename)
file.download(file_path) file.download(file_path)
tp = filetype.guess(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 = 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}) 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}' file_url = f'https://{config.s3_bucket}.s3.amazonaws.com/{filename}'
text += f'\n\n{file_url}'
if update.message.photo: if update.message.photo:
markup.inline_keyboard = [[ markup.inline_keyboard = [[
InlineKeyboardButton('Google IS', f'https://www.google.com/searchbyimage?image_url={file_url}'), 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('Yandex IS', f'https://yandex.ru/images/search?rpt=imageview&url={file_url}'),
InlineKeyboardButton('SauceNao', f'https://saucenao.com/search.php?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')