MP-Server/macro_manager.py

# Macro management and execution with command sequence support

import copy
import json
import os
import shlex
import tempfile
import threading
import uuid
import pyautogui
import subprocess
import time
from PIL import Image
from typing import Optional


# --- Payload validation for execution ---------------------------------------
# Allowlist of key names pyautogui knows how to press. Anything not in this set
# is rejected before being handed to pyautogui so arbitrary strings can't be
# passed through. Populated from pyautogui when available.
VALID_KEYS = frozenset(getattr(pyautogui, "KEYBOARD_KEYS", []) or [])

# Upper bound (milliseconds) for a single wait step so a macro can't hang the
# executor forever.
MAX_WAIT_MS = 60000

# Only persist last_used to disk at most this often (seconds) to avoid rewriting
# the whole JSON file on every single execution.
LAST_USED_SAVE_INTERVAL = 5.0


class MacroManager:
    """Manages macro storage, migration, and execution with command sequences.

    Thread-safe: all reads/writes of ``self.macros`` and all disk writes are
    guarded by ``self._lock``. Public methods acquire the lock and delegate to
    ``*_locked`` internal helpers so we never re-acquire a non-reentrant lock.
    """

    def __init__(self, data_file: str, images_dir: str, app_dir: str):
        self.data_file = data_file
        self.images_dir = images_dir
        self.app_dir = app_dir
        self.macros = {}
        self._lock = threading.Lock()
        # Timestamp of the last last_used-triggered save (for debouncing).
        self._last_used_save_time = 0.0
        self.load_macros()

    def load_macros(self):
        """Load macros from JSON file and migrate if needed."""
        with self._lock:
            try:
                if os.path.exists(self.data_file):
                    with open(self.data_file, "r") as file:
                        self.macros = json.load(file)

                    # Migrate old format macros
                    migrated = False
                    for macro_id, macro in list(self.macros.items()):
                        if macro.get("type") != "sequence":
                            self.macros[macro_id] = self._migrate_macro(macro)
                            migrated = True

                    if migrated:
                        self._save_macros_locked()
                        print("Migrated macros to new command sequence format")

            except Exception as e:
                print(f"Error loading macros: {e}")
                self.macros = {}

    def _migrate_macro(self, old_macro: dict) -> dict:
        """Convert old macro format to new command sequence format."""
        if old_macro.get("type") == "sequence":
            return old_macro

        commands = []
        modifiers = old_macro.get("modifiers", {})

        if old_macro.get("type") == "text":
            # Build held keys list
            held_keys = []
            for mod in ["ctrl", "alt", "shift"]:
                if modifiers.get(mod):
                    held_keys.append(mod)

            if held_keys:
                # Use hotkey for modified text
                commands.append({
                    "type": "hotkey",
                    "keys": held_keys + [old_macro.get("command", "")]
                })
            else:
                # Plain text
                commands.append({
                    "type": "text",
                    "value": old_macro.get("command", "")
                })

            # Add enter if requested
            if modifiers.get("enter"):
                commands.append({"type": "key", "value": "enter"})

        elif old_macro.get("type") == "app":
            commands.append({
                "type": "app",
                "command": old_macro.get("command", "")
            })

        return {
            "name": old_macro.get("name", "Unnamed"),
            "type": "sequence",
            "commands": commands,
            "category": old_macro.get("category", ""),
            "image_path": old_macro.get("image_path", ""),
            "last_used": old_macro.get("last_used", 0)
        }

    def save_macros(self):
        """Save macros to JSON file (thread-safe, atomic)."""
        with self._lock:
            self._save_macros_locked()

    def _save_macros_locked(self):
        """Atomically write macros to disk. Caller must hold ``self._lock``.

        Writes to a temp file in the same directory then os.replace() over the
        target so a crash mid-write can't leave a truncated/corrupt JSON file.
        """
        try:
            target_dir = os.path.dirname(self.data_file) or "."
            os.makedirs(target_dir, exist_ok=True)
            fd, tmp_path = tempfile.mkstemp(
                prefix=".macros-", suffix=".tmp", dir=target_dir
            )
            try:
                with os.fdopen(fd, "w") as file:
                    json.dump(self.macros, file, indent=4)
                    file.flush()
                    os.fsync(file.fileno())
                os.replace(tmp_path, self.data_file)
            except Exception:
                # Clean up the temp file on any failure
                try:
                    os.remove(tmp_path)
                except OSError:
                    pass
                raise
        except Exception as e:
            print(f"Error saving macros: {e}")

    def get_sorted_macros(self, sort_by: str = "name"):
        """Get macros sorted by specified criteria."""
        with self._lock:
            macro_list = list(self.macros.items())

        if sort_by == "name":
            macro_list.sort(key=lambda x: x[1].get("name", "").lower())
        elif sort_by == "type":
            # Sort by first command type in sequence
            def get_first_type(macro):
                cmds = macro.get("commands", [])
                return cmds[0].get("type", "") if cmds else ""
            macro_list.sort(
                key=lambda x: (get_first_type(x[1]), x[1].get("name", "").lower())
            )
        elif sort_by == "recent":
            macro_list.sort(
                key=lambda x: (x[1].get("last_used", 0), x[1].get("name", "").lower()),
                reverse=True
            )

        return macro_list

    def filter_macros_by_tab(self, macro_list: list, tab_name: str):
        """Filter macros based on tab/category name."""
        if tab_name == "All":
            return macro_list

        filtered = []
        for macro_id, macro in macro_list:
            # Check category match
            if macro.get("category") == tab_name:
                filtered.append((macro_id, macro))
            # Check first command type match
            elif macro.get("commands"):
                first_type = macro["commands"][0].get("type", "").title()
                if first_type == tab_name:
                    filtered.append((macro_id, macro))

        return filtered

    def get_unique_tabs(self) -> list:
        """Get list of unique tabs based on categories."""
        tabs = ["All"]
        categories = set()

        with self._lock:
            for macro in self.macros.values():
                if macro.get("category"):
                    categories.add(macro["category"])

        for category in sorted(categories):
            if category and category not in tabs:
                tabs.append(category)

        return tabs

    def add_macro(
        self,
        name: str,
        commands: list,
        category: str = "",
        image_path: str = ""
    ) -> str:
        """Add a new macro with command sequence."""
        macro_id = str(uuid.uuid4())

        # Process image if provided (file IO; does not touch self.macros)
        image_path_reference = self._process_image(image_path) if image_path else ""

        # Create macro data (deep copy commands to avoid reference issues)
        macro_data = {
            "name": name,
            "type": "sequence",
            "commands": copy.deepcopy(commands),
            "category": category,
            "image_path": image_path_reference,
            "last_used": 0
        }

        with self._lock:
            self.macros[macro_id] = macro_data
            self._save_macros_locked()
        return macro_id

    def update_macro(
        self,
        macro_id: str,
        name: str,
        commands: list,
        category: str = "",
        image_path: Optional[str] = None
    ) -> bool:
        """Update an existing macro."""
        with self._lock:
            if macro_id not in self.macros:
                return False
            existing_image = self.macros[macro_id].get("image_path", "")
            existing_last_used = self.macros[macro_id].get("last_used", 0)

        # Keep old image or update with new one (file IO outside the lock)
        if image_path is not None:
            image_path_reference = self._process_image(image_path) if image_path else ""
        else:
            image_path_reference = existing_image

        new_data = {
            "name": name,
            "type": "sequence",
            "commands": copy.deepcopy(commands),
            "category": category,
            "image_path": image_path_reference,
            "last_used": existing_last_used
        }

        with self._lock:
            if macro_id not in self.macros:
                return False
            self.macros[macro_id] = new_data
            self._save_macros_locked()
        return True

    def _process_image(self, image_path: str) -> str:
        """Process and store an image for a macro."""
        if not image_path:
            return ""

        try:
            file_ext = os.path.splitext(image_path)[1].lower()
            unique_filename = f"{uuid.uuid4().hex}{file_ext}"
            dest_path = os.path.join(self.images_dir, unique_filename)

            # Resize image to max 256x256
            with Image.open(image_path) as img:
                img.thumbnail((256, 256))
                img.save(dest_path)

            return os.path.join("macro_images", unique_filename)
        except Exception as e:
            print(f"Error processing image: {e}")
            return ""

    def delete_macro(self, macro_id: str) -> bool:
        """Delete a macro."""
        with self._lock:
            if macro_id not in self.macros:
                return False

            macro = self.macros[macro_id]
            image_path = macro.get("image_path")
            del self.macros[macro_id]
            self._save_macros_locked()

        # Delete associated image file (outside the lock; pure file IO)
        if image_path:
            try:
                img_path = os.path.join(self.app_dir, image_path)
                if os.path.exists(img_path):
                    os.remove(img_path)
            except Exception as e:
                print(f"Error removing image file: {e}")

        return True

    def execute_macro(self, macro_id: str) -> bool:
        """Execute a macro by ID.

        The macro's commands are copied under the lock, then executed *outside*
        the lock so long-running wait/keyboard steps don't block other threads
        (web server, GUI). last_used is updated in memory and only persisted at
        most every LAST_USED_SAVE_INTERVAL seconds to avoid rewriting the whole
        file on every press.
        """
        with self._lock:
            if macro_id not in self.macros:
                return False

            macro = self.macros[macro_id]
            commands = copy.deepcopy(macro.get("commands", []))

            # Update last_used in memory; debounce the disk write.
            now = time.time()
            macro["last_used"] = now
            if now - self._last_used_save_time >= LAST_USED_SAVE_INTERVAL:
                self._last_used_save_time = now
                self._save_macros_locked()

        try:
            for cmd in commands:
                self._execute_command(cmd)
            return True
        except Exception as e:
            print(f"Error executing macro: {e}")
            return False

    @staticmethod
    def _valid_keys(keys) -> list:
        """Filter a list of key names to those pyautogui recognises."""
        result = []
        for k in keys:
            if not isinstance(k, str):
                continue
            name = k.strip().lower()
            if not name:
                continue
            if VALID_KEYS and name not in VALID_KEYS:
                print(f"Skipping invalid key: {name!r}")
                continue
            result.append(name)
        return result

    def _execute_command(self, cmd: dict):
        """Execute a single command from a sequence."""
        cmd_type = cmd.get("type", "")

        if cmd_type == "text":
            # Type text string using clipboard for Unicode support
            value = cmd.get("value", "")
            if value:
                try:
                    import pyperclip
                    # Save current clipboard, paste text, restore (optional)
                    pyperclip.copy(value)
                    pyautogui.hotkey("ctrl", "v")
                    time.sleep(0.05)  # Small delay after paste
                except ImportError:
                    # Fallback to typewrite for ASCII-only
                    if len(value) == 1:
                        pyautogui.press(value)
                    else:
                        pyautogui.typewrite(value, interval=0.02)

        elif cmd_type == "key":
            # Press a single key (validated against pyautogui's known keys)
            key = cmd.get("value", "")
            valid = self._valid_keys([key]) if key else []
            if valid:
                pyautogui.press(valid[0])

        elif cmd_type == "hotkey":
            # Press key combination
            keys = cmd.get("keys", [])
            if keys:
                # Ensure keys is a list, not a string
                if isinstance(keys, str):
                    keys = [k.strip().lower() for k in keys.split(",")]
                valid = self._valid_keys(keys)
                # Only fire if every requested key was valid; a partial hotkey
                # is worse than nothing.
                if valid and len(valid) == len([k for k in keys if str(k).strip()]):
                    # Small delay before hotkey for reliability on Windows
                    time.sleep(0.05)
                    pyautogui.hotkey(*valid, interval=0.05)

        elif cmd_type == "wait":
            # Delay in milliseconds (bounded, non-negative)
            try:
                ms = int(cmd.get("ms", 0))
            except (TypeError, ValueError):
                ms = 0
            if ms > 0:
                ms = min(ms, MAX_WAIT_MS)
                time.sleep(ms / 1000.0)

        elif cmd_type == "app":
            # Launch application.
            # SECURITY: never use shell=True — that would allow shell-metacharacter
            # injection (; && | $() chaining, redirection). Both branches below run
            # without a shell, so the command can only launch a program with args.
            command = cmd.get("command", "")
            if command:
                try:
                    if os.name == "nt":
                        # Windows: pass the string so CreateProcess parses it
                        # (correctly handling quoted paths like "C:\Program
                        # Files\app.exe"). shell=False means no cmd.exe, so no
                        # metacharacter chaining.
                        subprocess.Popen(command)
                    else:
                        # POSIX: split into an argv list; no shell involved.
                        args = shlex.split(command)
                        if args:
                            subprocess.Popen(args)
                except Exception as e:
                    print(f"Error launching app command: {e}")

    # Legacy API compatibility methods
    def add_macro_legacy(
        self,
        name: str,
        macro_type: str,
        command: str,
        category: str = "",
        modifiers: Optional[dict] = None,
        image_path: str = ""
    ) -> str:
        """Add macro using legacy format (auto-converts to sequence)."""
        if modifiers is None:
            modifiers = {"ctrl": False, "alt": False, "shift": False, "enter": False}

        # Build command sequence
        commands = []
        held_keys = []
        for mod in ["ctrl", "alt", "shift"]:
            if modifiers.get(mod):
                held_keys.append(mod)

        if macro_type == "text":
            if held_keys:
                commands.append({"type": "hotkey", "keys": held_keys + [command]})
            else:
                commands.append({"type": "text", "value": command})

            if modifiers.get("enter"):
                commands.append({"type": "key", "value": "enter"})
        elif macro_type == "app":
            commands.append({"type": "app", "command": command})

        return self.add_macro(name, commands, category, image_path)

    def get_macro(self, macro_id: str) -> Optional[dict]:
        """Get a macro by ID."""
        with self._lock:
            macro = self.macros.get(macro_id)
            return copy.deepcopy(macro) if macro is not None else None

    def get_all_macros(self) -> dict:
        """Get all macros."""
        with self._lock:
            return copy.deepcopy(self.macros)