ComConfigCopy/src/communication/protocols/tftp.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
from typing import Optional, Callable
import tftpy

from core.config import AppConfig
from core.exceptions import TFTPError
from core.events import event_bus, Event, EventTypes

class TFTPProtocol:
    """Протокол для работы с TFTP сервером."""
    
    def __init__(self):
        self._server: Optional[tftpy.TftpServer] = None
        self._client: Optional[tftpy.TftpClient] = None
        self._server_running = False
        self._root_dir = AppConfig.CONFIGS_DIR
        self._port = AppConfig.TFTP_PORT
        self._timeout = AppConfig.TFTP_TIMEOUT
        self._retries = AppConfig.TFTP_RETRIES
    
    def configure(self, root_dir: Optional[str] = None, port: Optional[int] = None,
                 timeout: Optional[int] = None, retries: Optional[int] = None) -> None:
        """
        Конфигурация TFTP сервера/клиента.
        
        Args:
            root_dir: Корневая директория для файлов
            port: Порт TFTP сервера
            timeout: Таймаут операций
            retries: Количество попыток
        """
        if root_dir is not None:
            self._root_dir = root_dir
        if port is not None:
            self._port = port
        if timeout is not None:
            self._timeout = timeout
        if retries is not None:
            self._retries = retries
    
    def start_server(self, host: str = "0.0.0.0") -> None:
        """
        Запуск TFTP сервера.
        
        Args:
            host: IP-адрес для прослушивания
            
        Raises:
            TFTPError: При ошибке запуска сервера
        """
        if self._server_running:
            return
            
        try:
            # Создаем серверный объект
            self._server = tftpy.TftpServer(self._root_dir)
            
            # Запускаем сервер в отдельном потоке
            self._server.listen(host, self._port, timeout=self._timeout)
            self._server_running = True
            
            event_bus.publish(Event(EventTypes.TFTP_SERVER_STARTED, {
                "host": host,
                "port": self._port,
                "root_dir": self._root_dir
            }))
            
        except Exception as e:
            raise TFTPError(f"Ошибка запуска TFTP сервера: {e}")
    
    def stop_server(self) -> None:
        """Остановка TFTP сервера."""
        if self._server and self._server_running:
            self._server.stop()
            self._server = None
            self._server_running = False
            event_bus.publish(Event(EventTypes.TFTP_SERVER_STOPPED, None))
    
    def upload_file(self, filename: str, host: str, remote_filename: Optional[str] = None,
                   progress_callback: Optional[Callable[[int], None]] = None) -> None:
        """
        Загрузка файла на удаленное устройство.
        
        Args:
            filename: Путь к локальному файлу
            host: IP-адрес устройства
            remote_filename: Имя файла на устройстве
            progress_callback: Функция обратного вызова для отслеживания прогресса
            
        Raises:
            TFTPError: При ошибке загрузки файла
        """
        if not os.path.exists(filename):
            raise TFTPError(f"Файл не найден: {filename}")
            
        try:
            # Создаем клиентский объект
            self._client = tftpy.TftpClient(
                host,
                self._port,
                options={"timeout": self._timeout, "retries": self._retries}
            )
            
            # Определяем имя удаленного файла
            if not remote_filename:
                remote_filename = os.path.basename(filename)
            
            event_bus.publish(Event(EventTypes.TFTP_TRANSFER_STARTED, {
                "operation": "upload",
                "local_file": filename,
                "remote_file": remote_filename,
                "host": host
            }))
            
            # Загружаем файл
            self._client.upload(
                remote_filename,
                filename,
                progress_callback
            )
            
            event_bus.publish(Event(EventTypes.TFTP_TRANSFER_COMPLETED, {
                "operation": "upload",
                "local_file": filename,
                "remote_file": remote_filename,
                "host": host
            }))
            
        except Exception as e:
            event_bus.publish(Event(EventTypes.TFTP_ERROR, str(e)))
            raise TFTPError(f"Ошибка загрузки файла: {e}")
    
    def download_file(self, remote_filename: str, host: str, local_filename: Optional[str] = None,
                     progress_callback: Optional[Callable[[int], None]] = None) -> None:
        """
        Загрузка файла с удаленного устройства.
        
        Args:
            remote_filename: Имя файла на устройстве
            host: IP-адрес устройства
            local_filename: Путь для сохранения файла
            progress_callback: Функция обратного вызова для отслеживания прогресса
            
        Raises:
            TFTPError: При ошибке загрузки файла
        """
        try:
            # Создаем клиентский объект
            self._client = tftpy.TftpClient(
                host,
                self._port,
                options={"timeout": self._timeout, "retries": self._retries}
            )
            
            # Определяем имя локального файла
            if not local_filename:
                local_filename = os.path.join(self._root_dir, remote_filename)
            
            event_bus.publish(Event(EventTypes.TFTP_TRANSFER_STARTED, {
                "operation": "download",
                "local_file": local_filename,
                "remote_file": remote_filename,
                "host": host
            }))
            
            # Загружаем файл
            self._client.download(
                remote_filename,
                local_filename,
                progress_callback
            )
            
            event_bus.publish(Event(EventTypes.TFTP_TRANSFER_COMPLETED, {
                "operation": "download",
                "local_file": local_filename,
                "remote_file": remote_filename,
                "host": host
            }))
            
        except Exception as e:
            event_bus.publish(Event(EventTypes.TFTP_ERROR, str(e)))
            raise TFTPError(f"Ошибка загрузки файла: {e}")
    
    @property
    def is_server_running(self) -> bool:
        """Проверка состояния сервера."""
        return self._server_running