|
|
1
|
+#!/usr/bin/python3
|
|
|
2
|
+# encoding: utf-8
|
|
|
3
|
+# SPDX-FileCopyrightText: 2023 FC Stegerman <flx@obfusk.net>
|
|
|
4
|
+# SPDX-License-Identifier: GPL-3.0-or-later
|
|
|
5
|
+
|
|
|
6
|
+import struct
|
|
|
7
|
+import zipfile
|
|
|
8
|
+import zlib
|
|
|
9
|
+
|
|
|
10
|
+from typing import Any, Dict, Tuple
|
|
|
11
|
+
|
|
|
12
|
+# https://android.googlesource.com/platform/tools/base
|
|
|
13
|
+# profgen/profgen/src/main/kotlin/com/android/tools/profgen/ArtProfileSerializer.kt
|
|
|
14
|
+
|
|
|
15
|
+PROF_MAGIC = b"pro\x00"
|
|
|
16
|
+PROFM_MAGIC = b"prm\x00"
|
|
|
17
|
+
|
|
|
18
|
+PROF_001_N = b"001\x00"
|
|
|
19
|
+PROF_005_O = b"005\x00"
|
|
|
20
|
+PROF_009_O_MR1 = b"009\x00"
|
|
|
21
|
+PROF_010_P = b"010\x00"
|
|
|
22
|
+PROF_015_S = b"015\x00"
|
|
|
23
|
+
|
|
|
24
|
+PROFM_001_N = b"001\x00"
|
|
|
25
|
+PROFM_002 = b"002\x00"
|
|
|
26
|
+
|
|
|
27
|
+ASSET_PROF = "assets/dexopt/baseline.prof"
|
|
|
28
|
+ASSET_PROFM = "assets/dexopt/baseline.profm"
|
|
|
29
|
+
|
|
|
30
|
+ATTRS = ("compress_type", "create_system", "create_version", "date_time",
|
|
|
31
|
+ "external_attr", "extract_version", "flag_bits")
|
|
|
32
|
+LEVELS = (9, 6, 4, 1)
|
|
|
33
|
+
|
|
|
34
|
+
|
|
|
35
|
+class Error(RuntimeError):
|
|
|
36
|
+ pass
|
|
|
37
|
+
|
|
|
38
|
+
|
|
|
39
|
+# FIXME: is there a better alternative?
|
|
|
40
|
+class ReproducibleZipInfo(zipfile.ZipInfo):
|
|
|
41
|
+ """Reproducible ZipInfo hack."""
|
|
|
42
|
+
|
|
|
43
|
+ if "_compresslevel" not in zipfile.ZipInfo.__slots__: # type: ignore[attr-defined]
|
|
|
44
|
+ raise Error("zipfile.ZipInfo has no ._compresslevel")
|
|
|
45
|
+
|
|
|
46
|
+ _compresslevel: int
|
|
|
47
|
+ _override: Dict[str, Any] = {}
|
|
|
48
|
+
|
|
|
49
|
+ def __init__(self, zinfo: zipfile.ZipInfo, **override: Any) -> None:
|
|
|
50
|
+ # pylint: disable=W0231
|
|
|
51
|
+ if override:
|
|
|
52
|
+ self._override = {**self._override, **override}
|
|
|
53
|
+ for k in self.__slots__:
|
|
|
54
|
+ if hasattr(zinfo, k):
|
|
|
55
|
+ setattr(self, k, getattr(zinfo, k))
|
|
|
56
|
+
|
|
|
57
|
+ def __getattribute__(self, name: str) -> Any:
|
|
|
58
|
+ if name != "_override":
|
|
|
59
|
+ try:
|
|
|
60
|
+ return self._override[name]
|
|
|
61
|
+ except KeyError:
|
|
|
62
|
+ pass
|
|
|
63
|
+ return object.__getattribute__(self, name)
|
|
|
64
|
+
|
|
|
65
|
+
|
|
|
66
|
+def sort_baseline(input_file: str, output_file: str) -> None:
|
|
|
67
|
+ with open(input_file, "rb") as fhi:
|
|
|
68
|
+ data = _sort_baseline(fhi.read())
|
|
|
69
|
+ with open(output_file, "wb") as fho:
|
|
|
70
|
+ fho.write(data)
|
|
|
71
|
+
|
|
|
72
|
+
|
|
|
73
|
+def sort_baseline_apk(input_apk: str, output_apk: str) -> None:
|
|
|
74
|
+ with open(input_apk, "rb") as fh_raw:
|
|
|
75
|
+ with zipfile.ZipFile(input_apk) as zf_in:
|
|
|
76
|
+ with zipfile.ZipFile(output_apk, "w") as zf_out:
|
|
|
77
|
+ for info in zf_in.infolist():
|
|
|
78
|
+ attrs = {attr: getattr(info, attr) for attr in ATTRS}
|
|
|
79
|
+ zinfo = ReproducibleZipInfo(info, **attrs)
|
|
|
80
|
+ if info.compress_type == 8:
|
|
|
81
|
+ fh_raw.seek(info.header_offset)
|
|
|
82
|
+ n, m = struct.unpack("<HH", fh_raw.read(30)[26:30])
|
|
|
83
|
+ fh_raw.seek(info.header_offset + 30 + m + n)
|
|
|
84
|
+ ccrc = 0
|
|
|
85
|
+ size = info.compress_size
|
|
|
86
|
+ while size > 0:
|
|
|
87
|
+ ccrc = zlib.crc32(fh_raw.read(min(size, 4096)), ccrc)
|
|
|
88
|
+ size -= 4096
|
|
|
89
|
+ with zf_in.open(info) as fh_in:
|
|
|
90
|
+ comps = {lvl: zlib.compressobj(lvl, 8, -15) for lvl in LEVELS}
|
|
|
91
|
+ ccrcs = {lvl: 0 for lvl in LEVELS}
|
|
|
92
|
+ while True:
|
|
|
93
|
+ data = fh_in.read(4096)
|
|
|
94
|
+ if not data:
|
|
|
95
|
+ break
|
|
|
96
|
+ for lvl in LEVELS:
|
|
|
97
|
+ ccrcs[lvl] = zlib.crc32(comps[lvl].compress(data), ccrcs[lvl])
|
|
|
98
|
+ for lvl in LEVELS:
|
|
|
99
|
+ if ccrc == zlib.crc32(comps[lvl].flush(), ccrcs[lvl]):
|
|
|
100
|
+ zinfo._compresslevel = lvl
|
|
|
101
|
+ break
|
|
|
102
|
+ else:
|
|
|
103
|
+ raise Error(f"Unable to determine compresslevel for {info.filename!r}")
|
|
|
104
|
+ elif info.compress_type != 0:
|
|
|
105
|
+ raise Error(f"Unsupported compress_type {info.compress_type}")
|
|
|
106
|
+ if info.filename == ASSET_PROFM:
|
|
|
107
|
+ print(f"replacing {info.filename!r}...")
|
|
|
108
|
+ zf_out.writestr(zinfo, _sort_baseline(zf_in.read(info)))
|
|
|
109
|
+ else:
|
|
|
110
|
+ with zf_in.open(info) as fh_in:
|
|
|
111
|
+ with zf_out.open(zinfo, "w") as fh_out:
|
|
|
112
|
+ while True:
|
|
|
113
|
+ data = fh_in.read(4096)
|
|
|
114
|
+ if not data:
|
|
|
115
|
+ break
|
|
|
116
|
+ fh_out.write(data)
|
|
|
117
|
+
|
|
|
118
|
+
|
|
|
119
|
+# FIXME
|
|
|
120
|
+# Supported .prof: none
|
|
|
121
|
+# Supported .profm: 002
|
|
|
122
|
+# Unsupported .profm: 001 N
|
|
|
123
|
+def _sort_baseline(data: bytes) -> bytes:
|
|
|
124
|
+ magic, data = _split(data, 4)
|
|
|
125
|
+ version, data = _split(data, 4)
|
|
|
126
|
+ if magic == PROF_MAGIC:
|
|
|
127
|
+ raise Error(f"Unsupported prof version {version!r}")
|
|
|
128
|
+ elif magic == PROFM_MAGIC:
|
|
|
129
|
+ if version == PROFM_002:
|
|
|
130
|
+ return PROFM_MAGIC + PROFM_002 + sort_profm_002(data)
|
|
|
131
|
+ else:
|
|
|
132
|
+ raise Error(f"Unsupported profm version {version!r}")
|
|
|
133
|
+ else:
|
|
|
134
|
+ raise Error(f"Unsupported magic {magic!r}")
|
|
|
135
|
+
|
|
|
136
|
+
|
|
|
137
|
+def sort_profm_002(data: bytes) -> bytes:
|
|
|
138
|
+ num_dex_files, uncompressed_data_size, compressed_data_size, data = _unpack("<HII", data)
|
|
|
139
|
+ profiles = []
|
|
|
140
|
+ if len(data) != compressed_data_size:
|
|
|
141
|
+ raise Error("Compressed data size does not match")
|
|
|
142
|
+ data = zlib.decompress(data)
|
|
|
143
|
+ if len(data) != uncompressed_data_size:
|
|
|
144
|
+ raise Error("Uncompressed data size does not match")
|
|
|
145
|
+ for _ in range(num_dex_files):
|
|
|
146
|
+ profile = data[:4]
|
|
|
147
|
+ profile_idx, profile_key_size, data = _unpack("<HH", data)
|
|
|
148
|
+ profile_key, data = _split(data, profile_key_size)
|
|
|
149
|
+ profile += profile_key + data[:6]
|
|
|
150
|
+ num_type_ids, num_class_ids, data = _unpack("<IH", data)
|
|
|
151
|
+ class_ids, data = _split(data, num_class_ids * 2)
|
|
|
152
|
+ profile += class_ids
|
|
|
153
|
+ profiles.append((profile_key, profile))
|
|
|
154
|
+ if data:
|
|
|
155
|
+ raise Error("Expected end of data")
|
|
|
156
|
+ srtd = b"".join(int.to_bytes(i, 2, "little") + p[1][2:]
|
|
|
157
|
+ for i, p in enumerate(sorted(profiles)))
|
|
|
158
|
+ cdata = zlib.compress(srtd, 1)
|
|
|
159
|
+ hdr = struct.pack("<HII", num_dex_files, uncompressed_data_size, len(cdata))
|
|
|
160
|
+ return hdr + cdata
|
|
|
161
|
+
|
|
|
162
|
+
|
|
|
163
|
+def _unpack(fmt: str, data: bytes) -> Any:
|
|
|
164
|
+ assert all(c in "<BHI" for c in fmt)
|
|
|
165
|
+ size = fmt.count("B") + 2 * fmt.count("H") + 4 * fmt.count("I")
|
|
|
166
|
+ return struct.unpack(fmt, data[:size]) + (data[size:],)
|
|
|
167
|
+
|
|
|
168
|
+
|
|
|
169
|
+def _split(data: bytes, size: int) -> Tuple[bytes, bytes]:
|
|
|
170
|
+ return data[:size], data[size:]
|
|
|
171
|
+
|
|
|
172
|
+
|
|
|
173
|
+if __name__ == "__main__":
|
|
|
174
|
+ import argparse
|
|
|
175
|
+ parser = argparse.ArgumentParser(prog="sort-baseline.py")
|
|
|
176
|
+ parser.add_argument("--apk", action="store_true")
|
|
|
177
|
+ parser.add_argument("input_prof_or_apk", metavar="INPUT_PROF_OR_APK")
|
|
|
178
|
+ parser.add_argument("output_prof_or_apk", metavar="OUTPUT_PROF_OR_APK")
|
|
|
179
|
+ args = parser.parse_args()
|
|
|
180
|
+ if args.apk:
|
|
|
181
|
+ sort_baseline_apk(args.input_prof_or_apk, args.output_prof_or_apk)
|
|
|
182
|
+ else:
|
|
|
183
|
+ sort_baseline(args.input_prof_or_apk, args.output_prof_or_apk)
|
|
|
184
|
+
|
|
|
185
|
+# vim: set tw=80 sw=4 sts=4 et fdm=marker : |