#!/usr/bin/env python3
"""Minimal Java .class parser.

Reads every .class under a directory and prints a JSON model with the pieces we
need to identify obfuscated classes by *structure* rather than by name:
this class, super class, declared fields (name + type) and declared methods
(name + parameter/return types).

Usage:
    python3 tools/classdump.py <classes-dir> [--out model.json] [--only TextButton]
"""

import argparse
import json
import os
import struct
import sys

TYPE_CHARS = {
    "B": "byte", "C": "char", "D": "double", "F": "float", "I": "int",
    "J": "long", "S": "short", "Z": "boolean", "V": "void",
}


def parse_descriptor(desc, at=0):
    """Turn one JVM type descriptor into (readable name, next index)."""
    arrays = 0
    while at < len(desc) and desc[at] == "[":
        arrays += 1
        at += 1
    if at >= len(desc):
        return "?", at
    ch = desc[at]
    if ch == "L":
        end = desc.find(";", at)
        if end < 0:
            return "?", len(desc)
        name = desc[at + 1:end].replace("/", ".")
        at = end + 1
    else:
        name = TYPE_CHARS.get(ch, ch)
        at += 1
    return name + "[]" * arrays, at


def parse_method_descriptor(desc):
    params = []
    at = 0
    assert desc[at] == "("
    at += 1
    while desc[at] != ")":
        name, at = parse_descriptor(desc, at)
        params.append(name)
    at += 1
    result, _ = parse_descriptor(desc, at)
    return params, result


class Reader:
    def __init__(self, data):
        self.data = data
        self.pos = 0

    def u1(self):
        value = self.data[self.pos]
        self.pos += 1
        return value

    def u2(self):
        value = struct.unpack_from(">H", self.data, self.pos)[0]
        self.pos += 2
        return value

    def u4(self):
        value = struct.unpack_from(">I", self.data, self.pos)[0]
        self.pos += 4
        return value

    def skip(self, count):
        self.pos += count


def read_constant_pool(reader):
    count = reader.u2()
    pool = [None] * count
    index = 1
    while index < count:
        tag = reader.u1()
        if tag == 1:  # Utf8
            length = reader.u2()
            raw = reader.data[reader.pos:reader.pos + length]
            reader.skip(length)
            pool[index] = raw.decode("utf-8", "replace")
        elif tag in (3, 4):  # Integer, Float
            reader.skip(4)
        elif tag in (5, 6):  # Long, Double
            reader.skip(8)
            index += 1
        elif tag == 7:  # Class -> index of the Utf8 holding the internal name
            pool[index] = ("class", reader.u2())
        elif tag in (8, 16, 19, 20):  # String, MethodType, Module, Package
            reader.skip(2)
        elif tag in (9, 10, 11, 12, 17, 18):  # Field/Method/Interface refs, NameAndType, Dynamic
            reader.skip(4)
        elif tag == 15:  # MethodHandle
            reader.skip(3)
        else:
            raise ValueError("unknown constant pool tag %d" % tag)
        index += 1
    return pool


def parse_class(path):
    with open(path, "rb") as handle:
        data = handle.read()
    if len(data) < 10 or data[:4] != b"\xca\xfe\xba\xbe":
        return None
    reader = Reader(data)
    reader.skip(4)
    reader.u2()  # minor
    reader.u2()  # major
    pool = read_constant_pool(reader)

    def utf8(index):
        value = pool[index] if 0 < index < len(pool) else None
        return value if isinstance(value, str) else None

    def class_name(index):
        entry = pool[index] if 0 < index < len(pool) else None
        if isinstance(entry, tuple) and entry[0] == "class":
            entry = utf8(entry[1])
        if isinstance(entry, str):
            return entry.replace("/", ".")
        return None

    access_flags = reader.u2()
    this_class = class_name(reader.u2())
    super_class = class_name(reader.u2())
    interfaces = [class_name(reader.u2()) for _ in range(reader.u2())]

    fields = []
    for _ in range(reader.u2()):
        flags = reader.u2()
        name = utf8(reader.u2())
        desc = utf8(reader.u2())
        for _ in range(reader.u2()):
            reader.u2()
            reader.skip(reader.u4())
        fields.append({
            "name": name,
            "type": parse_descriptor(desc)[0] if desc else "?",
            "flags": flags,
        })

    methods = []
    for _ in range(reader.u2()):
        flags = reader.u2()
        name = utf8(reader.u2())
        desc = utf8(reader.u2())
        for _ in range(reader.u2()):
            reader.u2()
            reader.skip(reader.u4())
        if desc and desc.startswith("("):
            params, result = parse_method_descriptor(desc)
        else:
            params, result = [], "?"
        methods.append({"name": name, "params": params, "returns": result, "flags": flags})

    return {
        "name": this_class,
        "super": super_class,
        "interfaces": interfaces,
        "flags": access_flags,
        "fields": fields,
        "methods": methods,
    }


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("directory")
    parser.add_argument("--out")
    parser.add_argument("--only", help="only keep classes whose name contains this")
    args = parser.parse_args()

    model = {}
    failed = 0
    for dirpath, _, filenames in os.walk(args.directory):
        for filename in filenames:
            if not filename.endswith(".class"):
                continue
            path = os.path.join(dirpath, filename)
            try:
                entry = parse_class(path)
            except Exception:
                failed += 1
                continue
            if not entry or not entry["name"]:
                failed += 1
                continue
            if args.only and args.only not in entry["name"]:
                continue
            entry["path"] = path
            model[entry["name"]] = entry

    print("解析成功 %d 个类，失败 %d" % (len(model), failed), file=sys.stderr)
    if args.out:
        with open(args.out, "w") as handle:
            json.dump(model, handle)
        print("已写入 %s" % args.out, file=sys.stderr)


if __name__ == "__main__":
    main()