宜博 / pixelphone

"""
Device & app management MCP tools for the Android MCP server.

Independent module (E3 / AgentPhone parity). Registered into the existing
FastMCP instance by android_mcp._build_server() via a single line:

    from device_tools import register_device_tools
    register_device_tools(mcp)

This is the execution-layer breadth补齐 for phoneuse.md's plan P1 (System API)
/ P2 (Intent): AgentPhone exposes a wide App-control + device-state surface
(open_app / list_packages / force_stop / battery / volume / location / lock),
while phoneuse previously had only list_event_apps. These are all thin wrappers
over standard Android shell surfaces — pm / am / dumpsys / settings / cmd — so
the Root cost is effectively zero.

  open_app           — monkey launcher intent (fallback: resolve-activity + am start)
  list_packages      — pm list packages [-3|-s] [filter]
  get_foreground_app — dumpsys activity topResumedActivity (fallback: window mCurrentFocus)
  force_stop_app     — am force-stop          (guarded: refuses data-bearing apps)
  clear_app_data     — pm clear               (guarded: denylist + explicit confirm)
  get_battery        — dumpsys battery
  get_volume         — cmd media_session volume --get
  set_volume         — cmd media_session volume --set
  get_location       — dumpsys location (last known fix)
  get_lock_status    — dumpsys window + dumpsys power (screen on / keyguard)
  vibrate            — cmd vibrator_manager synced -f oneshot <ms>

Permission model — mirrors ui_control_tools.py / get_ui_tree:
  * On-device (Termux): commands run via `su -c` (Termux uid lacks the needed
    privileges for am force-stop / pm clear / dumpsys).
  * Off-device (e.g. a Mac driving the phone over USB): commands run via
    `adb shell`, whose uid already holds those permissions, so `su` is
    unnecessary. Detection is automatic via the Termux path check, so the same
    module is both deployable on-device and testable from a host.

Safety (hard boundary): force_stop_app and clear_app_data refuse any package on
_PROTECTED_PACKAGES (WeChat / Alipay / Telegram / banking / the system) — those
carry user data whose loss is irreversible. Demonstrate destructive tools on a
no-data app such as the calculator (com.google.android.calculator). clear_app_data
additionally requires confirm=True. Read-only tools carry no allowlist.
"""

from __future__ import annotations

import json
import os
import re
import subprocess

import ui_guard
from pathlib import Path

# ── shell routing (on-device su  /  off-device adb) ──────────────────────────

_TIMEOUT = 15
_ON_DEVICE = Path("/data/data/com.termux/files/usr").exists()


def _adb_prefix() -> list[str]:
    """`adb shell` invocation, preferring the repo-bundled platform-tools."""
    repo_adb = Path(__file__).resolve().parent / "platform-tools" / "adb"
    adb = str(repo_adb) if repo_adb.exists() else "adb"
    cmd = [adb]
    serial = os.environ.get("ANDROID_SERIAL", "")
    if serial:
        cmd += ["-s", serial]
    cmd.append("shell")
    return cmd


def _run(cmd: str, timeout: int = _TIMEOUT) -> subprocess.CompletedProcess:
    """Run a single shell command on the device.

    On-device (Termux) we need root for am/pm/dumpsys, so we go through
    `su -c`. Off-device we drive the phone through `adb shell`, whose uid
    already holds those permissions — no `su` needed.
    """
    if ui_guard.check(cmd):   # ui.paused：拒绝动屏幕的命令
        return subprocess.CompletedProcess(cmd, 75, "", "UI paused (ui.paused)")
    if _ON_DEVICE:
        argv = ["su", "-c", cmd]
    else:
        argv = _adb_prefix() + [cmd]
    return subprocess.run(argv, capture_output=True, text=True, timeout=timeout)


# ── destructive-action safety boundary ───────────────────────────────────────

# Packages whose data must never be force-stopped or cleared by the agent.
# Substring match against the package name keeps banking/payment families covered.
_PROTECTED_PACKAGES = (
    "com.tencent.mm",            # WeChat
    "com.tencent.mobileqq",      # QQ
    "com.eg.android.AlipayGphone",  # Alipay
    "org.telegram",              # Telegram (any variant)
    "com.whatsapp",
    "com.android.",              # system components
    "com.google.android.gms",    # Play services
    "com.google.android.gsf",
    "android",                   # the framework package itself
)


