Zynq_Flasher/src/tftp_manager.py

"""TFTP file transfer manager for Zynq Flasher GUI.

Handles uploading, downloading, and CRC verification of files
via TFTP to the Zynq device using a pure-Python TFTP client.
"""

from __future__ import annotations

import binascii
import socket
import struct
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import Callable

from config_manager import Config
from ip_verifier import ping_ip


@dataclass
class TftpResult:
    """Result of a TFTP operation."""

    step: str
    success: bool
    message: str
    local_path: str = ""
    remote_path: str = ""
    crc_local: int = 0
    crc_remote: int = 0


ProgressCallback = Callable[[str, str], None] | None

# TFTP opcodes
_TFTP_RRQ = 1
_TFTP_WRQ = 2
_TFTP_DATA = 3
_TFTP_ACK = 4
_TFTP_ERROR = 5
_TFTP_BLOCK_SIZE = 512


def _compute_crc32(file_path: str | Path) -> int:
    """Compute CRC32 checksum of a file."""
    crc = 0
    with open(file_path, "rb") as f:
        while True:
            chunk = f.read(65536)
            if not chunk:
                break
            crc = binascii.crc32(chunk, crc)
    return crc & 0xFFFFFFFF


def _tftp_put(host: str, file_path: str, remote_name: str,
              timeout: float = 10.0) -> tuple[bool, str]:
    """Upload a file to a TFTP server using pure Python.

    Args:
        host: TFTP server IP address.
        file_path: Local file path.
        remote_name: Remote filename.
        timeout: Socket timeout in seconds.

    Returns:
        (success, message).
    """
    data = Path(file_path).read_bytes()
    total = len(data)

    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.settimeout(timeout)

    try:
        # ── Send WRQ ──
        wrq = struct.pack("!H", _TFTP_WRQ)
        wrq += remote_name.encode("ascii") + b"\x00"
        wrq += b"octet\x00"
        sock.sendto(wrq, (host, 69))

        # ── Wait for ACK(0) to get server's data port ──
        ack_data, server_addr = sock.recvfrom(516)
        if len(ack_data) < 4:
            return False, "Invalid ACK from server"
        opcode = struct.unpack("!H", ack_data[:2])[0]
        if opcode == _TFTP_ERROR:
            err_msg = ack_data[4:].decode("ascii", errors="replace").rstrip("\x00")
            return False, f"TFTP error: {err_msg}"
        if opcode != _TFTP_ACK:
            return False, f"Expected ACK, got opcode {opcode}"

        # ── Send data blocks ──
        block = 1
        offset = 0
        while offset < total:
            chunk = data[offset:offset + _TFTP_BLOCK_SIZE]
            offset += len(chunk)
            data_pkt = struct.pack("!HH", _TFTP_DATA, block) + chunk
            sock.sendto(data_pkt, server_addr)

            # Wait for ACK
            try:
                ack_data, _ = sock.recvfrom(516)
                ack_op = struct.unpack("!H", ack_data[:2])[0]
                ack_block = struct.unpack("!H", ack_data[2:4])[0]
                if ack_op != _TFTP_ACK or ack_block != block:
                    return False, f"Bad ACK: op={ack_op} block={ack_block}"
            except socket.timeout:
                return False, f"Timeout waiting for ACK block {block}"
            block += 1

        return True, f"Uploaded {total} bytes in {block - 1} blocks"
    finally:
        sock.close()


