59 lines
2.5 KiB
Python
59 lines
2.5 KiB
Python
import hashlib
|
|
import struct
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
from rewrite_baked_lua_string import stage
|
|
|
|
|
|
def fixture(root, contents):
|
|
digest = hashlib.md5(contents).digest()
|
|
(root / (digest.hex() + ".bytes")).write_bytes(contents)
|
|
# One block and one module, with a 3-byte LuaV block trailer.
|
|
index = (struct.pack(">IIII", 255, 4, 1, 1) + bytes(8) + digest +
|
|
struct.pack(">III", 0, len(contents), 1) + bytes(8) +
|
|
struct.pack(">II", len(contents), 0) + b"\x00\x00\x80")
|
|
index_hash = hashlib.md5(index).digest()
|
|
(root / ("LuaV_" + index_hash.hex() + ".bytes")).write_bytes(index)
|
|
selector = bytearray(66)
|
|
selector[:4] = b"SRMI"
|
|
for i in range(4):
|
|
selector[28 + 4 * i:32 + 4 * i] = index_hash[4 * i:4 * i + 4][::-1]
|
|
struct.pack_into("<I", selector, 44, len(index))
|
|
(root / "M_LuaV.bytes").write_bytes(selector)
|
|
|
|
|
|
class RewriteTest(unittest.TestCase):
|
|
def test_changed_string_updates_digest_and_round_trips(self):
|
|
with tempfile.TemporaryDirectory() as folder:
|
|
root = Path(folder)
|
|
src, out = root / "src", root / "stage"
|
|
src.mkdir()
|
|
old, new = b"scene/a.prefab", b"scene/big.prefab"
|
|
fixture(src, b"prefix" + bytes([len(old)]) + old + b"suffix")
|
|
result = stage(src, out, old.decode(), new.decode())
|
|
changed = (out / (result["new_block_md5"] + ".bytes")).read_bytes()
|
|
self.assertIn(bytes([len(new)]) + new, changed)
|
|
self.assertEqual(result["module_new_size"] - result["module_old_size"], len(new) - len(old))
|
|
reverse = root / "reverse"
|
|
result2 = stage(out, reverse, new.decode(), old.decode())
|
|
self.assertEqual(result2["new_block_md5"], result["old_block_md5"])
|
|
self.assertEqual(result2["new_index_md5"], result["old_index_md5"])
|
|
self.assertEqual((reverse / "M_LuaV.bytes").read_bytes(), (src / "M_LuaV.bytes").read_bytes())
|
|
|
|
def test_duplicate_string_refuses_write(self):
|
|
with tempfile.TemporaryDirectory() as folder:
|
|
root = Path(folder)
|
|
src, out = root / "src", root / "stage"
|
|
src.mkdir()
|
|
old = b"a.prefab"
|
|
fixture(src, (bytes([len(old)]) + old) * 2)
|
|
with self.assertRaisesRegex(ValueError, "ambiguous"):
|
|
stage(src, out, old.decode(), "longer.prefab")
|
|
self.assertFalse(out.exists())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|