def _is_protected(package: str) -> bool:
    pkg = (package or "").strip()
    if not pkg:
        return True
    low = pkg.lower()
    # Any obvious bank/pay/wallet keyword is treated as data-bearing.
    if any(kw in low for kw in ("bank", "pay", "wallet", "alipay")):
        return True
    return any(pkg == p or pkg.startswith(p) for p in _PROTECTED_PACKAGES)


# ── parsing helpers ──────────────────────────────────────────────────────────

_FOCUS_PKG_RE = re.compile(r"([A-Za-z0-9_.]+)/([A-Za-z0-9_.$]+)")
_BATTERY_INT_FIELDS = ("level", "scale", "status", "health", "temperature",
                       "voltage", "plugged")

_BATTERY_STATUS = {1: "unknown", 2: "charging", 3: "discharging",
                   4: "not_charging", 5: "full"}
_BATTERY_HEALTH = {1: "unknown", 2: "good", 3: "overheat", 4: "dead",
                   5: "over_voltage", 6: "unspecified_failure", 7: "cold"}

# stream name → AudioManager stream index (cmd media_session --stream)
_STREAMS = {"voice_call": 0, "system": 1, "ring": 2, "media": 3,
            "alarm": 4, "notification": 5}
_VOL_RE = re.compile(r"volume is (\d+) in range \[(\d+)\.\.(\d+)\]")
_LOC_RE = re.compile(
    r"last location=Location\[(\w+)\s+(-?\d+\.\d+),(-?\d+\.\d+)[^\]]*?hAcc=(\d+\.?\d*)")


def _parse_foreground(activity_dump: str, window_dump: str) -> dict:
    """Resolve the foreground package/activity from dumpsys output.

    Primary signal is topResumedActivity (activity manager); window
    mCurrentFocus is the fallback (it can name a non-app system window).
    """
    for line in activity_dump.splitlines():
        if "topResumedActivity" in line or "mResumedActivity" in line:
            m = _FOCUS_PKG_RE.search(line.split("u0", 1)[-1] if "u0" in line else line)
            if m:
                return {"package": m.group(1), "activity": m.group(2), "source": "topResumedActivity"}
    for line in window_dump.splitlines():
        if "mCurrentFocus" in line:
            m = _FOCUS_PKG_RE.search(line)
            if m:
                return {"package": m.group(1), "activity": m.group(2), "source": "mCurrentFocus"}
    return {}


def _read_volume(stream_idx: int) -> dict:
    r = _run(f"cmd media_session volume --stream {stream_idx} --get", timeout=8)
    m = _VOL_RE.search((r.stdout or "") + (r.stderr or ""))
    if not m:
        return {"error": (r.stderr or r.stdout or "could not read volume").strip()[:160]}
    return {"volume": int(m.group(1)), "min": int(m.group(2)), "max": int(m.group(3))}


# ── tool registration ────────────────────────────────────────────────────────