def tftp_upload(
    config: Config,
    file_path: str | Path,
    remote_name: str | None = None,
    callback: ProgressCallback = None,
) -> TftpResult:
    """Upload a file to the Zynq device via TFTP (pure Python).

    Args:
        config: Application configuration.
        file_path: Path to the local file to upload.
        remote_name: Remote filename (defaults to config's tftp_upload_name).
        callback: Optional progress callback.

    Returns:
        TftpResult with upload status.
    """
    remote_name = remote_name or config.tftp_upload_name
    file_path = Path(file_path).resolve()

    if not file_path.exists():
        return TftpResult(
            step="upload", success=False,
            message=f"File not found: {file_path}",
            local_path=str(file_path), remote_path=remote_name,
        )

    if callback:
        callback("start",
                 f"Uploading {file_path.name} -> {remote_name}")

    # Verify device is reachable
    if not ping_ip(config.zynq_ip, config.ping_count, config.ping_timeout):
        return TftpResult(
            step="upload", success=False,
            message=f"Device {config.zynq_ip} not reachable",
            local_path=str(file_path), remote_path=remote_name,
        )

    try:
        ok, msg = _tftp_put(config.zynq_ip, str(file_path), remote_name)
    except OSError as e:
        return TftpResult(
            step="upload", success=False,
            message=f"Socket error: {e}",
            local_path=str(file_path), remote_path=remote_name,
        )

    if callback:
        callback("complete" if ok else "error",
                 "Upload " + ("succeeded" if ok else "failed"))

    return TftpResult(
        step="upload", success=ok,
        message=msg,
        local_path=str(file_path), remote_path=remote_name,
        crc_local=_compute_crc32(file_path) if ok else 0,
    )


def tftp_download(
    config: Config,
    remote_name: str,
    local_dir: str | Path | None = None,
    callback: ProgressCallback = None,
) -> TftpResult:
    """Download a file from the Zynq device via TFTP (pure Python).

    Note: This requires the Zynq to act as a TFTP client (U-Boot tftpboot),
    which typically means we need a TFTP server on the PC.  For simplicity,
    this operation is not natively supported — use ping-based verification
    or UART log analysis instead.

    Args:
        config: Application configuration.
        remote_name: Remote filename on the TFTP server.
        local_dir: Local directory to save the file.
        callback: Optional progress callback.

    Returns:
        TftpResult with status (always not-implemented).
    """
    return TftpResult(
        step="download",
        success=False,
        message="TFTP download requires PC-side server (not implemented)",
        remote_path=remote_name,
    )


def tftp_upload_verify(
    config: Config,
    file_path: str | Path,
    remote_name: str | None = None,
    callback: ProgressCallback = None,
) -> TftpResult:
    """Upload a file via TFTP, download it back, and verify CRC.

    Args:
        config: Application configuration.
        file_path: Path to the local file to upload.
        remote_name: Remote filename (defaults to config's tftp_upload_name).
        callback: Optional progress callback.

    Returns:
        TftpResult with upload, download, and CRC verification status.
    """
    remote_name = remote_name or config.tftp_upload_name

    # Step 1: Upload
    upload_result = tftp_upload(config, file_path, remote_name, callback)
    if not upload_result.success:
        return upload_result

    # Step 2: Download for verification
    download_result = tftp_download(config, remote_name, callback=callback)
    if not download_result.success:
        return download_result

    # Step 3: Compare CRC
    original_crc = _compute_crc32(file_path)
    downloaded_crc = download_result.crc_remote

    verified = original_crc == downloaded_crc

    if callback:
        if verified:
            callback("complete", f"CRC verified: {original_crc:#010x} == {downloaded_crc:#010x}")
        else:
            callback("error", f"CRC mismatch: {original_crc:#010x} != {downloaded_crc:#010x}")

    # Step 4: Clean up downloaded temp file regardless of result
    if download_result.local_path:
        try:
            Path(download_result.local_path).unlink(missing_ok=True)
            if callback:
                callback("complete", "Temp file cleaned up")
        except OSError:
            if callback:
                callback("progress", f"Temp file not removed: {download_result.local_path}")

    return TftpResult(
        step="upload_verify",
        success=verified,
        message=(
            f"CRC verified: {original_crc:#010x}"
            if verified
            else f"CRC mismatch: {original_crc:#010x} vs {downloaded_crc:#010x}"
        ),
        local_path=str(file_path),
        remote_path=remote_name,
        crc_local=original_crc,
        crc_remote=downloaded_crc,
    )