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
import os

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:
                # Setup Python path for pm3 module
                # The pm3 module is in proxmark3/client/pyscripts
                # The _pm3.so is in proxmark3/client/experimental_lib/example_py

                # Try multiple possible locations for pyscripts
                # The pm3 module is in proxmark3/client/pyscripts
                possible_pyscripts = [
                    os.path.expanduser("~/.pm3/proxmark3/client/pyscripts"),  # Pi install location
                    "/home/dt/.pm3/proxmark3/client/pyscripts",  # Pi install (explicit path)
                    os.path.expanduser("~/.pm3/client/pyscripts"),  # Alternative structure
                    "/usr/local/share/proxmark3/pyscripts",  # System install
                    "/usr/share/proxmark3/pyscripts",  # System install alt
                    os.path.join(os.path.dirname(__file__), "../../../.pm3-test/proxmark3/client/pyscripts"),  # Dev build
                ]

                pm3_pyscripts = None
                for path in possible_pyscripts:
                    pm3_py = os.path.join(path, "pm3.py")
                    if os.path.exists(pm3_py):
                        pm3_pyscripts = path
                        break

                if not pm3_pyscripts:
                    raise ImportError(f"Could not find pm3.py in any of: {possible_pyscripts}")

                # experimental_lib is at the same level as pyscripts (client/experimental_lib)
                pm3_lib = os.path.join(os.path.dirname(pm3_pyscripts), "experimental_lib/example_py")

                # Add to Python path if not already there
                if pm3_pyscripts not in sys.path:
                    sys.path.insert(0, pm3_pyscripts)
                if pm3_lib not in sys.path:
                    sys.path.insert(0, pm3_lib)

                # Import pm3 module
                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_internal(self):
        """Internal connect without locking (caller must hold lock)."""
        if self._connected:
            return

        try:
            pm3 = await self._import_pm3()

            # Run blocking pm3.pm3() constructor in executor
            loop = asyncio.get_event_loop()
            self._device = await loop.run_in_executor(
                None,
                pm3.pm3,
                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 connect(self):
        """Connect to the Proxmark3 device."""
        async with self._lock:
            await self._connect_internal()

    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 (use internal method since we hold the lock)
                if not self._connected:
                    await self._connect_internal()

                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:
                    # Helper function to run command and get output in same executor call
                    def run_command():
                        try:
                            self._device.console(command)
                            output = self._device.grabbed_output
                            return output
                        except Exception as e:
                            raise Exception(f"Error in run_command: {e}, device type: {type(self._device)}")

                    # Execute command and get output
                    output = await loop.run_in_executor(None, run_command)

                    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",
            "hw led": "[+] LED control command executed",
        }

    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
        )


class SubprocessPM3Worker(PM3Worker):
    """PM3 worker using subprocess to call pm3 CLI directly.

    This is a fallback for when the Python SWIG bindings aren't working.
    It executes pm3 CLI commands via subprocess which is more robust.
    """

    def __init__(self, device_path: str = None):
        super().__init__(device_path)
        self._pm3_path = None
        self._find_pm3_executable()

    def _find_pm3_executable(self):
        """Find the pm3 executable."""
        possible_paths = [
            "/usr/local/bin/pm3",
            os.path.expanduser("~/.pm3/proxmark3/client/proxmark3"),
            "/home/dt/.pm3/proxmark3/client/proxmark3",
            "/usr/bin/pm3",
        ]

        for path in possible_paths:
            if os.path.exists(path) and os.access(path, os.X_OK):
                self._pm3_path = path
                break

        if not self._pm3_path:
            # Try to find in PATH
            import shutil
            pm3_in_path = shutil.which("pm3") or shutil.which("proxmark3")
            if pm3_in_path:
                self._pm3_path = pm3_in_path

    async def connect(self):
        """Check that pm3 executable exists and device is available."""
        async with self._lock:
            if self._connected:
                return

            if not self._pm3_path:
                raise ConnectionError("pm3 executable not found")

            if not Path(self.device_path).exists():
                raise ConnectionError(f"PM3 device not found at {self.device_path}")

            self._connected = True

    async def disconnect(self):
        """Mark as disconnected."""
        async with self._lock:
            self._connected = False

    async def is_connected(self) -> bool:
        """Check if device exists."""
        return Path(self.device_path).exists() if self.device_path else False

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

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

        Returns:
            PM3CommandResult with success status, output, and optional error
        """
        async with self._lock:
            try:
                if not self._pm3_path:
                    return PM3CommandResult(
                        success=False,
                        output="",
                        error="pm3 executable not found"
                    )

                # Build command: pm3 -p /dev/ttyACM0 -c "command"
                cmd = [
                    self._pm3_path,
                    "-p", self.device_path,
                    "-c", command
                ]

                # Run subprocess
                process = await asyncio.create_subprocess_exec(
                    *cmd,
                    stdout=asyncio.subprocess.PIPE,
                    stderr=asyncio.subprocess.PIPE
                )

                try:
                    stdout, stderr = await asyncio.wait_for(
                        process.communicate(),
                        timeout=30.0  # 30 second timeout
                    )
                except asyncio.TimeoutError:
                    process.kill()
                    return PM3CommandResult(
                        success=False,
                        output="",
                        error="Command timed out after 30 seconds"
                    )

                output = stdout.decode("utf-8", errors="replace")
                error_output = stderr.decode("utf-8", errors="replace")

                if process.returncode == 0:
                    return PM3CommandResult(
                        success=True,
                        output=output,
                        error=None
                    )
                else:
                    return PM3CommandResult(
                        success=False,
                        output=output,
                        error=error_output or f"Command failed with exit code {process.returncode}"
                    )

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