#!/usr/bin/env python3
"""把 3.8.99 载荷里所有方法的访问标志加上 public。

为什么敢改访问标志：解密函数算 key 用的是 **(调用方类)** 的常量池大小
    key = (cpSize + 调用类名 + 调用方法名 + cpSize).hashCode()
其中 cpSize 每次都取 st[2]（调用解密函数的那个类），跟解密函数所在类无关。
而 access_flags 是方法头里的两个字节，改它不会动常量池，密文一个字节都不受影响。

改 public 的目的：我们可以造一个合成类，用**原始调用方的类名和方法名**去调解密函数，
从而让 key 自然对上，不需要自己实现 AES —— 载荷里有好几个内联副本，全开 public 最省事。
"""

import sys
import zipfile

ACC_PUBLIC = 0x0001
ACC_PRIVATE = 0x0002
ACC_PROTECTED = 0x0004


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

    def u1(self):
        v = self.d[self.p]
        self.p += 1
        return v

    def u2(self):
        v = (self.d[self.p] << 8) | self.d[self.p + 1]
        self.p += 2
        return v

    def u4(self):
        v = int.from_bytes(self.d[self.p:self.p + 4], "big")
        self.p += 4
        return v

    def skip(self, n):
        self.p += n


def parse_constant_pool(r):
    """返回 (utf8 列表, 索引->绝对偏移) 以及常量池结束位置。"""
    count = r.u2()
    utf8 = {}
    offsets = {}
    i = 1
    while i < count:
        offsets[i] = r.p
        tag = r.u1()
        if tag == 1:
            length = r.u2()
            raw = r.d[r.p:r.p + length]
            r.skip(length)
            utf8[i] = raw
        elif tag in (3, 4):
            r.skip(4)
        elif tag in (5, 6):
            r.skip(8)
            i += 1
        elif tag in (7, 8, 16, 19, 20):
            r.skip(2)
        elif tag in (9, 10, 11, 12, 17, 18):
            r.skip(4)
        elif tag == 15:
            r.skip(3)
        else:
            raise ValueError("未知常量池 tag %d @%d" % (tag, r.p - 1))
        i += 1
    return utf8, offsets, count


def skip_attributes(r):
    n = r.u2()
    for _ in range(n):
        r.skip(2)
        r.skip(r.u4())


def patch(data):
    """把类里所有方法的访问标志加上 ACC_PUBLIC，就地修改 data（bytearray）。"""
    r = Reader(data)
    if r.u4() != 0xCAFEBABE:
        raise ValueError("不是 class 文件")
    r.u2()
    r.u2()
    utf8, offsets, cp_count = parse_constant_pool(r)
    r.skip(2)  # access_flags（类的）
    r.skip(2)  # this_class
    r.skip(2)  # super_class
    r.skip(2 * r.u2())  # interfaces

    for _ in range(r.u2()):  # fields
        r.skip(6)
        skip_attributes(r)

    changed = 0
    for _ in range(r.u2()):  # methods
        access_pos = r.p
        access = r.u2()
        r.skip(4)  # name_index / descriptor_index
        skip_attributes(r)
        # 不能直接 OR：private static final 是 0x1A，OR 上 0x01 会变成 0x1B（public+private，非法）
        new_access = (access & ~(ACC_PRIVATE | ACC_PROTECTED)) | ACC_PUBLIC
        if new_access != access:
            data[access_pos] = (new_access >> 8) & 0xFF
            data[access_pos + 1] = new_access & 0xFF
            changed += 1
    return changed, cp_count


def main():
    source, target = sys.argv[1], sys.argv[2]
    patched_classes = 0
    patched_methods = 0
    with zipfile.ZipFile(source) as zin:
        with zipfile.ZipFile(target, "w", zipfile.ZIP_DEFLATED) as zout:
            for item in zin.infolist():
                raw = zin.read(item.filename)
                if item.filename.endswith(".class"):
                    raw = bytearray(raw)
                    try:
                        changed, cp_count = patch(raw)
                        patched_classes += 1
                        patched_methods += changed
                    except Exception as error:
                        # 载荷里有几个不以 CAFEBABE 开头的 .class（比如 wvv.class 是签名数据），跳过
                        print("跳过 %s : %s" % (item.filename, error))
                zout.writestr(item, raw)
    print("已把 %d 个类里的 %d 个方法改成 public" % (patched_classes, patched_methods))
    print("已写出 " + target)


if __name__ == "__main__":
    main()