def register_device_tools(mcp) -> None:
    """Register the App-management & device-state tools onto an existing FastMCP."""

    # ── open_app ─────────────────────────────────────────────────────────────
    @mcp.tool()
    def open_app(package: str) -> str:
        """Launch an app by package name (brings it to the foreground).

        Uses the monkey launcher-intent path, falling back to resolving the
        LAUNCHER activity and starting it explicitly. Returns whether the app
        is foreground afterwards (per get_foreground_app).

        Args:
            package: the app's package name, e.g. 'com.google.android.calculator'.
        """
        try:
            pkg = (package or "").strip()
            if not pkg:
                return json.dumps({"ok": False, "error": "package is required"})
            r = _run(f"monkey -p {pkg} -c android.intent.category.LAUNCHER 1", timeout=12)
            launched = "Events injected: 1" in (r.stdout or "")
            method = "monkey"
            if not launched:
                # Fallback: resolve the launcher activity and start it directly.
                resolve = _run(
                    f"cmd package resolve-activity --brief -c "
                    f"android.intent.category.LAUNCHER {pkg}", timeout=8)
                comp = ""
                for line in (resolve.stdout or "").splitlines():
                    if "/" in line and pkg in line:
                        comp = line.strip()
                if comp:
                    s = _run(f"am start -n {comp}", timeout=10)
                    launched = "Error" not in (s.stderr or "") and s.returncode == 0
                    method = "am_start"
            fg = _parse_foreground(
                _run("dumpsys activity activities", timeout=8).stdout,
                _run("dumpsys window", timeout=8).stdout)
            return json.dumps({
                "ok": launched or fg.get("package") == pkg,
                "package": pkg,
                "method": method,
                "foreground": fg,
                "now_foreground": fg.get("package") == pkg,
            }, ensure_ascii=False)
        except subprocess.TimeoutExpired:
            return json.dumps({"ok": False, "error": "open_app timed out"})
        except Exception as e:
            return json.dumps({"ok": False, "error": str(e)})

    # ── list_packages ────────────────────────────────────────────────────────
    @mcp.tool()
    def list_packages(filter: str = "", kind: str = "all", limit: int = 100) -> str:
        """List installed package names via `pm list packages`.

        Args:
            filter: optional substring to match package names against.
            kind: 'all' (default), 'third_party' (-3), or 'system' (-s).
            limit: max package names to return in the sample (default 100, max 1000).
        """
        try:
            flag = {"third_party": " -3", "system": " -s"}.get(kind, "")
            safe = re.sub(r"[^\w.\-]", "", filter or "")
            cmd = f"pm list packages{flag}" + (f" {safe}" if safe else "")
            r = _run(cmd, timeout=12)
            if r.returncode != 0 and (r.stderr or "").strip():
                return json.dumps({"error": r.stderr.strip(), "cmd": cmd})
            pkgs = [l.split(":", 1)[1].strip() for l in (r.stdout or "").splitlines()
                    if l.startswith("package:")]
            limit = min(max(1, limit), 1000)
            return json.dumps({
                "kind": kind,
                "filter": filter,
                "count": len(pkgs),
                "packages": sorted(pkgs)[:limit],
            }, ensure_ascii=False)
        except subprocess.TimeoutExpired:
            return json.dumps({"error": "list_packages timed out"})
        except Exception as e:
            return json.dumps({"error": str(e)})

    # ── get_foreground_app ───────────────────────────────────────────────────
    @mcp.tool()
    def get_foreground_app() -> str:
        """Return the package/activity of the app currently in the foreground.

        Reads topResumedActivity from `dumpsys activity activities`, falling
        back to mCurrentFocus from `dumpsys window`.
        """
        try:
            fg = _parse_foreground(
                _run("dumpsys activity activities", timeout=8).stdout,
                _run("dumpsys window", timeout=8).stdout)
            if not fg:
                return json.dumps({"ok": False, "error": "could not determine foreground app"})
            return json.dumps({"ok": True, **fg}, ensure_ascii=False)
        except subprocess.TimeoutExpired:
            return json.dumps({"ok": False, "error": "get_foreground_app timed out"})
        except Exception as e:
            return json.dumps({"ok": False, "error": str(e)})

    # ── force_stop_app ───────────────────────────────────────────────────────
    @mcp.tool()
    def force_stop_app(package: str) -> str:
        """Force-stop a running app (`am force-stop`).

        SAFETY: refuses data-bearing / system packages (WeChat, Alipay,
        Telegram, banking, com.android.*, the framework). Use a no-data app
        such as com.google.android.calculator to exercise this safely.

        Args:
            package: the app's package name to stop.
        """
        try:
            pkg = (package or "").strip()
            if _is_protected(pkg):
                return json.dumps({"ok": False, "refused": True, "package": pkg,
                                   "reason": "protected data-bearing or system package"})
            r = _run(f"am force-stop {pkg}", timeout=10)
            ok = r.returncode == 0 and "Error" not in (r.stderr or "")
            return json.dumps({"ok": ok, "package": pkg, "action": "force-stop",
                               "stderr": (r.stderr or "").strip()[:160]}, ensure_ascii=False)
        except subprocess.TimeoutExpired:
            return json.dumps({"ok": False, "error": "force_stop_app timed out"})
        except Exception as e:
            return json.dumps({"ok": False, "error": str(e)})

    # ── clear_app_data ───────────────────────────────────────────────────────
    @mcp.tool()
    def clear_app_data(package: str, confirm: bool = False) -> str:
        """Clear all of an app's data & cache (`pm clear`). DESTRUCTIVE.

        SAFETY: refuses data-bearing / system packages (same denylist as
        force_stop_app) AND requires confirm=True. This wipes accounts,
        databases, and settings irreversibly. Only ever run on a no-data app.

        Args:
            package: the app's package name to clear.
            confirm: must be True to proceed (guards against accidental wipes).
        """
        try:
            pkg = (package or "").strip()
            if _is_protected(pkg):
                return json.dumps({"ok": False, "refused": True, "package": pkg,
                                   "reason": "protected data-bearing or system package"})
            if not confirm:
                return json.dumps({"ok": False, "refused": True, "package": pkg,
                                   "reason": "destructive: pass confirm=True to clear app data"})
            r = _run(f"pm clear {pkg}", timeout=12)
            ok = "Success" in (r.stdout or "")
            return json.dumps({"ok": ok, "package": pkg, "action": "clear",
                               "stdout": (r.stdout or "").strip()[:120],
                               "stderr": (r.stderr or "").strip()[:160]}, ensure_ascii=False)
        except subprocess.TimeoutExpired:
            return json.dumps({"ok": False, "error": "clear_app_data timed out"})
        except Exception as e:
            return json.dumps({"ok": False, "error": str(e)})

    # ── get_battery ──────────────────────────────────────────────────────────
    @mcp.tool()
    def get_battery() -> str:
        """Return battery level, charging status, health, temperature and voltage.

        Parses `dumpsys battery`.
        """
        try:
            r = _run("dumpsys battery", timeout=8)
            fields: dict = {}
            for line in (r.stdout or "").splitlines():
                if ":" not in line:
                    continue
                k, _, v = line.strip().partition(":")
                k = k.strip().lower().replace(" ", "_")
                v = v.strip()
                if k in _BATTERY_INT_FIELDS:
                    try:
                        fields[k] = int(v)
                    except ValueError:
                        fields[k] = v
                elif k in ("ac_powered", "usb_powered", "wireless_powered", "dock_powered"):
                    fields[k] = v.lower() == "true"
                elif k == "technology":
                    fields[k] = v
            out = {
                "level": fields.get("level"),
                "scale": fields.get("scale", 100),
                "status": _BATTERY_STATUS.get(fields.get("status"), fields.get("status")),
                "health": _BATTERY_HEALTH.get(fields.get("health"), fields.get("health")),
                "charging": bool(fields.get("ac_powered") or fields.get("usb_powered")
                                 or fields.get("wireless_powered") or fields.get("dock_powered")),
                "ac_powered": fields.get("ac_powered"),
                "usb_powered": fields.get("usb_powered"),
                "wireless_powered": fields.get("wireless_powered"),
                "temperature_c": (fields["temperature"] / 10.0)
                                 if isinstance(fields.get("temperature"), int) else None,
                "voltage_mv": fields.get("voltage"),
                "technology": fields.get("technology"),
            }
            return json.dumps({"ok": out["level"] is not None, **out}, ensure_ascii=False)
        except subprocess.TimeoutExpired:
            return json.dumps({"ok": False, "error": "get_battery timed out"})
        except Exception as e:
            return json.dumps({"ok": False, "error": str(e)})

    # ── get_volume ───────────────────────────────────────────────────────────
    @mcp.tool()
    def get_volume(stream: str = "") -> str:
        """Read current volume levels for audio streams.

        Args:
            stream: one of voice_call/system/ring/media/alarm/notification.
                    Empty (default) returns all streams. Note Android reports
                    the active-output-device volume here.
        """
        try:
            if stream:
                idx = _STREAMS.get(stream)
                if idx is None:
                    return json.dumps({"ok": False, "error": f"unknown stream {stream!r}",
                                       "valid": list(_STREAMS)})
                return json.dumps({"ok": True, "stream": stream, **_read_volume(idx)},
                                  ensure_ascii=False)
            streams = {name: _read_volume(idx) for name, idx in _STREAMS.items()}
            return json.dumps({"ok": True, "streams": streams}, ensure_ascii=False)
        except subprocess.TimeoutExpired:
            return json.dumps({"ok": False, "error": "get_volume timed out"})
        except Exception as e:
            return json.dumps({"ok": False, "error": str(e)})

    # ── set_volume ───────────────────────────────────────────────────────────
    @mcp.tool()
    def set_volume(level: int, stream: str = "media") -> str:
        """Set the volume of an audio stream and return the read-back value.

        Args:
            level: target volume index (clamped to the stream's [min..max]).
            stream: one of voice_call/system/ring/media/alarm/notification
                    (default 'media'). Note: on some Android versions --set
                    targets the default output device, so the active-device
                    read-back may differ.
        """
        try:
            idx = _STREAMS.get(stream)
            if idx is None:
                return json.dumps({"ok": False, "error": f"unknown stream {stream!r}",
                                   "valid": list(_STREAMS)})
            before = _read_volume(idx)
            lo, hi = before.get("min", 0), before.get("max", 15)
            level = max(lo, min(int(level), hi))
            _run(f"cmd media_session volume --stream {idx} --set {level}", timeout=8)
            after = _read_volume(idx)
            return json.dumps({
                "ok": "error" not in after,
                "stream": stream,
                "requested": level,
                "before": before.get("volume"),
                "after": after.get("volume"),
                "range": {"min": lo, "max": hi},
            }, ensure_ascii=False)
        except subprocess.TimeoutExpired:
            return json.dumps({"ok": False, "error": "set_volume timed out"})
        except Exception as e:
            return json.dumps({"ok": False, "error": str(e)})

    # ── get_location ─────────────────────────────────────────────────────────
    @mcp.tool()
    def get_location() -> str:
        """Return the device's last known location fix (latitude/longitude).

        Parses the most accurate `last location=...` line from
        `dumpsys location`. No active fix is requested (read-only).
        """
        try:
            r = _run("dumpsys location", timeout=10)
            best = None
            for m in _LOC_RE.finditer(r.stdout or ""):
                acc = float(m.group(4))
                cand = {"provider": m.group(1), "latitude": float(m.group(2)),
                        "longitude": float(m.group(3)), "accuracy_m": acc}
                if best is None or acc < best["accuracy_m"]:
                    best = cand
            if not best:
                return json.dumps({"ok": False,
                                   "error": "no last-known location available"})
            return json.dumps({"ok": True, **best}, ensure_ascii=False)
        except subprocess.TimeoutExpired:
            return json.dumps({"ok": False, "error": "get_location timed out"})
        except Exception as e:
            return json.dumps({"ok": False, "error": str(e)})

    # ── get_lock_status ──────────────────────────────────────────────────────
    @mcp.tool()
    def get_lock_status() -> str:
        """Return whether the screen is on and whether the keyguard is locked.

        Reads mWakefulness from `dumpsys power` and the keyguard-showing flags
        from `dumpsys window`.
        """
        try:
            power = _run("dumpsys power", timeout=8).stdout or ""
            window = _run("dumpsys window", timeout=8).stdout or ""
            wm = re.search(r"mWakefulness=(\w+)", power)
            wakefulness = wm.group(1) if wm else ""
            screen_on = wakefulness.lower() == "awake"
            locked = None
            km = re.search(r"(?:isKeyguardShowing|mShowingDream|KeyguardShowing)=(\w+)", window)
            ks = re.search(r"mDreamingLockscreen=(\w+)", window)
            if km:
                locked = km.group(1).lower() == "true"
            elif ks:
                locked = ks.group(1).lower() == "true"
            return json.dumps({
                "ok": True,
                "screen_on": screen_on,
                "wakefulness": wakefulness,
                "locked": locked,
            }, ensure_ascii=False)
        except subprocess.TimeoutExpired:
            return json.dumps({"ok": False, "error": "get_lock_status timed out"})
        except Exception as e:
            return json.dumps({"ok": False, "error": str(e)})

    # ── vibrate ──────────────────────────────────────────────────────────────
    @mcp.tool()
    def vibrate(duration_ms: int = 200) -> str:
        """Trigger a one-shot vibration for duration_ms milliseconds.

        Uses `cmd vibrator_manager synced -f oneshot <ms>` (the -f forces the
        effect even under do-not-disturb / vibration-off settings).

        Args:
            duration_ms: vibration length in milliseconds (default 200, max 5000).
        """
        try:
            d = max(1, min(int(duration_ms), 5000))
            r = _run(f"cmd vibrator_manager synced -f oneshot {d}", timeout=8)
            err = (r.stderr or "").strip()
            ok = r.returncode == 0 and "Exception" not in err
            return json.dumps({"ok": ok, "duration_ms": d,
                               "error": err[:160] if not ok else None}, ensure_ascii=False)
        except subprocess.TimeoutExpired:
            return json.dumps({"ok": False, "error": "vibrate timed out"})
        except Exception as e:
            return json.dumps({"ok": False, "error": str(e)})