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
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,
              callback: Callable[[str, str], None] | None = None) -> 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.
        callback: Optional progress callback (status, message).

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

    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.settimeout(timeout)
    # Allow reuse of local port (prevents "Address already in use")
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

    try:
        # ── Send WRQ to port 69 ──
        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) — discover 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(0), got opcode {opcode} from {server_addr}"

        srv_ip, srv_port = server_addr
        import sys
        print(f"[TFTP] ACK(0) received from {srv_ip}:{srv_port} "
              f"(file={total}B, blocks={-(-total // block_size)})", file=sys.stderr)

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

            if block % 100 == 0 or block == -(-total // block_size):
                pct = offset * 100 // total
                print(f"[TFTP] Block {block}/{-(-total // block_size)} ({pct}%) sent to {srv_ip}:{srv_port}",
                      file=sys.stderr)
                if callback:
                    callback("progress", f"{pct}%")

            # 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} (expected {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_get(host: str, remote_name: str, local_dir: str | Path,
              timeout: float = 30.0,
              callback: Callable[[str, str], None] | None = None,
              expected_size: int = 0) -> tuple[bool, str, Path]:
    """Download a file from a TFTP server using pure Python.

    Args:
        host: TFTP server IP address.
        remote_name: Remote filename.
        local_dir: Local directory to save the downloaded file.
        timeout: Socket timeout in seconds.
        callback: Optional progress callback (status, message).
        expected_size: Expected file size in bytes (for progress percentage).

    Returns:
        (success, message, local_path).
    """
    local_dir = Path(local_dir)
    local_dir.mkdir(parents=True, exist_ok=True)
    local_path = local_dir / remote_name

    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.settimeout(timeout)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

    try:
        # ── Send RRQ to port 69 ──
        rrq = struct.pack("!H", _TFTP_RRQ)
        rrq += remote_name.encode("ascii") + b"\x00"
        rrq += b"octet\x00"
        sock.sendto(rrq, (host, 69))

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

        srv_ip, srv_port = server_addr
        block = struct.unpack("!H", data_pkt[2:4])[0]
        chunk = data_pkt[4:]
        import sys
        print(f"[TFTP] DATA block {block} received from {srv_ip}:{srv_port} "
              f"({len(chunk)} bytes)", file=sys.stderr)

        # ── ACK this block with the actual block number ──
        file_data = bytearray(chunk)
        ack = struct.pack("!HH", _TFTP_ACK, block)
        sock.sendto(ack, server_addr)

        if callback:
            pct = len(file_data) * 100 // expected_size if expected_size > 0 else -1
            callback("progress", f"{pct}%")

        # ── Receive subsequent blocks while last chunk was full-size ──
        while len(chunk) == _TFTP_BLOCK_SIZE:
            try:
                data_pkt, _ = sock.recvfrom(516)
                if len(data_pkt) < 4:
                    return False, "Incomplete DATA packet", local_path
                opcode = struct.unpack("!H", data_pkt[:2])[0]
                if opcode == _TFTP_ERROR:
                    err_msg = data_pkt[4:].decode("ascii", errors="replace").rstrip("\x00")
                    return False, f"TFTP error: {err_msg}", local_path
                if opcode != _TFTP_DATA:
                    return False, f"Expected DATA, got opcode {opcode}", local_path

                new_block = struct.unpack("!H", data_pkt[2:4])[0]
                chunk = data_pkt[4:]
                file_data += chunk
                block = new_block

                if block % 100 == 0:
                    print(f"[TFTP] Block {block} ({len(file_data)} bytes) "
                          f"received from {srv_ip}:{srv_port}", file=sys.stderr)

                # ACK this block
                ack = struct.pack("!HH", _TFTP_ACK, block)
                sock.sendto(ack, server_addr)

                if callback:
                    pct = len(file_data) * 100 // expected_size if expected_size > 0 else -1
                    callback("progress", f"{pct}%")
            except socket.timeout:
                return False, f"Timeout waiting for DATA block {block + 1}", local_path

        # ── Write file ──
        total = len(file_data)
        local_path.write_bytes(bytes(file_data))
        return True, f"Downloaded {total} bytes in {block} blocks to {local_path}", local_path
    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,
                            callback=callback)
    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,
    expected_size: int = 0,
) -> TftpResult:
    """Download a file from the Zynq device via TFTP (pure Python).

    Args:
        config: Application configuration.
        remote_name: Remote filename on the TFTP server.
        local_dir: Local directory to save the file (default: project .tmp/tftp/).
        callback: Optional progress callback.
        expected_size: Expected file size in bytes (for progress percentage).

    Returns:
        TftpResult with download status and CRC.
    """
    if local_dir is None:
        # Use project-local temp dir for cross-platform compatibility
        local_dir = Path(__file__).resolve().parent.parent / ".tmp" / "tftp"
    local_dir = Path(local_dir)

    if callback:
        callback("start", f"Downloading {remote_name} from {config.zynq_ip}")

    # 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",
            remote_path=remote_name,
        )

    try:
        ok, msg, local_path = _tftp_get(config.zynq_ip, remote_name, local_dir,
                                        callback=callback, expected_size=expected_size)
    except OSError as e:
        return TftpResult(
            step="download", success=False,
            message=f"Socket error: {e}",
            remote_path=remote_name,
        )

    crc = _compute_crc32(local_path) if ok else 0

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

    return TftpResult(
        step="download", success=ok,
        message=msg,
        local_path=str(local_path), remote_path=remote_name,
        crc_local=crc,
    )


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