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.
"""

from __future__ import annotations

import binascii
import os
import platform
import shutil
import subprocess
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


def _compute_crc32(file_path: str | Path) -> int:
    """Compute CRC32 checksum of a file.

    Args:
        file_path: Path to the file.

    Returns:
        CRC32 value as unsigned integer.
    """
    import binascii

    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 _find_tftp_client() -> str | None:
    """Find the tftp client executable.

    Uses shutil.which() for cross-platform compatibility.

    Returns:
        Path to tftp executable, or None if not found.
    """
    return shutil.which("tftp")


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.

    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 first
    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,
        )

    tftp_cmd = _find_tftp_client()
    if not tftp_cmd:
        return TftpResult(
            step="upload",
            success=False,
            message="tftp client not found",
            local_path=str(file_path),
            remote_path=remote_name,
        )

    if callback:
        callback("progress", f"Running: {tftp_cmd}")

    # Build TFTP command
    cmd = [tftp_cmd, "-m", "binary", config.zynq_ip, "put", str(file_path), remote_name]

    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=60,
        )
        success = result.returncode == 0
        output = (result.stdout + result.stderr).strip()

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

        return TftpResult(
            step="upload",
            success=success,
            message=output or ("Upload succeeded" if success else "Upload failed"),
            local_path=str(file_path),
            remote_path=remote_name,
            crc_local=_compute_crc32(file_path),
        )
    except subprocess.TimeoutExpired:
        if callback:
            callback("error", "Upload timed out")
        return TftpResult(
            step="upload",
            success=False,
            message="Upload timed out",
            local_path=str(file_path),
            remote_path=remote_name,
            crc_local=_compute_crc32(file_path),
        )
    except OSError as e:
        return TftpResult(
            step="upload",
            success=False,
            message=f"OS error: {e}",
            local_path=str(file_path),
            remote_path=remote_name,
        )


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.

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

    Returns:
        TftpResult with download status.
    """
    local_dir = Path(local_dir or tempfile.gettempdir())
    local_path = local_dir / remote_name

    if callback:
        callback("start", f"Downloading {remote_name} -> {local_path}")

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

    tftp_cmd = _find_tftp_client()
    if not tftp_cmd:
        return TftpResult(
            step="download",
            success=False,
            message="tftp client not found",
            local_path=str(local_path),
            remote_path=remote_name,
        )

    cmd = [tftp_cmd, "-m", "binary", config.zynq_ip, "get", remote_name, str(local_path)]

    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=60,
        )
        success = result.returncode == 0
        output = (result.stdout + result.stderr).strip()

        if callback:
            callback("complete" if success else "error",
                     "Download " + ("succeeded" if success else "failed"))

        crc = _compute_crc32(local_path) if success and local_path.exists() else 0

        return TftpResult(
            step="download",
            success=success,
            message=output or ("Download succeeded" if success else "Download failed"),
            local_path=str(local_path),
            remote_path=remote_name,
            crc_remote=crc,
        )
    except subprocess.TimeoutExpired:
        return TftpResult(
            step="download",
            success=False,
            message="Download timed out",
            local_path=str(local_path),
            remote_path=remote_name,
        )
    except OSError as e:
        return TftpResult(
            step="download",
            success=False,
            message=f"OS error: {e}",
            local_path=str(local_path),
            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,
    )