add robot module

2021-03-15 21:59:34 +03:00
parent ddb82c4269
commit 6ca5ff666e
10 changed files with 162 additions and 22 deletions
--- a/bots/modules/robot.py
+++ b/bots/modules/robot.py
@@ -0,0 +1,55 @@
+import os
+from tempfile import TemporaryDirectory
+
+import boto3
+from django.db import models
+from telegram import Update, InlineKeyboardMarkup, InlineKeyboardButton
+from telegram.ext import CallbackContext, Dispatcher, MessageHandler, Filters
+from djconfig import config
+
+from bots.models import TelegramBotModuleConfig, BotUser
+
+
+class RobotBotModuleConfig(TelegramBotModuleConfig):
+    users = models.ManyToManyField(BotUser)
+
+    MODULE_NAME = 'RoBOT'
+
+    def build_dispatcher(self, dispatcher: Dispatcher):
+        dispatcher.add_handler(MessageHandler(Filters.audio | Filters.photo | Filters.video | Filters.document,
+                                              self.save_file))
+        return dispatcher
+
+    def save_file(self, update: Update, ctx: CallbackContext):
+        if self.users.count() and not self.users.filter(user_id=update.effective_user.id).count():
+            update.effective_message.reply_text('GTFO')
+            return
+
+        markup = InlineKeyboardMarkup([])
+
+        if update.message.photo:
+            file_id = update.message.photo[-1].file_id
+        elif update.message.audio:
+            file_id = update.message.audio.file_id
+        elif update.message.video:
+            file_id = update.message.video.file_id
+        elif update.message.document:
+            file_id = update.message.document.file_id
+        else:
+            return
+
+        file = ctx.bot.get_file(file_id)
+        filename = os.path.basename(file.file_path)
+        with TemporaryDirectory() as d:
+            file_path = os.path.join(d, filename)
+            file.download(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)
+        file_url = f'https://{config.s3_bucket}.s3.amazonaws.com/{filename}'
+        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)