# -*- coding: utf-8 -*- """ Обработчик файлов: скачивание файлов по URL (из OOB), сохранение во временную директорию, постановка задачи в очередь индексации. Ключевое изменение для групповой работы: - Если файл пришёл в групповой чат (room_jid не None), он автоматически индексируется в базу знаний комнаты, без необходимости команды !learn. - Для личных чатов – поведение зависит от режима обучения (learning_mode) и флага --global. """ import os import asyncio import logging import urllib.parse import requests from datetime import datetime import re logger = logging.getLogger(__name__) class FileHandler: """ Отвечает за скачивание файлов, переданных через ссылку в сообщении (XEP-0066 OOB). """ def __init__(self, bot): self.bot = bot async def _wait_for_queue(self): """ Ожидает создания очереди индексации (защита от гонки при старте бота). Проблема: очередь `self.bot.indexing_queue` создаётся в `session_start` после подключения к БД. Если сообщение с файлом придёт ДО этого момента, `indexing_queue` будет `None`, и `put()` вызовет исключение. Решение: ждём, пока очередь станет доступна, проверяя каждые 0.5 секунды. """ while self.bot.indexing_queue is None: logger.debug("Очередь индексации ещё не создана, ожидание...") await asyncio.sleep(0.5) async def handle_file(self, msg, url: str, user_jid: str, room_jid: str = None): """ Скачивает файл и создаёт задачу для indexing_worker. Аргументы: msg: stanza сообщения (нужна для получения user, ответных сообщений). url (str): HTTP/HTTPS ссылка на файл (из атрибута url элемента oob). user_jid (str): JID отправителя (без ресурса). room_jid (str | None): JID комнаты, если сообщение из groupchat; иначе None. Логика: 1. Извлечь имя файла из URL, декодировать, сгенерировать уникальное имя во временной папке. 2. Скачать файл (синхронно, в отдельном потоке, с таймаутом из конфига). 3. Определить режим индексации: - Если room_jid не None -> mode = 'room', is_global = False (всегда). - Если личный чат -> проверить learning_mode пользователя. 4. Извлечь --title из тела сообщения (если есть). 5. Положить задачу в очередь индексации (asyncio.Queue). 6. Ответить пользователю, что файл принят в обработку. """ self.bot.metrics_handler.increment("files_processed") try: # --- 1. Извлечение имени файла из URL --- raw_name = url.split('/')[-1].split('?')[0] original_fname = urllib.parse.unquote(raw_name) ext = os.path.splitext(original_fname)[1].lower() fname = f"{msg['from'].user}_{int(datetime.now().timestamp())}{ext}" target_path = os.path.join(self.bot.config.temp_dir, fname) # --- 2. Скачивание файла --- # HEAD-запрос для получения размера try: head_resp = requests.head(url, timeout=self.bot.config.http_download_timeout) content_length = head_resp.headers.get('Content-Length') if content_length is not None and int(content_length) > self.bot.config.max_file_size_mb * 1024 * 1024: reply = msg.reply(f"❌ Файл превышает максимальный размер {self.bot.config.max_file_size_mb} МБ.") if reply: reply.send() return except Exception: pass # Скачивание с прерыванием по лимиту def download_with_limit(): r = requests.get(url, stream=True, timeout=self.bot.config.http_download_timeout) if r.status_code != 200: return False downloaded = 0 limit = self.bot.config.max_file_size_mb * 1024 * 1024 with open(target_path, 'wb') as f: for chunk in r.iter_content(8192): f.write(chunk) downloaded += len(chunk) if downloaded > limit: os.unlink(target_path) raise Exception("File too large") return True try: success = await asyncio.to_thread(download_with_limit) if not success: reply = msg.reply("❌ Не удалось загрузить файл. Проверьте доступность ссылки.") if reply: reply.send() return except Exception as e: if "File too large" in str(e): reply = msg.reply(f"❌ Файл превышает максимальный размер {self.bot.config.max_file_size_mb} МБ.") else: reply = msg.reply("❌ Ошибка при скачивании файла.") if reply: reply.send() return # --- 3. Определение режима индексации --- if room_jid: mode = 'room' is_global = False title = None else: mode_info = self.bot.learning_mode.get(user_jid) if mode_info is None: mode = None is_global = False title = None else: if isinstance(mode_info, tuple): mode, saved_room = mode_info is_global = (mode == 'global') else: mode = mode_info is_global = (mode == 'global') title = None # --- 4. Парсинг --title из тела сообщения --- raw_body = msg['body'].strip() if msg['body'] else "" if raw_body: match = re.search(r'--title "([^"]+)"', raw_body) if match: title = match.group(1) # --- 5. Создание задачи для очереди --- await self._wait_for_queue() task = { "user_jid": user_jid, "room_jid": room_jid, "file_path": target_path, "original_fname": original_fname, "title": title, "is_global": is_global, "mode": mode, "ext": ext } await self.bot.indexing_queue.put(task) # --- 6. Уведомление пользователя --- if room_jid: reply = msg.reply(f"📥 Файл '{original_fname}' принят в обработку для КОМНАТНОЙ БЗ.") else: scope = "глобальную" if is_global else "личную" reply = msg.reply(f"📥 Файл '{original_fname}' принят в обработку для {scope} БЗ.") if reply: reply.send() except Exception as e: logger.exception(f"Ошибка при скачивании файла: {e}") self.bot.metrics_handler.increment("errors_total") reply = msg.reply("❌ Ошибка при обработке файла. Проверьте логи.") if reply: reply.send()