Initial commit: Telegram CLI Bot with AI memory, tools and multi-provider support
This commit is contained in:
mirivlad
@@ -0,0 +1,22 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Утилиты бота."""
|
||||
|
||||
from bot.utils.cleaners import clean_ansi_codes, normalize_output
|
||||
from bot.utils.formatters import escape_html, escape_markdown, split_message, send_long_message, format_long_output, MAX_MESSAGE_LENGTH
|
||||
from bot.utils.decorators import check_access
|
||||
from bot.utils.ssh_readers import detect_input_type, read_ssh_output, read_pty_output
|
||||
|
||||
__all__ = [
|
||||
"clean_ansi_codes",
|
||||
"normalize_output",
|
||||
"escape_html",
|
||||
"escape_markdown",
|
||||
"split_message",
|
||||
"send_long_message",
|
||||
"format_long_output",
|
||||
"MAX_MESSAGE_LENGTH",
|
||||
"check_access",
|
||||
"detect_input_type",
|
||||
"read_ssh_output",
|
||||
"read_pty_output",
|
||||
]
|
||||
@@ -0,0 +1,82 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Утилиты для очистки текста (ANSI-коды, нормализация)."""
|
||||
|
||||
import re
|
||||
|
||||
|
||||
def clean_ansi_codes(text: str) -> str:
|
||||
"""
|
||||
Очистка ANSI-кодов и мусора из вывода терминала.
|
||||
Обрабатывает:
|
||||
- ANSI escape-последовательности \x1b[...m
|
||||
- «Битые» ANSI-коды без escape-символа (например [33m, [0m)
|
||||
- Символы замены Unicode ()
|
||||
- Кириллические имитации ANSI-кодов (например [0м)
|
||||
"""
|
||||
# Удаляем ANSI escape-последовательности
|
||||
text = re.sub(r'\x1b\[[0-9;?]*[a-zA-Z]', '', text)
|
||||
|
||||
# Удаляем «битые» ANSI-коды: [33m, [0m, [1m и т.д. (латиница и кириллица)
|
||||
text = re.sub(r'\[\d+[мm]', '', text)
|
||||
|
||||
# Удаляем символы замены Unicode
|
||||
text = text.replace('\ufffd', '')
|
||||
|
||||
return text
|
||||
|
||||
|
||||
def normalize_output(text: str) -> str:
|
||||
"""
|
||||
Нормализовать вывод: обработать \r и убрать пустые строки.
|
||||
\r используется для перезаписи строки (прогресс-баров).
|
||||
"""
|
||||
# Заменяем \r\n на \n
|
||||
text = text.replace('\r\n', '\n')
|
||||
|
||||
# Обрабатываем \r (возврат каретки) — строки с \r перезаписывают друг друга
|
||||
lines = []
|
||||
for line in text.split('\n'):
|
||||
if '\r' in line:
|
||||
# Разбиваем по \r и берём последнюю часть (финальное состояние)
|
||||
parts = line.split('\r')
|
||||
line = parts[-1]
|
||||
lines.append(line)
|
||||
|
||||
text = '\n'.join(lines)
|
||||
|
||||
# Разбиваем на строки, убираем пустые и trailing пробелы
|
||||
lines = text.split('\n')
|
||||
lines = [line.rstrip() for line in lines if line.strip()]
|
||||
|
||||
# Очищаем прогресс-бары вида "Текст… 0%Текст… 50%Текст… 100%"
|
||||
# И дублирующийся текст
|
||||
cleaned_lines = []
|
||||
for line in lines:
|
||||
# Ищем повторяющийся паттерн "текст… цифры%"
|
||||
progress_pattern = re.compile(r'((?:.+?\.{3})\d+%)+')
|
||||
match = progress_pattern.search(line)
|
||||
if match:
|
||||
# Берём последнее вхождение
|
||||
items = re.findall(r'(.+?\.{3})(\d+)%', match.group(0))
|
||||
if items:
|
||||
last_text, last_percent = items[-1]
|
||||
line = line[:match.start()] + f'{last_text}{last_percent}%' + line[match.end():]
|
||||
|
||||
# СНАЧАЛА удаляем остатки ANSI-кодов из строки
|
||||
# line = re.sub(r'.', '', line) # ← ЭТО УДАЛЯЛО ВСЁ! Закомментировал
|
||||
line = clean_ansi_codes(line) # ← Используем правильную функцию
|
||||
|
||||
# Удаляем дублирующийся текст вида "0% [текст] 0% [текст]"
|
||||
dup_pattern = re.compile(r'(\d+%\s*\[.+?\])(?:\s*\d+%\s*\[.+?\])+')
|
||||
match = dup_pattern.search(line)
|
||||
if match:
|
||||
# Оставляем только первое вхождение
|
||||
line = line[:match.start()] + match.group(1) + line[match.end():]
|
||||
|
||||
# Удаляем ведущие пробелы (артефакты терминала)
|
||||
line = line.lstrip()
|
||||
|
||||
if line:
|
||||
cleaned_lines.append(line)
|
||||
|
||||
return '\n'.join(cleaned_lines)
|
||||
@@ -0,0 +1,36 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Декораторы для бота."""
|
||||
|
||||
import logging
|
||||
from functools import wraps
|
||||
from telegram import Update
|
||||
from telegram.ext import ContextTypes
|
||||
|
||||
# Импортируем config для проверки доступа
|
||||
from bot.config import config
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def check_access(func):
|
||||
"""Декоратор для проверки прав доступа пользователя."""
|
||||
@wraps(func)
|
||||
async def wrapper(update: Update, context: ContextTypes.DEFAULT_TYPE, *args, **kwargs):
|
||||
user_id = update.effective_user.id
|
||||
|
||||
# Если доступ не ограничен — пропускаем всех
|
||||
if not config.is_access_restricted:
|
||||
return await func(update, context, *args, **kwargs)
|
||||
|
||||
if user_id not in config.allowed_users:
|
||||
logger.warning(f"Попытка доступа от запрещённого пользователя {user_id}")
|
||||
await update.message.reply_text(
|
||||
"❌ **Доступ запрещён**\n\n"
|
||||
"Ваш ID не добавлен в список разрешённых пользователей.\n"
|
||||
f"Ваш ID: `{user_id}`",
|
||||
parse_mode="Markdown"
|
||||
)
|
||||
return
|
||||
|
||||
return await func(update, context, *args, **kwargs)
|
||||
return wrapper
|
||||
@@ -0,0 +1,346 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Утилиты для форматирования и отправки сообщений."""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import re
|
||||
from typing import List, Tuple
|
||||
from telegram import Update
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Лимиты Telegram
|
||||
MAX_MESSAGE_LENGTH = 4096 # Максимальная длина сообщения
|
||||
RESERVED_FOR_HEADER = 20 # Резервируем место для "(N/N) "
|
||||
|
||||
|
||||
def escape_code_block_content(text: str) -> str:
|
||||
"""
|
||||
Экранировать спецсимволы Markdown внутри блоков кода.
|
||||
Нужно для случаев когда сообщение содержит ```блок кода``` и текст с Markdown.
|
||||
Telegram пытается интерпретировать [text](url) внутри блока кода как ссылку.
|
||||
"""
|
||||
# Находим все блоки кода
|
||||
parts = text.split("```")
|
||||
result_parts = []
|
||||
|
||||
for i, part in enumerate(parts):
|
||||
if i % 2 == 1:
|
||||
# Внутри блока кода — экранируем [ и ]
|
||||
part = part.replace('[', '\\[').replace(']', '\\]')
|
||||
result_parts.append(part)
|
||||
|
||||
return "```".join(result_parts)
|
||||
|
||||
|
||||
def escape_markdown(text: str) -> str:
|
||||
"""
|
||||
Экранирование специальных символов Markdown для Telegram API.
|
||||
|
||||
Telegram Markdown v1 использует: * _ ` [ ] ( )
|
||||
Эти символы нужно экранировать обратным слэшем.
|
||||
|
||||
ВАЖНО: Не экранирует содержимое блоков кода (```).
|
||||
"""
|
||||
if not text:
|
||||
return text
|
||||
|
||||
# Разбиваем текст на части: внутри и снаружи блоков кода
|
||||
parts = []
|
||||
last_end = 0
|
||||
in_code = False
|
||||
|
||||
# Находим все блоки кода
|
||||
code_pattern = re.compile(r'`[^`]+`|```[\s\S]*?```')
|
||||
|
||||
for match in code_pattern.finditer(text):
|
||||
# Экранируем текст перед блоком кода
|
||||
text_before = text[last_end:match.start()]
|
||||
parts.append(_escape_markdown_chars(text_before))
|
||||
# Добавляем блок кода без экранирования
|
||||
parts.append(match.group())
|
||||
last_end = match.end()
|
||||
|
||||
# Экранируем оставшийся текст после последнего блока кода
|
||||
remaining = text[last_end:]
|
||||
parts.append(_escape_markdown_chars(remaining))
|
||||
|
||||
return ''.join(parts)
|
||||
|
||||
|
||||
def escape_html(text: str) -> str:
|
||||
"""
|
||||
Экранирование специальных символов для Telegram Markdown.
|
||||
Алиас на escape_markdown для обратной совместимости.
|
||||
"""
|
||||
return escape_markdown(text)
|
||||
|
||||
|
||||
def _escape_markdown_chars(text: str) -> str:
|
||||
"""Экранировать специальные символы Markdown (вспомогательная функция)."""
|
||||
# Порядок важен: сначала экранируем обратные слэши
|
||||
text = text.replace('\\', '\\\\')
|
||||
text = text.replace('`', '\\`')
|
||||
text = text.replace('*', '\\*')
|
||||
text = text.replace('_', '\\_')
|
||||
text = text.replace('[', '\\[')
|
||||
text = text.replace(']', '\\]')
|
||||
text = text.replace('(', '\\(')
|
||||
text = text.replace(')', '\\)')
|
||||
return text
|
||||
|
||||
|
||||
def find_code_blocks(text: str) -> List[Tuple[int, int]]:
|
||||
"""
|
||||
Найти все блоки кода (```) в тексте.
|
||||
Возвращает список кортежей (start, end) для каждого блока.
|
||||
"""
|
||||
blocks = []
|
||||
pattern = re.compile(r'```')
|
||||
matches = list(pattern.finditer(text))
|
||||
|
||||
# Пары start-end для каждого блока
|
||||
i = 0
|
||||
while i < len(matches) - 1:
|
||||
start = matches[i].start()
|
||||
end = matches[i + 1].end()
|
||||
blocks.append((start, end))
|
||||
i += 2
|
||||
|
||||
return blocks
|
||||
|
||||
|
||||
def split_message(text: str, max_length: int = MAX_MESSAGE_LENGTH) -> List[Tuple[str, bool, bool, bool]]:
|
||||
"""
|
||||
Умно разбить длинный текст на сообщения <= max_length символов.
|
||||
|
||||
Возвращает список кортежей (text, has_code, code_opened, code_closed):
|
||||
- text: текст сообщения
|
||||
- has_code: True если сообщение содержит часть блока кода
|
||||
- code_opened: True если блок кода ОТКРЫТ в этом сообщении (есть открывающий ```)
|
||||
- code_closed: True если блок кода ЗАКРЫТ в этом сообщении (есть закрывающий ```)
|
||||
|
||||
Алгоритм:
|
||||
1. Разбиваем текст на строки
|
||||
2. Накапливаем строки до достижения лимита
|
||||
3. Отслеживаем состояние блока кода (внутри/снаружи)
|
||||
"""
|
||||
def calc_code_flags(txt: str) -> Tuple[bool, bool, bool]:
|
||||
"""Вычислить флаги code для данного текста."""
|
||||
has_code = '```' in txt
|
||||
backtick_count = txt.count('```')
|
||||
code_opened = backtick_count >= 1
|
||||
code_closed = backtick_count >= 2 or (backtick_count == 1 and txt.rstrip().endswith('```'))
|
||||
return has_code, code_opened, code_closed
|
||||
|
||||
if len(text) <= max_length:
|
||||
has_code, code_opened, code_closed = calc_code_flags(text)
|
||||
return [(text, has_code, code_opened, code_closed)]
|
||||
|
||||
parts = []
|
||||
lines = text.split('\n')
|
||||
current = ""
|
||||
in_code_block = False # Состояние: внутри блока кода или нет
|
||||
|
||||
for line in lines:
|
||||
# Проверяем, содержит ли строка ```
|
||||
backticks_in_line = line.count('```')
|
||||
|
||||
# Если строка содержит нечётное количество ```, она меняет состояние
|
||||
if backticks_in_line % 2 == 1:
|
||||
# Эта строка содержит ``` который меняет состояние
|
||||
if in_code_block:
|
||||
# Были внутри блока — эта строка закрывает его
|
||||
test_line = current + ('\n' if current else '') + line
|
||||
if len(test_line) > max_length - RESERVED_FOR_HEADER:
|
||||
# Сначала отправляем текущее (блок не закрыт в этой части!)
|
||||
if current:
|
||||
has_code, code_opened, _ = calc_code_flags(current)
|
||||
# code_closed=False потому что блок продолжится
|
||||
parts.append((current, has_code, code_opened, False))
|
||||
current = line
|
||||
else:
|
||||
current = test_line
|
||||
in_code_block = False
|
||||
else:
|
||||
# Были снаружи — эта строка открывает блок
|
||||
test_line = current + ('\n' if current else '') + line
|
||||
if len(test_line) > max_length - RESERVED_FOR_HEADER:
|
||||
if current:
|
||||
has_code, code_opened, code_closed = calc_code_flags(current)
|
||||
parts.append((current, has_code, code_opened, code_closed))
|
||||
current = line
|
||||
else:
|
||||
current = test_line
|
||||
in_code_block = True
|
||||
else:
|
||||
# Строка не меняет состояние
|
||||
test_line = current + ('\n' if current else '') + line
|
||||
if len(test_line) > max_length - RESERVED_FOR_HEADER:
|
||||
if current:
|
||||
has_code, code_opened, _ = calc_code_flags(current)
|
||||
# Если мы внутри блока кода — block не закрыт
|
||||
code_closed = not in_code_block
|
||||
parts.append((current, has_code, code_opened, code_closed))
|
||||
current = line
|
||||
else:
|
||||
current = test_line
|
||||
|
||||
if current:
|
||||
has_code, code_opened, _ = calc_code_flags(current)
|
||||
code_closed = not in_code_block
|
||||
parts.append((current, has_code, code_opened, code_closed))
|
||||
|
||||
return parts
|
||||
|
||||
|
||||
async def send_long_message(update: Update, text: str, parse_mode: str = None, pause_every: int = 3, start_from: int = 0):
|
||||
"""
|
||||
Отправить длинный текст, разбив на несколько сообщений.
|
||||
|
||||
Поддерживает:
|
||||
- Update с update.message (обычные сообщения)
|
||||
- CallbackQuery (query.edit_message_text / query.message.reply_text)
|
||||
|
||||
Args:
|
||||
start_from: Номер сообщения с которого начать (для продолжения после кнопки)
|
||||
"""
|
||||
from telegram import InlineKeyboardButton, InlineKeyboardMarkup
|
||||
from bot.config import state_manager
|
||||
|
||||
# Определяем тип объекта и получаем метод для отправки
|
||||
# CallbackQuery имеет from_user и answer(), но не имеет message.reply_text
|
||||
is_callback_query = hasattr(update, 'answer') and hasattr(update, 'from_user')
|
||||
|
||||
if is_callback_query:
|
||||
query = update
|
||||
message = query.message
|
||||
send_method = message.reply_text if message else query.edit_message_text
|
||||
user_id = query.from_user.id # Для CallbackQuery используем from_user
|
||||
else:
|
||||
message = update.message
|
||||
send_method = message.reply_text if message else None
|
||||
user_id = update.effective_user.id # Для Update используем effective_user
|
||||
|
||||
if not send_method:
|
||||
logger.error("send_long_message: не удалось определить метод отправки")
|
||||
return False
|
||||
|
||||
parts = split_message(text)
|
||||
total = len(parts)
|
||||
state = state_manager.get(user_id)
|
||||
|
||||
# Восстанавливаем состояние блока кода при продолжении
|
||||
prev_code_closed = state.output_prev_code_closed if start_from > 0 else True
|
||||
|
||||
# Начинаем с указанного сообщения
|
||||
for i in range(start_from, total):
|
||||
part, has_code, code_opened, code_closed = parts[i]
|
||||
|
||||
# Проверяем был ли блок кода открыт в предыдущем сообщении
|
||||
# При продолжении используем сохранённое состояние
|
||||
if i == start_from and start_from > 0:
|
||||
prev_closed = prev_code_closed
|
||||
else:
|
||||
prev_closed = parts[i-1][3] if i > 0 else True
|
||||
|
||||
# Определяем находимся ли внутри блока кода (между ``` и ```)
|
||||
in_code_block = not prev_closed or (code_opened and not code_closed)
|
||||
|
||||
# Проверяем будем ли добавлять ``` к этому сообщению
|
||||
will_add_opening = total > 1 and i > 0 and not prev_closed and not code_closed
|
||||
will_add_closing = total > 1 and i < total - 1 and not code_closed
|
||||
|
||||
# КЛЮЧЕВОЕ ИСПРАВЛЕНИЕ:
|
||||
# Отключаем parse_mode если:
|
||||
# 1. Это промежуточная часть блока кода (in_code_block=True, но нет ``` в этом сообщении)
|
||||
# 2. Это сообщение содержит ``` (has_code=True)
|
||||
# 3. Мы добавляем искусственные ``` (will_add_opening или will_add_closing)
|
||||
if in_code_block or has_code or will_add_opening or will_add_closing:
|
||||
actual_parse_mode = None
|
||||
else:
|
||||
# Обычный текст без блоков кода — используем HTML
|
||||
actual_parse_mode = parse_mode
|
||||
|
||||
# Логика работы с блоками кода между сообщениями
|
||||
if total > 1 and i > 0 and not prev_closed:
|
||||
part = "```\n" + part
|
||||
|
||||
if total > 1 and i < total - 1 and not code_closed:
|
||||
part = part + "\n```"
|
||||
|
||||
# Добавляем номер части
|
||||
if total > 1:
|
||||
header = f"({i+1}/{total}) "
|
||||
if len(header) + len(part) <= MAX_MESSAGE_LENGTH:
|
||||
part = header + part
|
||||
|
||||
try:
|
||||
await send_method(part, parse_mode=actual_parse_mode)
|
||||
except Exception as e:
|
||||
logger.debug(f"Ошибка Markdown, отправляем без разметки: {e}")
|
||||
await send_method(part)
|
||||
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
# КАЖДЫЕ pause_every сообщений — спрашивать продолжать ли
|
||||
if pause_every > 0 and (i + 1) % pause_every == 0 and i < total - 1:
|
||||
remaining = total - (i + 1)
|
||||
keyboard = InlineKeyboardMarkup([
|
||||
[
|
||||
InlineKeyboardButton("▶️ Продолжить", callback_data=f"continue_output_{remaining}_{i+1}"),
|
||||
InlineKeyboardButton("❌ Отменить", callback_data="cancel_output")
|
||||
]
|
||||
])
|
||||
|
||||
wait_msg = await send_method(
|
||||
f"📊 **Отправлено {i + 1} из {total} сообщений**\n\n"
|
||||
f"Осталось ещё {remaining} сообщений.\n\n"
|
||||
f"Продолжить вывод?",
|
||||
parse_mode="Markdown",
|
||||
reply_markup=keyboard
|
||||
)
|
||||
|
||||
# Сохраняем состояние и ВОЗВРАЩАЕМ УПРАВЛЕНИЕ
|
||||
state.waiting_for_output_control = True
|
||||
state.output_remaining = remaining
|
||||
state.output_wait_message = wait_msg
|
||||
state.output_next_index = i + 1 # С какого сообщения продолжить
|
||||
state.output_text = text # Сохраняем текст для продолжения
|
||||
state.output_parse_mode = parse_mode
|
||||
state.output_prev_code_closed = code_closed # Сохраняем состояние блока кода
|
||||
|
||||
logger.info(f"send_long_message: пауза после {i+1}/{total}, ждём кнопки (user_id={user_id})")
|
||||
return True # Возвращаем True — есть продолжение
|
||||
|
||||
# Все сообщения отправлены
|
||||
state.waiting_for_output_control = False
|
||||
state.output_remaining = None
|
||||
state.output_wait_message = None
|
||||
state.output_next_index = None
|
||||
state.output_text = None
|
||||
return False # Возвращаем False — продолжения нет
|
||||
|
||||
|
||||
def format_long_output(text: str, max_lines: int = 100, head_lines: int = 50, tail_lines: int = 50) -> str:
|
||||
"""
|
||||
Форматировать длинный вывод: показать первые и последние строки.
|
||||
По умолчанию: первые 50 + последние 50 строк = 100 строк максимум.
|
||||
"""
|
||||
lines = text.split('\n')
|
||||
total_lines = len(lines)
|
||||
|
||||
if total_lines <= max_lines:
|
||||
return text
|
||||
|
||||
# Показываем первые head_lines и последние tail_lines
|
||||
head = lines[:head_lines]
|
||||
tail = lines[-tail_lines:]
|
||||
|
||||
skipped = total_lines - head_lines - tail_lines
|
||||
|
||||
result = '\n'.join(head)
|
||||
result += f'\n\n... ({skipped} строк пропущено) ...\n'
|
||||
result += '\n'.join(tail)
|
||||
|
||||
return result
|
||||
@@ -0,0 +1,593 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Qwen OAuth 2.0 Device Flow клиент.
|
||||
Реализует авторизацию через Device Authorization Grant (RFC 8628).
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import json
|
||||
import hashlib
|
||||
import secrets
|
||||
import time
|
||||
import aiohttp
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Optional, Dict, Any
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Qwen OAuth константы
|
||||
QWEN_OAUTH_BASE_URL = 'https://chat.qwen.ai'
|
||||
QWEN_OAUTH_DEVICE_CODE_ENDPOINT = f'{QWEN_OAUTH_BASE_URL}/api/v1/oauth2/device/code'
|
||||
QWEN_OAUTH_TOKEN_ENDPOINT = f'{QWEN_OAUTH_BASE_URL}/api/v1/oauth2/token'
|
||||
QWEN_OAUTH_CLIENT_ID = 'f0304373b74a44d2b584a3fb70ca9e56'
|
||||
QWEN_OAUTH_SCOPE = 'openid profile email model.completion'
|
||||
|
||||
# Пути для хранения токенов (как в qwen-code CLI)
|
||||
QWEN_CONFIG_DIR = Path.home() / '.qwen'
|
||||
QWEN_CREDENTIALS_FILE = QWEN_CONFIG_DIR / 'oauth_creds.json'
|
||||
|
||||
|
||||
@dataclass
|
||||
class QwenCredentials:
|
||||
"""OAuth токены Qwen."""
|
||||
access_token: str = ''
|
||||
refresh_token: str = ''
|
||||
token_type: str = 'Bearer'
|
||||
expiry_date: int = 0 # Unix timestamp в миллисекундах
|
||||
resource_url: str = 'portal.qwen.ai'
|
||||
|
||||
def is_expired(self, buffer_minutes: int = 5) -> bool:
|
||||
"""Проверка истечения токена с буфером."""
|
||||
if not self.expiry_date:
|
||||
return True
|
||||
expiry_ms = self.expiry_date - (buffer_minutes * 60 * 1000)
|
||||
return int(datetime.now().timestamp() * 1000) >= expiry_ms
|
||||
|
||||
|
||||
@dataclass
|
||||
class DeviceAuthorizationResponse:
|
||||
"""Ответ устройства авторизации."""
|
||||
device_code: str = ''
|
||||
user_code: str = ''
|
||||
verification_uri: str = ''
|
||||
verification_uri_complete: str = ''
|
||||
expires_in: int = 0
|
||||
interval: int = 5 # Polling interval в секундах
|
||||
|
||||
@property
|
||||
def authorization_url(self) -> str:
|
||||
"""Полная ссылка для авторизации."""
|
||||
return self.verification_uri_complete
|
||||
|
||||
@property
|
||||
def is_expired(self) -> bool:
|
||||
"""Истёк ли срок действия device code."""
|
||||
return time.time() > (self.expires_in - 30) # 30 сек буфер
|
||||
|
||||
|
||||
class QwenOAuthClient:
|
||||
"""Qwen OAuth 2.0 клиент."""
|
||||
|
||||
# Как в qwen-code: 30 секунд буфер + 5 секунд интервал проверки
|
||||
TOKEN_REFRESH_BUFFER_MS = 30 * 1000
|
||||
CACHE_CHECK_INTERVAL_MS = 5 * 1000
|
||||
|
||||
def __init__(self, credentials_path: Optional[Path] = None):
|
||||
self.credentials_path = credentials_path or QWEN_CREDENTIALS_FILE
|
||||
self._credentials: Optional[QwenCredentials] = None
|
||||
self._file_mod_time: float = 0 # Время последней модификации файла
|
||||
self._last_check_time: float = 0 # Время последней проверки
|
||||
self._load_credentials()
|
||||
|
||||
def _get_file_mod_time(self) -> float:
|
||||
"""Получить время модификации файла токенов."""
|
||||
try:
|
||||
if self.credentials_path.exists():
|
||||
return self.credentials_path.stat().st_mtime
|
||||
except Exception:
|
||||
pass
|
||||
return 0
|
||||
|
||||
def _load_credentials(self, force: bool = False) -> None:
|
||||
"""
|
||||
Загрузить токены из файла.
|
||||
|
||||
Args:
|
||||
force: Принудительная перезагрузка даже если файл не изменился
|
||||
"""
|
||||
current_mod_time = self._get_file_mod_time()
|
||||
now = time.time()
|
||||
|
||||
# Если force=True — пропускаем все проверки и загружаем всегда
|
||||
if not force:
|
||||
# Проверяем не слишком ли часто проверяем (как в qwen-code)
|
||||
if (now - self._last_check_time) < (self.CACHE_CHECK_INTERVAL_MS / 1000):
|
||||
return
|
||||
|
||||
# Проверяем изменился ли файл (как в qwen-code)
|
||||
if current_mod_time <= self._file_mod_time:
|
||||
# Файл не изменился — используем кэш
|
||||
return
|
||||
|
||||
self._last_check_time = now
|
||||
|
||||
logger.debug(f"Загрузка токенов из {self.credentials_path} (mod_time={current_mod_time}, force={force})")
|
||||
if self.credentials_path.exists():
|
||||
try:
|
||||
with open(self.credentials_path, 'r') as f:
|
||||
data = json.load(f)
|
||||
self._credentials = QwenCredentials(**data)
|
||||
self._file_mod_time = current_mod_time
|
||||
logger.info(f"Токены загружены из {self.credentials_path}")
|
||||
logger.debug(f"Access token: {self._credentials.access_token[:20] if self._credentials.access_token else 'None'}..., expiry: {self._credentials.expiry_date}")
|
||||
# НЕ вызываем has_valid_token() здесь — это вызывает рекурсию!
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка загрузки токенов: {e}")
|
||||
self._credentials = None
|
||||
self._file_mod_time = 0
|
||||
else:
|
||||
logger.debug("Файл с токенами не найден")
|
||||
self._credentials = None
|
||||
self._file_mod_time = 0
|
||||
|
||||
# Загружаем code_verifier из device_code.json если есть
|
||||
device_code_file = QWEN_CONFIG_DIR / 'device_code.json'
|
||||
if device_code_file.exists():
|
||||
try:
|
||||
with open(device_code_file, 'r') as f:
|
||||
data = json.load(f)
|
||||
code_verifier = data.get('code_verifier', '')
|
||||
if code_verifier:
|
||||
self._code_verifier = code_verifier
|
||||
logger.info(f"Code verifier загружен из {device_code_file}")
|
||||
except Exception as e:
|
||||
logger.debug(f"Ошибка загрузки code_verifier: {e}")
|
||||
|
||||
def _save_credentials(self) -> None:
|
||||
"""Сохранить токены в файл."""
|
||||
if self._credentials:
|
||||
self.credentials_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(self.credentials_path, 'w') as f:
|
||||
json.dump(self._credentials.__dict__, f, indent=2)
|
||||
# Устанавливаем права 600 (только владелец)
|
||||
os.chmod(self.credentials_path, 0o600)
|
||||
logger.info(f"Токены сохранены в {self.credentials_path}")
|
||||
|
||||
def has_valid_token(self) -> bool:
|
||||
"""Проверка наличия валидного токена."""
|
||||
# Принудительно перезагружаем токены из файла
|
||||
# Это нужно потому что токены могут быть обновлены другим процессом (qwen-code CLI)
|
||||
self._load_credentials(force=True)
|
||||
return self._credentials is not None and not self._credentials.is_expired()
|
||||
|
||||
async def get_access_token(self) -> Optional[str]:
|
||||
"""Получить access token (обновляет если истёк)."""
|
||||
if self.has_valid_token():
|
||||
return self._credentials.access_token
|
||||
|
||||
if self._credentials and self._credentials.refresh_token:
|
||||
# Пробуем обновить токен
|
||||
if await self._refresh_token():
|
||||
return self._credentials.access_token
|
||||
|
||||
return None
|
||||
|
||||
async def request_device_authorization(self) -> DeviceAuthorizationResponse:
|
||||
"""
|
||||
Запросить Device Authorization.
|
||||
|
||||
Returns:
|
||||
DeviceAuthorizationResponse с данными для авторизации
|
||||
"""
|
||||
# Проверяем есть ли сохранённый code_verifier
|
||||
if not hasattr(self, '_code_verifier') or not self._code_verifier:
|
||||
# Генерируем PKCE code verifier и challenge
|
||||
self._code_verifier = secrets.token_urlsafe(32)
|
||||
|
||||
code_challenge = hashlib.sha256(self._code_verifier.encode()).hexdigest()
|
||||
|
||||
payload = {
|
||||
'client_id': QWEN_OAUTH_CLIENT_ID,
|
||||
'scope': QWEN_OAUTH_SCOPE,
|
||||
'code_challenge': code_challenge,
|
||||
'code_challenge_method': 'S256'
|
||||
}
|
||||
|
||||
headers = {
|
||||
'Content-Type': 'application/x-www-form-urlencoded',
|
||||
'Accept': 'application/json',
|
||||
'x-request-id': secrets.token_hex(16),
|
||||
'User-Agent': 'qwen-code-cli/0.11.0'
|
||||
}
|
||||
|
||||
form_data = '&'.join(f'{k}={v}' for k, v in payload.items())
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.post(
|
||||
QWEN_OAUTH_DEVICE_CODE_ENDPOINT,
|
||||
data=form_data,
|
||||
headers=headers
|
||||
) as resp:
|
||||
if resp.status != 200:
|
||||
text = await resp.text()
|
||||
raise Exception(f"Device authorization failed: {resp.status} - {text}")
|
||||
|
||||
data = await resp.json()
|
||||
|
||||
return DeviceAuthorizationResponse(
|
||||
device_code=data.get('device_code', ''),
|
||||
user_code=data.get('user_code', ''),
|
||||
verification_uri=data.get('verification_uri', ''),
|
||||
verification_uri_complete=data.get('verification_uri_complete', ''),
|
||||
expires_in=data.get('expires_in', 900),
|
||||
interval=data.get('interval', 5)
|
||||
)
|
||||
|
||||
async def poll_for_token(self, device_code: str, timeout_seconds: int = 900) -> bool:
|
||||
"""
|
||||
Опрос сервера для получения токена после авторизации пользователем.
|
||||
|
||||
Args:
|
||||
device_code: Device code из request_device_authorization
|
||||
timeout_seconds: Максимальное время ожидания
|
||||
|
||||
Returns:
|
||||
True если авторизация успешна
|
||||
"""
|
||||
if not hasattr(self, '_code_verifier'):
|
||||
raise Exception("Code verifier not set. Call request_device_authorization first.")
|
||||
|
||||
start_time = time.time()
|
||||
interval = 5 # Начальный интервал
|
||||
|
||||
while time.time() - start_time < timeout_seconds:
|
||||
payload = {
|
||||
'grant_type': 'urn:ietf:params:oauth:grant-type:device_code',
|
||||
'device_code': device_code,
|
||||
'code_verifier': self._code_verifier,
|
||||
'client_id': QWEN_OAUTH_CLIENT_ID
|
||||
}
|
||||
|
||||
headers = {
|
||||
'Content-Type': 'application/x-www-form-urlencoded',
|
||||
'Accept': 'application/json',
|
||||
'User-Agent': 'qwen-code-cli/0.11.0'
|
||||
}
|
||||
|
||||
form_data = '&'.join(f'{k}={v}' for k, v in payload.items())
|
||||
|
||||
logger.debug(f"Polling payload: grant_type={payload['grant_type']}, device_code={device_code[:20]}..., code_verifier={self._code_verifier[:20]}...")
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.post(
|
||||
QWEN_OAUTH_TOKEN_ENDPOINT,
|
||||
data=form_data,
|
||||
headers=headers
|
||||
) as resp:
|
||||
data = await resp.json()
|
||||
logger.info(f"Polling response: status={resp.status}, data={data}")
|
||||
|
||||
if resp.status == 200:
|
||||
# Успех! Сохраняем токены
|
||||
self._credentials = QwenCredentials(
|
||||
access_token=data.get('access_token', ''),
|
||||
refresh_token=data.get('refresh_token', ''),
|
||||
token_type=data.get('token_type', 'Bearer'),
|
||||
expiry_date=int(datetime.now().timestamp() * 1000) + (data.get('expires_in', 3600) * 1000),
|
||||
resource_url=data.get('resource_url', 'portal.qwen.ai')
|
||||
)
|
||||
self._save_credentials()
|
||||
logger.info("Авторизация успешна!")
|
||||
return True
|
||||
|
||||
error = data.get('error', '')
|
||||
|
||||
if error == 'authorization_pending':
|
||||
# Пользователь ещё не авторизовался
|
||||
logger.debug("Ожидание авторизации пользователя...")
|
||||
await asyncio.sleep(interval)
|
||||
continue
|
||||
|
||||
elif error == 'slow_down':
|
||||
# Сервер просит увеличить интервал
|
||||
interval += 5
|
||||
logger.debug(f"Увеличиваем интервал до {interval} сек")
|
||||
await asyncio.sleep(interval)
|
||||
continue
|
||||
|
||||
elif error == 'expired_token':
|
||||
logger.error("Device code истёк")
|
||||
return False
|
||||
|
||||
elif error == 'access_denied':
|
||||
logger.error("Пользователь отклонил авторизацию")
|
||||
return False
|
||||
|
||||
else:
|
||||
logger.error(f"Неизвестная ошибка: {error}")
|
||||
return False
|
||||
|
||||
logger.error("Таймаут авторизации")
|
||||
return False
|
||||
|
||||
async def _refresh_token(self) -> bool:
|
||||
"""Обновить access token используя refresh token."""
|
||||
if not self._credentials or not self._credentials.refresh_token:
|
||||
return False
|
||||
|
||||
payload = {
|
||||
'grant_type': 'refresh_token',
|
||||
'refresh_token': self._credentials.refresh_token,
|
||||
'client_id': QWEN_OAUTH_CLIENT_ID
|
||||
}
|
||||
|
||||
headers = {
|
||||
'Content-Type': 'application/x-www-form-urlencoded',
|
||||
'Accept': 'application/json',
|
||||
'User-Agent': 'qwen-code-cli/0.11.0'
|
||||
}
|
||||
|
||||
form_data = '&'.join(f'{k}={v}' for k, v in payload.items())
|
||||
|
||||
try:
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.post(
|
||||
QWEN_OAUTH_TOKEN_ENDPOINT,
|
||||
data=form_data,
|
||||
headers=headers
|
||||
) as resp:
|
||||
if resp.status == 200:
|
||||
data = await resp.json()
|
||||
self._credentials.access_token = data.get('access_token', '')
|
||||
self._credentials.refresh_token = data.get('refresh_token', self._credentials.refresh_token)
|
||||
self._credentials.expiry_date = int(datetime.now().timestamp() * 1000) + (data.get('expires_in', 3600) * 1000)
|
||||
self._save_credentials()
|
||||
logger.info("Токен обновлён")
|
||||
return True
|
||||
else:
|
||||
logger.error(f"Ошибка обновления токена: {resp.status}")
|
||||
self._credentials = None # Очищаем неверные токены
|
||||
return False
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка обновления токена: {e}")
|
||||
return False
|
||||
|
||||
def clear_credentials(self) -> None:
|
||||
"""Очистить сохранённые токены."""
|
||||
self._credentials = None
|
||||
if self.credentials_path.exists():
|
||||
self.credentials_path.unlink()
|
||||
logger.info("Токены очищены")
|
||||
|
||||
|
||||
# Глобальный клиент (singleton)
|
||||
_oauth_client: Optional[QwenOAuthClient] = None
|
||||
|
||||
|
||||
def get_oauth_client() -> QwenOAuthClient:
|
||||
"""Получить OAuth клиент (singleton)."""
|
||||
global _oauth_client
|
||||
if _oauth_client is None:
|
||||
_oauth_client = QwenOAuthClient()
|
||||
return _oauth_client
|
||||
|
||||
|
||||
async def get_authorization_url() -> Optional[str]:
|
||||
"""
|
||||
Получить URL для авторизации.
|
||||
|
||||
Returns:
|
||||
URL для авторизации или None если ошибка
|
||||
"""
|
||||
try:
|
||||
# Проверяем есть ли активный device code
|
||||
device_code_file = QWEN_CONFIG_DIR / 'device_code.json'
|
||||
if device_code_file.exists():
|
||||
with open(device_code_file, 'r') as f:
|
||||
data = json.load(f)
|
||||
start_time = data.get('start_time', 0)
|
||||
expires_in = data.get('expires_in', 900)
|
||||
code_verifier = data.get('code_verifier', '')
|
||||
auth_url = data.get('authorization_url', '')
|
||||
device_code = data.get('device_code', '')
|
||||
|
||||
# Если device code ещё активен — используем его
|
||||
if time.time() - start_time < expires_in - 60 and code_verifier and auth_url and device_code:
|
||||
logger.info("Используем существующий device code")
|
||||
return auth_url
|
||||
|
||||
# Генерируем PKCE пару как в qwen-code CLI
|
||||
import base64
|
||||
import hashlib
|
||||
|
||||
code_verifier = base64.urlsafe_b64encode(os.urandom(32)).decode('utf-8').rstrip('=')
|
||||
code_challenge = base64.urlsafe_b64encode(hashlib.sha256(code_verifier.encode()).digest()).decode('utf-8').rstrip('=')
|
||||
|
||||
# Запрос device authorization
|
||||
payload = {
|
||||
'client_id': QWEN_OAUTH_CLIENT_ID,
|
||||
'scope': QWEN_OAUTH_SCOPE,
|
||||
'code_challenge': code_challenge,
|
||||
'code_challenge_method': 'S256'
|
||||
}
|
||||
|
||||
headers = {
|
||||
'Content-Type': 'application/x-www-form-urlencoded',
|
||||
'Accept': 'application/json',
|
||||
'x-request-id': secrets.token_hex(16),
|
||||
'User-Agent': 'qwen-code-cli/0.11.0'
|
||||
}
|
||||
|
||||
form_data = '&'.join(f'{k}={v}' for k, v in payload.items())
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.post(
|
||||
QWEN_OAUTH_DEVICE_CODE_ENDPOINT,
|
||||
data=form_data,
|
||||
headers=headers
|
||||
) as resp:
|
||||
if resp.status != 200:
|
||||
text = await resp.text()
|
||||
raise Exception(f"Device authorization failed: {resp.status} - {text}")
|
||||
|
||||
data = await resp.json()
|
||||
|
||||
auth_url = data.get('verification_uri_complete', '')
|
||||
device_code = data.get('device_code', '')
|
||||
expires_in = data.get('expires_in', 900)
|
||||
|
||||
# Сохраняем device code и code verifier для polling
|
||||
device_code_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(device_code_file, 'w') as f:
|
||||
json.dump({
|
||||
'device_code': device_code,
|
||||
'code_verifier': code_verifier, # Сохраняем тот же code_verifier!
|
||||
'expires_in': expires_in,
|
||||
'start_time': time.time(),
|
||||
'authorization_url': auth_url
|
||||
}, f)
|
||||
|
||||
logger.info(f"Device code получен: {device_code[:20]}..., code_verifier сохранён")
|
||||
return auth_url
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка получения URL авторизации: {e}")
|
||||
return None
|
||||
|
||||
|
||||
async def check_authorization_complete() -> bool:
|
||||
"""
|
||||
Проверить завершение авторизации (polling).
|
||||
|
||||
Returns:
|
||||
True если авторизация завершена успешно
|
||||
"""
|
||||
try:
|
||||
# Читаем device code
|
||||
device_code_file = QWEN_CONFIG_DIR / 'device_code.json'
|
||||
if not device_code_file.exists():
|
||||
logger.debug("device_code.json не найден")
|
||||
return False
|
||||
|
||||
with open(device_code_file, 'r') as f:
|
||||
data = json.load(f)
|
||||
|
||||
device_code = data.get('device_code', '')
|
||||
code_verifier = data.get('code_verifier', '') # Получаем code_verifier из файла
|
||||
start_time = data.get('start_time', time.time())
|
||||
expires_in = data.get('expires_in', 900)
|
||||
|
||||
logger.info(f"Device code: {device_code[:20]}..., code_verifier: {code_verifier[:20]}...")
|
||||
|
||||
# Проверяем не истёк ли timeout
|
||||
if time.time() - start_time > expires_in:
|
||||
logger.warning("Device code истёк")
|
||||
device_code_file.unlink()
|
||||
return False
|
||||
|
||||
# Polling с code_verifier из файла
|
||||
logger.info("Запуск polling для получения токена...")
|
||||
|
||||
payload = {
|
||||
'grant_type': 'urn:ietf:params:oauth:grant-type:device_code',
|
||||
'device_code': device_code,
|
||||
'code_verifier': code_verifier,
|
||||
'client_id': QWEN_OAUTH_CLIENT_ID
|
||||
}
|
||||
|
||||
headers = {
|
||||
'Content-Type': 'application/x-www-form-urlencoded',
|
||||
'Accept': 'application/json',
|
||||
'User-Agent': 'qwen-code-cli/0.11.0'
|
||||
}
|
||||
|
||||
form_data = '&'.join(f'{k}={v}' for k, v in payload.items())
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.post(
|
||||
QWEN_OAUTH_TOKEN_ENDPOINT,
|
||||
data=form_data,
|
||||
headers=headers
|
||||
) as resp:
|
||||
data = await resp.json()
|
||||
logger.info(f"Polling response: status={resp.status}, data={data}")
|
||||
|
||||
if resp.status == 200:
|
||||
# Успех! Сохраняем токены
|
||||
credentials = {
|
||||
'access_token': data.get('access_token', ''),
|
||||
'refresh_token': data.get('refresh_token', ''),
|
||||
'token_type': data.get('token_type', 'Bearer'),
|
||||
'expiry_date': int(time.time() * 1000) + (data.get('expires_in', 3600) * 1000),
|
||||
'resource_url': data.get('resource_url', 'portal.qwen.ai')
|
||||
}
|
||||
|
||||
# Сохраняем токены в файл
|
||||
QWEN_CONFIG_DIR.mkdir(parents=True, exist_ok=True)
|
||||
with open(QWEN_CREDENTIALS_FILE, 'w') as f:
|
||||
json.dump(credentials, f, indent=2)
|
||||
os.chmod(QWEN_CREDENTIALS_FILE, 0o600)
|
||||
|
||||
device_code_file.unlink()
|
||||
logger.info("Авторизация успешна! Токены сохранены.")
|
||||
return True
|
||||
|
||||
error = data.get('error', '')
|
||||
|
||||
if error == 'authorization_pending':
|
||||
logger.debug("Ожидание авторизации пользователя...")
|
||||
return False
|
||||
|
||||
elif error == 'slow_down':
|
||||
logger.debug("Сервер просит увеличить интервал")
|
||||
await asyncio.sleep(5)
|
||||
return False
|
||||
|
||||
elif error == 'expired_token':
|
||||
logger.error("Device code истёк")
|
||||
device_code_file.unlink()
|
||||
return False
|
||||
|
||||
elif error == 'access_denied':
|
||||
logger.error("Пользователь отклонил авторизацию")
|
||||
device_code_file.unlink()
|
||||
return False
|
||||
|
||||
elif error == 'invalid_request':
|
||||
error_desc = data.get('error_description', '')
|
||||
logger.error(f"Invalid request: {error_desc}")
|
||||
# Проверяем не истёк ли code_verifier
|
||||
if 'code_verifier' in error_desc.lower():
|
||||
logger.error("Code verifier не совпадает — удаляем device_code.json")
|
||||
device_code_file.unlink()
|
||||
return False
|
||||
|
||||
else:
|
||||
logger.error(f"Неизвестная ошибка: {error}")
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка проверки авторизации: {e}", exc_info=True)
|
||||
return False
|
||||
|
||||
|
||||
async def is_authorized() -> bool:
|
||||
"""Проверить авторизован ли пользователь."""
|
||||
client = get_oauth_client()
|
||||
return client.has_valid_token()
|
||||
|
||||
|
||||
async def get_access_token() -> Optional[str]:
|
||||
"""Получить access token."""
|
||||
client = get_oauth_client()
|
||||
return await client.get_access_token()
|
||||
|
||||
|
||||
def clear_authorization() -> None:
|
||||
"""Очистить авторизацию."""
|
||||
client = get_oauth_client()
|
||||
client.clear_credentials()
|
||||
@@ -0,0 +1,293 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Утилиты для чтения вывода SSH и PTY."""
|
||||
|
||||
import asyncio
|
||||
import fcntl
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import select
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import asyncssh
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Импортируем паттерны из session
|
||||
from bot.models.session import INPUT_PATTERNS
|
||||
|
||||
|
||||
def detect_input_type(text: str) -> Optional[str]:
|
||||
"""Определить тип запроса ввода по тексту."""
|
||||
text = text.strip()
|
||||
|
||||
# Проверка на пароль
|
||||
for pattern in INPUT_PATTERNS["password"]:
|
||||
if re.search(pattern, text, re.MULTILINE):
|
||||
return "password"
|
||||
|
||||
# Проверка на подтверждение
|
||||
for pattern in INPUT_PATTERNS["confirm"]:
|
||||
if re.search(pattern, text, re.MULTILINE):
|
||||
return "confirm"
|
||||
|
||||
# Проверка на приглашение оболочки
|
||||
for pattern in INPUT_PATTERNS["shell_prompt"]:
|
||||
if re.search(pattern, text, re.MULTILINE):
|
||||
return "prompt"
|
||||
|
||||
return None
|
||||
|
||||
|
||||
async def read_ssh_output(process: asyncssh.SSHClientProcess, timeout: float = 2.0, wait_for_completion: bool = False) -> Tuple[str, bool]:
|
||||
"""
|
||||
Чтение вывода из SSH-процесса с таймаутом.
|
||||
|
||||
Args:
|
||||
process: SSH процесс для чтения
|
||||
timeout: Таймаут для чтения данных (сек)
|
||||
wait_for_completion: Если True, дождаться завершения процесса через process.wait()
|
||||
|
||||
Returns:
|
||||
(вывод, завершён_ли_процесс)
|
||||
"""
|
||||
output = ""
|
||||
error_output = ""
|
||||
is_done = False
|
||||
|
||||
try:
|
||||
# Используем read() для чтения доступных данных
|
||||
while True:
|
||||
try:
|
||||
# read() читает данные до EOF
|
||||
data = await asyncio.wait_for(process.stdout.read(), timeout=timeout)
|
||||
if data:
|
||||
if isinstance(data, bytes):
|
||||
output += data.decode('utf-8', errors='replace')
|
||||
else:
|
||||
output += str(data)
|
||||
logger.debug(f"Прочитано stdout: {len(data)} байт, всего: {len(output)}")
|
||||
else:
|
||||
# EOF
|
||||
logger.debug("SSH stdout EOF")
|
||||
is_done = True
|
||||
break
|
||||
except asyncio.TimeoutError:
|
||||
# Данные закончились по таймауту
|
||||
logger.debug(f"Timeout stdout ({timeout} сек), прочитано: {len(output)} байт")
|
||||
break
|
||||
except UnicodeDecodeError as e:
|
||||
logger.debug(f"Ошибка декодирования UTF-8: {e}")
|
||||
continue
|
||||
except Exception as e:
|
||||
# Конец потока
|
||||
logger.debug(f"Конец потока stdout: {type(e).__name__}: {e}")
|
||||
is_done = True
|
||||
break
|
||||
except Exception as e:
|
||||
logger.debug(f"Ошибка чтения SSH stdout: {type(e).__name__}: {e}")
|
||||
is_done = True
|
||||
|
||||
# Читаем stderr если есть
|
||||
try:
|
||||
while True:
|
||||
try:
|
||||
data = await asyncio.wait_for(process.stderr.read(), timeout=0.5)
|
||||
if data:
|
||||
if isinstance(data, bytes):
|
||||
error_output += data.decode('utf-8', errors='replace')
|
||||
else:
|
||||
error_output += str(data)
|
||||
else:
|
||||
break
|
||||
except (asyncio.TimeoutError, Exception):
|
||||
break
|
||||
except Exception as e:
|
||||
logger.debug(f"Ошибка чтения SSH stderr: {type(e).__name__}: {e}")
|
||||
|
||||
# Объединяем stdout и stderr
|
||||
if error_output:
|
||||
output = output + error_output if output else error_output
|
||||
|
||||
logger.info(f"read_ssh_output: output={len(output)} байт, is_done={is_done}, returncode={process.returncode}")
|
||||
return output, is_done
|
||||
|
||||
|
||||
async def wait_and_read_ssh(process: asyncssh.SSHClientProcess, timeout: float = 30.0) -> Tuple[str, str, int]:
|
||||
"""
|
||||
Чтение вывода SSH-процесса с ожиданием полного завершения.
|
||||
Аналог asyncio.subprocess.communicate() для asyncssh.
|
||||
|
||||
Эта функция решает проблему с returncode, который становится доступен
|
||||
только после завершения процесса. Читает stdout и stderr параллельно
|
||||
с выполнением команды.
|
||||
|
||||
Args:
|
||||
process: SSH процесс
|
||||
timeout: Максимальное время ожидания выполнения (сек)
|
||||
|
||||
Returns:
|
||||
(stdout, stderr, returncode)
|
||||
"""
|
||||
stdout_data = ""
|
||||
stderr_data = ""
|
||||
|
||||
async def read_stream(stream, is_stdout=True):
|
||||
"""Читает поток до EOF."""
|
||||
data = ""
|
||||
try:
|
||||
while True:
|
||||
chunk = await stream.read()
|
||||
if not chunk:
|
||||
break
|
||||
if isinstance(chunk, bytes):
|
||||
data += chunk.decode('utf-8', errors='replace')
|
||||
else:
|
||||
data += str(chunk)
|
||||
stream_name = "stdout" if is_stdout else "stderr"
|
||||
logger.debug(f"{stream_name}: прочитано {len(chunk)} байт")
|
||||
except Exception as e:
|
||||
stream_name = "stdout" if is_stdout else "stderr"
|
||||
logger.debug(f"{stream_name} завершен: {type(e).__name__}: {e}")
|
||||
return data
|
||||
|
||||
try:
|
||||
# Читаем stdout и stderr параллельно с ожиданием завершения
|
||||
logger.debug(f"wait_and_read_ssh: запуск чтения (timeout={timeout})")
|
||||
|
||||
# Создаём задачи для чтения stdout и stderr
|
||||
stdout_task = asyncio.create_task(read_stream(process.stdout, is_stdout=True))
|
||||
stderr_task = asyncio.create_task(read_stream(process.stderr, is_stdout=False))
|
||||
|
||||
# Ждём завершения процесса с таймаутом
|
||||
await asyncio.wait_for(process.wait(), timeout=timeout)
|
||||
logger.debug(f"wait_and_read_ssh: процесс завершился, returncode={process.returncode}")
|
||||
|
||||
# Ждём завершения чтения с коротким таймаутом
|
||||
try:
|
||||
stdout_data = await asyncio.wait_for(stdout_task, timeout=2.0)
|
||||
except asyncio.TimeoutError:
|
||||
logger.warning("wait_and_read_ssh: таймаут чтения stdout")
|
||||
stdout_task.cancel()
|
||||
try:
|
||||
await stdout_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
try:
|
||||
stderr_data = await asyncio.wait_for(stderr_task, timeout=2.0)
|
||||
except asyncio.TimeoutError:
|
||||
logger.warning("wait_and_read_ssh: таймаут чтения stderr")
|
||||
stderr_task.cancel()
|
||||
try:
|
||||
await stderr_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
logger.error(f"wait_and_read_ssh: таймаут выполнения команды ({timeout} сек)")
|
||||
# Отменяем задачи чтения
|
||||
stdout_task.cancel()
|
||||
stderr_task.cancel()
|
||||
try:
|
||||
await stdout_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
try:
|
||||
await stderr_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
raise
|
||||
|
||||
returncode = process.returncode if process.returncode is not None else 0
|
||||
|
||||
logger.info(f"wait_and_read_ssh: stdout={len(stdout_data)} байт, stderr={len(stderr_data)} байт, returncode={returncode}")
|
||||
return stdout_data, stderr_data, returncode
|
||||
|
||||
|
||||
def read_pty_output(master_fd: int, timeout: float = 2.0) -> Tuple[str, bool]:
|
||||
"""
|
||||
Чтение вывода из PTY с таймаутом.
|
||||
Возвращает (вывод, завершён_ли_процесс).
|
||||
"""
|
||||
output = ""
|
||||
is_done = False
|
||||
total_waited = 0
|
||||
consecutive_errors = 0 # Счётчик последовательных ошибок
|
||||
MAX_ERRORS = 10 # Максимальное количество ошибок перед выходом
|
||||
|
||||
try:
|
||||
# Устанавливаем non-blocking режим
|
||||
flags = fcntl.fcntl(master_fd, fcntl.F_GETFL)
|
||||
fcntl.fcntl(master_fd, fcntl.F_SETFL, flags | os.O_NONBLOCK)
|
||||
|
||||
while total_waited < timeout:
|
||||
try:
|
||||
# Ждём данные с коротким таймаутом
|
||||
ready, _, _ = select.select([master_fd], [], [], 0.2)
|
||||
if ready:
|
||||
try:
|
||||
data = os.read(master_fd, 4096)
|
||||
if data:
|
||||
output += data.decode('utf-8', errors='replace')
|
||||
logger.debug(f"Прочитано из PTY: {len(data)} байт")
|
||||
# Сбрасываем таймер если есть данные
|
||||
total_waited = 0
|
||||
consecutive_errors = 0 # Сбрасываем счётчик ошибок
|
||||
else:
|
||||
# EOF - процесс завершился
|
||||
is_done = True
|
||||
break
|
||||
except BlockingIOError:
|
||||
# Нет данных, продолжаем ждать
|
||||
pass
|
||||
except OSError as e:
|
||||
# Ошибка чтения (например, EIO - процесс умер)
|
||||
logger.warning(f"OSError при чтении PTY: {e} (ошибка {e.errno})")
|
||||
consecutive_errors += 1
|
||||
if consecutive_errors >= MAX_ERRORS:
|
||||
logger.warning(f"Слишком много ошибок чтения PTY ({consecutive_errors}), считаем процесс завершённым")
|
||||
is_done = True
|
||||
break
|
||||
# При ошибке чтения сразу считаем что процесс завершился
|
||||
is_done = True
|
||||
break
|
||||
else:
|
||||
# Timeout - проверяем не завершился ли процесс
|
||||
try:
|
||||
_, status = os.waitpid(-1, os.WNOHANG)
|
||||
if status != 0:
|
||||
logger.debug(f"Процесс завершился со статусом: {status}")
|
||||
is_done = True
|
||||
break
|
||||
except ChildProcessError:
|
||||
# Процесс уже завершён
|
||||
is_done = True
|
||||
break
|
||||
|
||||
# Если уже что-то прочитали и есть запрос ввода - выходим
|
||||
if output and detect_input_type(output):
|
||||
logger.debug(f"Обнаружен запрос ввода")
|
||||
break
|
||||
|
||||
total_waited += 0.2
|
||||
|
||||
except OSError as e:
|
||||
# Ошибка select (например, Bad file descriptor)
|
||||
logger.warning(f"OSError при select PTY: {e}")
|
||||
is_done = True
|
||||
break
|
||||
except Exception as e:
|
||||
logger.debug(f"Ошибка при чтении PTY: {e}")
|
||||
consecutive_errors += 1
|
||||
if consecutive_errors >= MAX_ERRORS:
|
||||
is_done = True
|
||||
break
|
||||
total_waited += 0.2
|
||||
|
||||
except Exception as e:
|
||||
logger.debug(f"Ошибка чтения PTY: {e}")
|
||||
is_done = True
|
||||
|
||||
logger.debug(f"read_pty_output: output={len(output)} байт, is_done={is_done}")
|
||||
return output, is_done
|
||||
Reference in New Issue
Block a user