pi-pm3/app/backend/workers/pm3_worker.py

"""Proxmark3 worker for executing commands."""
import asyncio
from dataclasses import dataclass
from typing import Optional
from pathlib import Path
import sys

from .. import config


@dataclass
class PM3CommandResult:
    """Result from a PM3 command execution."""
    success: bool
    output: str
    error: Optional[str] = None


class PM3Worker:
    """Worker for executing Proxmark3 commands using the built-in pm3 module."""

    def __init__(self, device_path: str = None):
        """Initialize the PM3 worker.

        Args:
            device_path: Path to PM3 device (default: from config)
        """
        self.device_path = device_path or config.PM3_DEVICE
        self._device = None
        self._connected = False
        self._lock = asyncio.Lock()
        self._pm3_module = None

    async def _import_pm3(self):
        """Import the pm3 module (lazy loading)."""
        if self._pm3_module is None:
            try:
                # The pm3 module should be available after pm3 client installation
                import pm3
                self._pm3_module = pm3
            except ImportError as e:
                raise ImportError(
                    "pm3 module not found. Ensure Proxmark3 client is installed with Python support. "
                    f"Error: {e}"
                )
        return self._pm3_module

    async def connect(self):
        """Connect to the Proxmark3 device."""
        async with self._lock:
            if self._connected:
                return

            try:
                pm3 = await self._import_pm3()

                # Run blocking pm3.open() in executor
                loop = asyncio.get_event_loop()
                self._device = await loop.run_in_executor(
                    None,
                    pm3.open,
                    self.device_path
                )
                self._connected = True
            except Exception as e:
                raise ConnectionError(f"Failed to connect to PM3 at {self.device_path}: {e}")

    async def disconnect(self):
        """Disconnect from the Proxmark3 device."""
        async with self._lock:
            if self._device:
                try:
                    # Close the device connection
                    # Note: pm3 module may not have explicit close, connection closes when object is deleted
                    self._device = None
                    self._connected = False
                except Exception as e:
                    print(f"Error disconnecting from PM3: {e}")

    async def is_connected(self) -> bool:
        """Check if connected to PM3 device."""
        if not self._connected:
            # Try to check if device exists
            return Path(self.device_path).exists()
        return self._connected

    async def execute_command(self, command: str) -> PM3CommandResult:
        """Execute a Proxmark3 command.

        Args:
            command: PM3 command to execute (e.g., "hw version")

        Returns:
            PM3CommandResult with success status, output, and optional error
        """
        async with self._lock:
            try:
                # Ensure we're connected
                if not self._connected:
                    await self.connect()

                if not self._device:
                    return PM3CommandResult(
                        success=False,
                        output="",
                        error="Not connected to PM3 device"
                    )

                # Execute command in executor (blocking call)
                loop = asyncio.get_event_loop()

                try:
                    # Use .cmd() method to execute command
                    output = await loop.run_in_executor(
                        None,
                        self._device.cmd,
                        command
                    )

                    # pm3.cmd() returns the output string
                    return PM3CommandResult(
                        success=True,
                        output=str(output) if output else "",
                        error=None
                    )
                except Exception as cmd_error:
                    return PM3CommandResult(
                        success=False,
                        output="",
                        error=f"Command execution failed: {cmd_error}"
                    )

            except Exception as e:
                return PM3CommandResult(
                    success=False,
                    output="",
                    error=str(e)
                )


# Mock PM3 Worker for testing without actual hardware
class MockPM3Worker(PM3Worker):
    """Mock PM3 worker for testing without hardware."""

    def __init__(self, device_path: str = None):
        super().__init__(device_path)
        self._mock_responses = {
            "hw version": "Proxmark3 RFID instrument\n  client: RRG/Iceman/master/v4.14831",
            "hw status": "Device: PM3 GENERIC\nBootrom: master/v4.14831\nOS: master/v4.14831",
            "hw tune": "Measuring antenna characteristics, please wait...\n# LF antenna: 50.00 V @ 125.00 kHz",
        }

    async def connect(self):
        """Mock connect."""
        self._connected = True

    async def disconnect(self):
        """Mock disconnect."""
        self._connected = False

    async def is_connected(self) -> bool:
        """Mock connection check."""
        return True

    async def execute_command(self, command: str) -> PM3CommandResult:
        """Mock command execution."""
        await asyncio.sleep(0.1)  # Simulate command delay

        # Return mock response if available
        for cmd_prefix, response in self._mock_responses.items():
            if command.startswith(cmd_prefix):
                return PM3CommandResult(
                    success=True,
                    output=response,
                    error=None
                )

        # Default response for unknown commands
        return PM3CommandResult(
            success=True,
            output=f"Mock response for: {command}",
            error=None
        )