Lowa
3126811f09
- Enhance server stop mechanism with more robust socket and thread management - Add better handling of active transfers during server shutdown - Implement additional safety checks and timeout handling - Improve logging and error reporting for server stop process - Prevent potential deadlocks and resource leaks during server termination
388 lines
19 KiB
Python
388 lines
19 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
TFTP сервер для передачи прошивки с компьютера на коммутатор.
|
||
|
||
- Создает сервер по заданному IP и порту.
|
||
- Расшаривает папку Firmware.
|
||
- Показывает текущее состояние сервера и статус передачи файла:
|
||
- кому (IP устройства),
|
||
- сколько осталось байт,
|
||
- сколько передано байт,
|
||
- время передачи.
|
||
"""
|
||
|
||
import os
|
||
import socket
|
||
import struct
|
||
import threading
|
||
import time
|
||
|
||
class TFTPServer:
|
||
def __init__(self, share_folder):
|
||
"""
|
||
Инициализация TFTP сервера.
|
||
|
||
:param share_folder: Путь к папке, содержащей файлы (например, папка 'Firmware')
|
||
"""
|
||
self.share_folder = share_folder
|
||
self.log_callback = None
|
||
self.running = False
|
||
self.server_socket = None
|
||
self.lock = threading.Lock()
|
||
self.transfer_sockets = set() # Множество для хранения всех активных сокетов передачи
|
||
# Словарь активных передач для мониторинга их статуса.
|
||
# Ключ – адрес клиента, значение – словарь с информацией о передаче.
|
||
self.active_transfers = {}
|
||
|
||
def set_log_callback(self, callback):
|
||
"""
|
||
Установка функции обратного вызова для логирования сообщений.
|
||
|
||
:param callback: Функция, принимающая строку сообщения.
|
||
"""
|
||
self.log_callback = callback
|
||
|
||
def log(self, message):
|
||
"""
|
||
Функция логирования: вызывает callback (если задан) или выводит сообщение в консоль.
|
||
|
||
:param message: Строка с сообщением для логирования.
|
||
"""
|
||
if self.log_callback:
|
||
self.log_callback(message)
|
||
else:
|
||
print(message)
|
||
|
||
def start_server(self, ip, port):
|
||
"""
|
||
Запуск TFTP сервера на указанном IP и порту.
|
||
|
||
:param ip: IP-адрес для привязки сервера.
|
||
:param port: Порт для TFTP сервера.
|
||
"""
|
||
if self.running:
|
||
self.log("[WARN] Сервер уже запущен")
|
||
return
|
||
|
||
self.running = True
|
||
try:
|
||
self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
||
self.server_socket.bind((ip, port))
|
||
self.log(f"[INFO] TFTP сервер запущен на {ip}:{port}")
|
||
|
||
while self.running:
|
||
try:
|
||
self.server_socket.settimeout(1.0)
|
||
data, client_addr = self.server_socket.recvfrom(2048)
|
||
if data and self.running:
|
||
threading.Thread(target=self.handle_request, args=(data, client_addr), daemon=True).start()
|
||
except socket.timeout:
|
||
continue
|
||
except socket.error as e:
|
||
if self.running: # Логируем ошибку только если сервер еще запущен
|
||
self.log(f"[ERROR] Ошибка получения данных: {str(e)}")
|
||
break
|
||
except Exception as e:
|
||
if self.running: # Логируем ошибку только если сервер еще запущен
|
||
self.log(f"[ERROR] Ошибка получения данных: {str(e)}")
|
||
break
|
||
except Exception as e:
|
||
self.log(f"[ERROR] Ошибка запуска сервера: {str(e)}")
|
||
finally:
|
||
self.running = False
|
||
if self.server_socket:
|
||
try:
|
||
self.server_socket.close()
|
||
except:
|
||
pass
|
||
self.server_socket = None
|
||
self.log("[INFO] TFTP сервер остановлен.")
|
||
|
||
def stop_server(self):
|
||
"""
|
||
Остановка TFTP сервера.
|
||
"""
|
||
if not self.running:
|
||
return
|
||
|
||
self.log("[INFO] Остановка TFTP сервера...")
|
||
self.running = False
|
||
|
||
try:
|
||
# Закрываем основной сокет сервера первым
|
||
if self.server_socket:
|
||
try:
|
||
# Создаем временный сокет и отправляем пакет самому себе,
|
||
# чтобы разблокировать recvfrom
|
||
temp_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
||
try:
|
||
server_address = self.server_socket.getsockname()
|
||
temp_socket.sendto(b'', server_address)
|
||
except:
|
||
pass
|
||
finally:
|
||
try:
|
||
temp_socket.close()
|
||
except:
|
||
pass
|
||
|
||
try:
|
||
self.server_socket.close()
|
||
except Exception as e:
|
||
self.log(f"[WARN] Ошибка при закрытии основного сокета: {str(e)}")
|
||
finally:
|
||
self.server_socket = None
|
||
|
||
# Закрываем все активные сокеты передачи
|
||
with self.lock:
|
||
active_sockets = list(self.transfer_sockets)
|
||
self.transfer_sockets.clear()
|
||
active_transfers = dict(self.active_transfers)
|
||
self.active_transfers.clear()
|
||
|
||
# Закрываем сокеты передачи после очистки множества
|
||
for sock in active_sockets:
|
||
try:
|
||
if sock:
|
||
sock.close()
|
||
except Exception as e:
|
||
self.log(f"[WARN] Ошибка при закрытии сокета передачи: {str(e)}")
|
||
|
||
# Отправляем сообщения об остановке для активных передач
|
||
for client_addr, transfer_info in active_transfers.items():
|
||
try:
|
||
self.send_error(client_addr, 0, "Сервер остановлен")
|
||
except:
|
||
pass
|
||
|
||
except Exception as e:
|
||
self.log(f"[ERROR] Ошибка при остановке сервера: {str(e)}")
|
||
finally:
|
||
self.running = False # Гарантируем, что флаг running будет False
|
||
self.log("[INFO] TFTP сервер остановлен.")
|
||
|
||
def handle_request(self, data, client_addr):
|
||
"""
|
||
Обработка входящего запроса от клиента.
|
||
|
||
:param data: Полученные данные (UDP-пакет).
|
||
:param client_addr: Адрес клиента, отправившего пакет.
|
||
"""
|
||
if len(data) < 2:
|
||
self.log(f"[WARN] Получен некорректный пакет от {client_addr}")
|
||
return
|
||
opcode = struct.unpack("!H", data[:2])[0]
|
||
if opcode == 1: # RRQ (Read Request) – запрос на чтение файла
|
||
self.handle_rrq(data, client_addr)
|
||
else:
|
||
self.log(f"[WARN] Неподдерживаемый запрос (опкод {opcode}) от {client_addr}")
|
||
|
||
def handle_rrq(self, data, client_addr):
|
||
"""
|
||
Обработка запроса на чтение файла (RRQ).
|
||
|
||
:param data: Данные запроса.
|
||
:param client_addr: Адрес клиента.
|
||
"""
|
||
try:
|
||
# RRQ формата: 2 байта опкода, затем строка имени файла, за которой следует 0,
|
||
# затем строка режима (например, "octet"), и завершается 0.
|
||
parts = data[2:].split(b'\0')
|
||
if len(parts) < 2:
|
||
self.log(f"[WARN] Некорректный RRQ пакет от {client_addr}")
|
||
return
|
||
filename = parts[0].decode('utf-8')
|
||
mode = parts[1].decode('utf-8').lower()
|
||
self.log(f"[INFO] Получен RRQ от {client_addr}: файл '{filename}', режим '{mode}'")
|
||
if mode != "octet":
|
||
self.send_error(client_addr, 0, "Поддерживается только octet режим")
|
||
return
|
||
file_path = os.path.join(self.share_folder, filename)
|
||
if not os.path.isfile(file_path):
|
||
self.send_error(client_addr, 1, "Файл не найден")
|
||
return
|
||
# Запускаем передачу файла в новом потоке.
|
||
threading.Thread(target=self.send_file, args=(file_path, client_addr), daemon=True).start()
|
||
except Exception as e:
|
||
self.log(f"[ERROR] Ошибка обработки RRQ: {str(e)}")
|
||
|
||
def send_error(self, client_addr, error_code, error_message):
|
||
"""
|
||
Отправка сообщения об ошибке клиенту.
|
||
|
||
:param client_addr: Адрес клиента.
|
||
:param error_code: Код ошибки.
|
||
:param error_message: Текст ошибки.
|
||
"""
|
||
# Формируем TFTP пакет ошибки: 2 байта опкода (5), 2 байта кода ошибки, сообщение об ошибке и завершающий 0.
|
||
packet = struct.pack("!HH", 5, error_code) + error_message.encode('utf-8') + b'\0'
|
||
self.server_socket.sendto(packet, client_addr)
|
||
self.log(f"[INFO] Отправлено сообщение об ошибке '{error_message}' клиенту {client_addr}")
|
||
|
||
def send_file(self, file_path, client_addr):
|
||
"""
|
||
Передача файла клиенту по протоколу TFTP.
|
||
|
||
В процессе передачи обновляется состояние активной передачи, которое логируется с информацией:
|
||
- адрес клиента,
|
||
- имя файла,
|
||
- размер файла,
|
||
- количество переданных байт,
|
||
- оставшиеся байты,
|
||
- время передачи.
|
||
|
||
:param file_path: Полный путь к файлу для передачи.
|
||
:param client_addr: Адрес клиента.
|
||
"""
|
||
BLOCK_SIZE = 512
|
||
transfer_socket = None
|
||
try:
|
||
if not os.path.exists(file_path):
|
||
self.log(f"[ERROR] Файл '{file_path}' не существует")
|
||
self.send_error(client_addr, 1, "Файл не найден")
|
||
return
|
||
|
||
filesize = os.path.getsize(file_path)
|
||
if filesize == 0:
|
||
self.log(f"[ERROR] Файл '{file_path}' пуст")
|
||
self.send_error(client_addr, 0, "Файл пуст")
|
||
return
|
||
|
||
start_time = time.time()
|
||
file_basename = os.path.basename(file_path)
|
||
|
||
# Регистрируем активную передачу
|
||
with self.lock:
|
||
self.active_transfers[client_addr] = {
|
||
'filename': file_basename,
|
||
'filesize': filesize,
|
||
'bytes_sent': 0,
|
||
'start_time': start_time
|
||
}
|
||
|
||
self.log(f"[INFO] Начало передачи файла '{file_basename}' клиенту {client_addr}. Размер файла: {filesize} байт.")
|
||
|
||
# Создаем отдельный сокет для передачи файла
|
||
transfer_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
||
transfer_socket.settimeout(5.0) # Таймаут ожидания ACK
|
||
|
||
# Добавляем сокет в множество активных сокетов
|
||
with self.lock:
|
||
self.transfer_sockets.add(transfer_socket)
|
||
|
||
block_num = 1
|
||
bytes_sent = 0
|
||
last_progress_time = time.time()
|
||
|
||
try:
|
||
with open(file_path, "rb") as f:
|
||
while self.running: # Проверяем флаг running
|
||
try:
|
||
data_block = f.read(BLOCK_SIZE)
|
||
if not data_block: # Достигнут конец файла
|
||
break
|
||
|
||
# Проверяем флаг running перед отправкой блока
|
||
if not self.running:
|
||
raise Exception("Передача прервана: сервер остановлен")
|
||
|
||
# Формируем TFTP пакет данных
|
||
packet = struct.pack("!HH", 3, block_num) + data_block
|
||
attempts = 0
|
||
ack_received = False
|
||
|
||
# Попытка отправки текущего блока (до 3 повторных попыток)
|
||
while attempts < 3 and not ack_received and self.running:
|
||
if transfer_socket is None:
|
||
raise Exception("Сокет передачи закрыт")
|
||
|
||
try:
|
||
transfer_socket.sendto(packet, client_addr)
|
||
|
||
# Логируем прогресс каждую секунду
|
||
current_time = time.time()
|
||
if current_time - last_progress_time >= 1.0:
|
||
elapsed_time = current_time - start_time
|
||
remaining = filesize - bytes_sent
|
||
self.log(f"[STATUS] Клиент: {client_addr} | Файл: {file_basename} | "
|
||
f"Отправлено: {bytes_sent}/{filesize} байт | "
|
||
f"Осталось: {remaining} байт | "
|
||
f"Время: {elapsed_time:.2f} сек.")
|
||
last_progress_time = current_time
|
||
|
||
# Ожидаем подтверждение
|
||
ack_data, addr = transfer_socket.recvfrom(4)
|
||
if addr == client_addr:
|
||
ack_opcode, ack_block = struct.unpack("!HH", ack_data)
|
||
if ack_opcode == 4 and ack_block == block_num:
|
||
ack_received = True
|
||
bytes_sent += len(data_block)
|
||
with self.lock:
|
||
if client_addr in self.active_transfers:
|
||
self.active_transfers[client_addr]['bytes_sent'] = bytes_sent
|
||
else:
|
||
self.log(f"[WARN] Неверный ACK от {client_addr}. "
|
||
f"Ожидался блок {block_num}, получен {ack_block}.")
|
||
except socket.timeout:
|
||
attempts += 1
|
||
self.log(f"[WARN] Таймаут ожидания ACK для блока {block_num} "
|
||
f"от {client_addr}. Попытка {attempts+1}.")
|
||
except socket.error as e:
|
||
if not self.running:
|
||
raise Exception("Передача прервана: сервер остановлен")
|
||
self.log(f"[ERROR] Ошибка сокета при отправке блока {block_num}: {str(e)}")
|
||
attempts += 1
|
||
except Exception as e:
|
||
if not self.running:
|
||
raise Exception("Передача прервана: сервер остановлен")
|
||
self.log(f"[ERROR] Ошибка при отправке блока {block_num}: {str(e)}")
|
||
attempts += 1
|
||
|
||
if not ack_received:
|
||
raise Exception(f"Не удалось получить подтверждение для блока {block_num}")
|
||
|
||
block_num = (block_num + 1) % 65536
|
||
|
||
except Exception as e:
|
||
if not self.running:
|
||
raise Exception("Передача прервана: сервер остановлен")
|
||
raise
|
||
|
||
if bytes_sent == filesize:
|
||
elapsed_time = time.time() - start_time
|
||
self.log(f"[INFO] Передача файла '{file_basename}' клиенту {client_addr} "
|
||
f"завершена за {elapsed_time:.2f} сек. Всего отправлено {bytes_sent} байт.")
|
||
except Exception as e:
|
||
if not self.running:
|
||
self.log(f"[INFO] Передача файла '{file_basename}' клиенту {client_addr} прервана: сервер остановлен")
|
||
else:
|
||
self.log(f"[ERROR] Ошибка при передаче файла '{file_basename}' "
|
||
f"клиенту {client_addr}: {str(e)}")
|
||
raise
|
||
|
||
except Exception as e:
|
||
if not self.running:
|
||
self.log(f"[INFO] Передача файла прервана: сервер остановлен")
|
||
else:
|
||
self.log(f"[ERROR] Ошибка при передаче файла '{os.path.basename(file_path)}' "
|
||
f"клиенту {client_addr}: {str(e)}")
|
||
try:
|
||
self.send_error(client_addr, 0, str(e))
|
||
except:
|
||
pass
|
||
finally:
|
||
# Закрываем сокет передачи
|
||
if transfer_socket:
|
||
try:
|
||
with self.lock:
|
||
self.transfer_sockets.discard(transfer_socket)
|
||
transfer_socket.close()
|
||
transfer_socket = None
|
||
except:
|
||
pass
|
||
|
||
# Удаляем информацию о передаче
|
||
with self.lock:
|
||
if client_addr in self.active_transfers:
|
||
del self.active_transfers[client_addr] |