# -*- coding: utf-8 -*-
"""修 Android 端审计出来的问题（第二版：修掉第一版脚本自身的 bug）。

  严重
   1. 4 个按钮永久失效（动作名与服务端对不上）
   2. RemoteClient 读取循环 NPE（未捕获会杀掉整个进程）
   3. Json 解析无深度上限（局域网任意主机可远程触发栈溢出）
   4. 对话框没判 isFinishing（退出后弹窗崩）
   5. readAll 无体积上限（大文件 OOM，而 OOM 是 Error 抓不住）
   6. moveTabIndicator 无限自我重投（主线程空转、发烫）
  中等
   7. connect 握手失败不 close（socket 泄漏 + 误报已连接）
   8. pids 越界
   9. 断线不清功能缓存
"""
from __future__ import annotations

import sys
from pathlib import Path

SRC = Path(r"C:\Users\shu\Desktop\deepseekagent\LanRemote\android\src"
           r"\com\lanremote\watch")
MAIN = SRC / "MainActivity.java"
CLIENT = SRC / "RemoteClient.java"
JSON = SRC / "Json.java"

FIXES = [
    # 1. 动作名
    (MAIN, 'media("play_pause")', 'media("playpause")', "播放/暂停 动作名"),
    (MAIN, 'media("vol_up")', 'media("volup")', "音量+ 动作名"),
    (MAIN, 'media("vol_down")', 'media("voldown")', "音量- 动作名"),
    (MAIN, 'confirmPower("reboot", "重启")', 'confirmPower("restart", "重启")',
     "重启 动作名"),

    # 2. 读取循环 NPE + 单行上限
    (CLIENT,
     '''        try {
            String line;
            while (running && (line = reader.readLine()) != null) {''',
     '''        try {
            // close() 会把 reader 置 null，而 running 判断和 readLine 之间有竞态。
            // 先取到局部变量并判空 —— 读取线程里的未捕获 NPE 会杀掉整个进程。
            BufferedReader localReader = this.reader;
            if (localReader == null) {
                return;
            }
            String line;
            while (running && (line = localReader.readLine()) != null) {
                if (line.length() > MAX_LINE_CHARS) {
                    continue;      // 对端发超长单行，丢掉防止占满内存
                }''',
     "读取循环 NPE + 单行上限"),

    # 3. connect 失败要 close
    (CLIENT,
     '''        Map<String, Object> reply = call(handshake, timeoutMs);
        if (Json.bool(reply, "ok", false)) {''',
     '''        Map<String, Object> reply;
        try {
            reply = call(handshake, timeoutMs);
        } catch (IOException failed) {
            // 握手失败必须关掉，否则 socket 和读取线程泄漏，
            // 而且 isConnected() 一直是 true，界面以为已经连上了。
            close();
            throw failed;
        }
        if (Json.bool(reply, "ok", false)) {''',
     "connect 失败 close"),

    # 4. Json 深度上限（字段）
    (JSON,
     '''    private static final class Parser {
        private final String src;
        private int pos;
''',
     '''    private static final class Parser {
        /** 最大嵌套深度。不加限制的话，一行 "[[[[…" 就能把栈撑爆，
         *  而 UDP 发现的入口是局域网里任何主机都能打的，等于远程 DoS。 */
        private static final int MAX_DEPTH = 64;

        private final String src;
        private int pos;
        private int depth;
''',
     "Json 深度字段"),

    # 5. Json 深度检查
    (JSON,
     '''        Object readValue() {
            if (atEnd()) {
                throw new IllegalArgumentException("JSON 意外结束");
            }
            char c = src.charAt(pos);''',
     '''        Object readValue() {
            if (atEnd()) {
                throw new IllegalArgumentException("JSON 意外结束");
            }
            if (depth > MAX_DEPTH) {
                throw new IllegalArgumentException("JSON 嵌套太深");
            }
            char c = src.charAt(pos);''',
     "Json 深度检查"),

    # 6. Json 深度计数：进 readObject / readArray 各加一次
    (JSON,
     '''        private Map<String, Object> readObject() {
            Map<String, Object> map = new LinkedHashMap<String, Object>();
            pos++;                       // '{'
''',
     '''        private Map<String, Object> readObject() {
            Map<String, Object> map = new LinkedHashMap<String, Object>();
            pos++;                       // '{'
            depth++;
''',
     "Json readObject 深度+"),

    # 7. 对话框判 isFinishing
    (MAIN,
     '''    private void showChooser(String title, List<String> items, final Chooser chooser) {
        final String[] array = items.toArray(new String[0]);''',
     '''    private void showChooser(String title, List<String> items, final Chooser chooser) {
        if (isFinishing() || isDestroyed()) {
            return;      // 界面已经关了就别弹，否则 BadTokenException
        }
        final String[] array = items.toArray(new String[0]);''',
     "showChooser 判 isFinishing"),

    # 8. readAll 体积上限
    (MAIN,
     '''    private byte[] readAll(Uri uri) {
        InputStream input = null;
        ByteArrayOutputStream buffer = new ByteArrayOutputStream();
        try {
            input = getContentResolver().openInputStream(uri);
            if (input == null) {
                return null;
            }
            byte[] chunk = new byte[64 * 1024];
            int read;
            while ((read = input.read(chunk)) > 0) {
                buffer.write(chunk, 0, read);
            }''',
     '''    /** 单个文件最大 32MB。整文件读进内存再 Base64 会放大好几倍，
     *  选个大视频就直接 OOM —— 而 OOM 是 Error，catch(IOException) 抓不住。 */
    private static final int MAX_UPLOAD_BYTES = 32 * 1024 * 1024;

    private byte[] readAll(Uri uri) {
        InputStream input = null;
        ByteArrayOutputStream buffer = new ByteArrayOutputStream();
        try {
            input = getContentResolver().openInputStream(uri);
            if (input == null) {
                return null;
            }
            byte[] chunk = new byte[64 * 1024];
            int read;
            while ((read = input.read(chunk)) > 0) {
                if (buffer.size() + read > MAX_UPLOAD_BYTES) {
                    return null;       // 超限，上层会提示文件太大
                }
                buffer.write(chunk, 0, read);
            }''',
     "readAll 体积上限"),

    # 9. moveTabIndicator 有限重试
    (MAIN,
     '''        if (target.getWidth() == 0) {
            // 布局还没完成，等一帧再试
            tabIndicator.post(new Runnable() {
                @Override
                public void run() {
                    moveTabIndicator(target);
                }
            });
            return;
        }''',
     '''        if (target.getWidth() == 0) {
            // 布局还没完成，等一帧再试 —— 但必须有限次。
            // 面板被 GONE 掉时父容器不再 layout，宽度永远是 0，
            // 无条件重投会让主线程空转，100% CPU 发烫耗电。
            if (indicatorRetries >= 5) {
                indicatorRetries = 0;
                return;
            }
            indicatorRetries++;
            tabIndicator.post(new Runnable() {
                @Override
                public void run() {
                    moveTabIndicator(target);
                }
            });
            return;
        }
        indicatorRetries = 0;''',
     "moveTabIndicator 有限重试"),

    (MAIN,
     "    private View tabIndicator;",
     "    private View tabIndicator;\n    private int indicatorRetries = 0;",
     "indicatorRetries 字段"),

    # 10. pids 越界
    (MAIN,
     '''                                public void onPick(int index) {
                                    confirmKill(names.get(index), pids.get(index));
                                }''',
     '''                                public void onPick(int index) {
                                    // pids 可能比 names 短（服务端漏发或被截断），
                                    // 不校验直接 get 会崩在 UI 线程
                                    if (index < 0 || index >= names.size()) {
                                        return;
                                    }
                                    String pid = index < pids.size()
                                            ? pids.get(index) : "";
                                    confirmKill(names.get(index), pid);
                                }''',
     "pids 越界保护"),
]


