#!/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