294 lines
12 KiB
Python
294 lines
12 KiB
Python
#!/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:
|
||
# Используем readany() для чтения доступных данных
|
||
while True:
|
||
try:
|
||
# readany() читает любые доступные данные
|
||
data = await asyncio.wait_for(process.stdout.readany(), 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.readany(), 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
|