def apply(path: Path, old: str, new: str, tag: str) -> bool:
    if not path.is_file():
        print(f"  [跳过] {tag} —— 文件不存在")
        return False
    text = path.read_text(encoding="utf-8")
    if old not in text:
        print(f"  [跳过] {tag} —— 锚点没找到")
        return False
    if text.count(old) > 1:
        print(f"  [注意] {tag} 匹配 {text.count(old)} 处，只改第一处")
    path.write_text(text.replace(old, new, 1), encoding="utf-8")
    print(f"  {tag}: OK")
    return True


def main() -> int:
    print("应用修复：")
    applied = 0
    for path, old, new, tag in FIXES:
        if apply(path, old, new, tag):
            applied += 1

    # ---- RemoteClient：加单行长度常量 ----
    ctext = CLIENT.read_text(encoding="utf-8")
    if "MAX_LINE_CHARS" in ctext and "private static final int MAX_LINE_CHARS" not in ctext:
        ctext = ctext.replace(
            "    private boolean running",
            "    /** 单行上限，和服务端的 8MB 对齐 */\n"
            "    private static final int MAX_LINE_CHARS = 8 * 1024 * 1024;\n\n"
            "    private boolean running", 1)
        CLIENT.write_text(ctext, encoding="utf-8")
        print("  MAX_LINE_CHARS 常量: OK")

    # ---- Json：readArray 的深度计数 + 出口减一 ----
    jtext = JSON.read_text(encoding="utf-8")
    if "depth++;" in jtext and jtext.count("depth++") < 2:
        jtext = jtext.replace(
            "        private List<Object> readArray() {",
            "        private List<Object> readArray() {", 1)
        # 找 readArray 的开头插入 depth++
        marker = "pos++;                       // '['"
        if marker in jtext:
            jtext = jtext.replace(marker, marker + "\n            depth++;", 1)
        JSON.write_text(jtext, encoding="utf-8")
        print("  Json readArray 深度+: OK")

    # ---- 断线清功能缓存 ----
    text = MAIN.read_text(encoding="utf-8")
    hooked = 0
    for anchor in ("    private void doDisconnect() {",
                   "    public void onDisconnected(final String reason) {"):
        if anchor not in text:
            continue
        head, _, tail = text.partition(anchor)
        if "featureCache = null;" in tail[:500]:
            continue
        newline = tail.find("\n")
        insertion = ("\n        // 换连另一台电脑后，「功能」页不能还显示上一台的分组\n"
                     "        featureCache = null;\n"
                     "        featureGroups = null;")
        text = head + anchor + tail[:newline + 1] + insertion + tail[newline + 1:]
        hooked += 1
    if hooked:
        MAIN.write_text(text, encoding="utf-8")
        print(f"  断线清功能缓存: OK（{hooked} 处）")

    print(f"\n  应用了 {applied} 处")
    return 0


if __name__ == "__main__":
    sys.exit(main())