146 lines
6.7 KiB
Python
146 lines
6.7 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
Обработчик файлов: скачивание файлов по URL (из OOB), сохранение во временную директорию,
|
||
постановка задачи в очередь индексации.
|
||
|
||
Ключевое изменение для групповой работы:
|
||
- Если файл пришёл в групповой чат (room_jid не None), он автоматически индексируется
|
||
в базу знаний комнаты, без необходимости команды !learn.
|
||
- Для личных чатов – поведение зависит от режима обучения (learning_mode) и флага --global.
|
||
|
||
ДОБАВЛЕНО:
|
||
- Гарантированное удаление временных файлов через tempfile.TemporaryDirectory.
|
||
- Обработка ошибок скачивания с очисткой.
|
||
"""
|
||
|
||
import os
|
||
import asyncio
|
||
import logging
|
||
import urllib.parse
|
||
import requests
|
||
import tempfile
|
||
import shutil
|
||
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):
|
||
"""
|
||
Ожидает создания очереди индексации (защита от гонки при старте бота).
|
||
"""
|
||
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.
|
||
Использует временную директорию для гарантированной очистки.
|
||
"""
|
||
self.bot.metrics_handler.increment("files_processed")
|
||
try:
|
||
raw_name = url.split('/')[-1].split('?')[0]
|
||
original_fname = urllib.parse.unquote(raw_name)
|
||
ext = os.path.splitext(original_fname)[1].lower()
|
||
|
||
# Создаём временную директорию (автоматически удалится)
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
fname = f"{msg['from'].user}_{int(datetime.now().timestamp())}{ext}"
|
||
target_path = os.path.join(tmpdir, fname)
|
||
|
||
# Проверка размера через 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:
|
||
raise Exception("HTTP error")
|
||
limit = self.bot.config.max_file_size_mb * 1024 * 1024
|
||
with open(target_path, 'wb') as f:
|
||
downloaded = 0
|
||
for chunk in r.iter_content(8192):
|
||
f.write(chunk)
|
||
downloaded += len(chunk)
|
||
if downloaded > limit:
|
||
raise Exception("File too large")
|
||
|
||
try:
|
||
await asyncio.to_thread(download_with_limit)
|
||
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
|
||
|
||
# Определение режима индексации
|
||
if room_jid:
|
||
mode = 'room'
|
||
is_global = False
|
||
else:
|
||
mode_info = self.bot.learning_mode.get(user_jid)
|
||
if mode_info is None:
|
||
mode = None
|
||
is_global = False
|
||
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
|
||
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)
|
||
|
||
# Создаём задачу для фонового индексатора
|
||
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)
|
||
|
||
scope = "комнатную" if room_jid else ("глобальную" 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() |