275 lines
9.4 KiB
Python
275 lines
9.4 KiB
Python
#!/usr/bin/env python3
|
|
# SPDX-License-Identifier: AGPL-3.0-or-later
|
|
"""Strict structural mutations used by archive-authentication tests."""
|
|
|
|
import argparse
|
|
import pathlib
|
|
import sys
|
|
|
|
|
|
ARCHIVE_HEADER_SIZE = 64
|
|
FOOTER_SIZE = 32
|
|
AIT_SIZE = 32
|
|
BLOCK_DATA = 0x00
|
|
BLOCK_INDEX = 0x02
|
|
BLOCK_ENC_HEADER = 0x03
|
|
BLOCK_DEDUP_REF = 0x04
|
|
BLOCK_COMMENT = 0x05
|
|
BLOCK_FLAG_ENCRYPTED = 0x01
|
|
|
|
|
|
class ArchiveError(ValueError):
|
|
pass
|
|
|
|
|
|
def _u16le(data, offset):
|
|
return int.from_bytes(data[offset:offset + 2], "little")
|
|
|
|
|
|
def _u64le(data, offset):
|
|
return int.from_bytes(data[offset:offset + 8], "little")
|
|
|
|
|
|
def _read_varint(data, offset, limit):
|
|
value = 0
|
|
for byte_number in range(10):
|
|
if offset >= limit:
|
|
raise ArchiveError("truncated varint")
|
|
byte = data[offset]
|
|
offset += 1
|
|
if byte_number == 9 and byte > 1:
|
|
raise ArchiveError("varint exceeds uint64")
|
|
value |= (byte & 0x7f) << (7 * byte_number)
|
|
if (byte & 0x80) == 0:
|
|
if byte_number and value < (1 << (7 * byte_number)):
|
|
raise ArchiveError("non-canonical varint")
|
|
return value, offset
|
|
raise ArchiveError("unterminated varint")
|
|
|
|
|
|
def _parse_frame(data, offset, limit):
|
|
start = offset
|
|
if limit - offset < 17:
|
|
raise ArchiveError("truncated block header")
|
|
if data[offset:offset + 2] != b"\xbb\x01":
|
|
raise ArchiveError("invalid block magic")
|
|
block_type = data[offset + 2]
|
|
codec = _u16le(data, offset + 3)
|
|
flags = _u16le(data, offset + 5)
|
|
offset += 7
|
|
uncompressed_size, offset = _read_varint(data, offset, limit)
|
|
compressed_size, offset = _read_varint(data, offset, limit)
|
|
if limit - offset < 8:
|
|
raise ArchiveError("truncated block checksum")
|
|
checksum = _u64le(data, offset)
|
|
payload_start = offset + 8
|
|
if compressed_size > limit - payload_start:
|
|
raise ArchiveError("block payload exceeds structural boundary")
|
|
end = payload_start + compressed_size
|
|
return {
|
|
"type": block_type,
|
|
"codec": codec,
|
|
"flags": flags,
|
|
"uncompressed_size": uncompressed_size,
|
|
"compressed_size": compressed_size,
|
|
"checksum": checksum,
|
|
"start": start,
|
|
"payload_start": payload_start,
|
|
"end": end,
|
|
}
|
|
|
|
|
|
def _parse_current_archive(data):
|
|
minimum = ARCHIVE_HEADER_SIZE + FOOTER_SIZE + AIT_SIZE
|
|
if len(data) < minimum:
|
|
raise ArchiveError("archive is too short")
|
|
if data[:6] != b"ZUPT\x1a\x00":
|
|
raise ArchiveError("invalid archive magic")
|
|
|
|
footer_start = len(data) - FOOTER_SIZE - AIT_SIZE
|
|
if data[footer_start + 24:footer_start + 28] != b"ZEND":
|
|
raise ArchiveError("current footer before AIT not found")
|
|
if int.from_bytes(data[footer_start + 28:footer_start + 32],
|
|
"little") != 1:
|
|
raise ArchiveError("unsupported footer version")
|
|
|
|
index_offset = _u64le(data, footer_start)
|
|
if index_offset < ARCHIVE_HEADER_SIZE or index_offset >= footer_start:
|
|
raise ArchiveError("index offset is outside the archive body")
|
|
|
|
frames = []
|
|
offset = ARCHIVE_HEADER_SIZE
|
|
while offset < index_offset:
|
|
frame = _parse_frame(data, offset, index_offset)
|
|
if frame["type"] == BLOCK_INDEX:
|
|
raise ArchiveError("index frame occurs before footer index offset")
|
|
frames.append(frame)
|
|
offset = frame["end"]
|
|
if offset != index_offset:
|
|
raise ArchiveError("archive body does not end at index offset")
|
|
|
|
index = _parse_frame(data, index_offset, footer_start)
|
|
if index["type"] != BLOCK_INDEX:
|
|
raise ArchiveError("footer does not point to an index frame")
|
|
if index["end"] != footer_start:
|
|
raise ArchiveError("bytes remain between index and footer")
|
|
|
|
return {
|
|
"frames": frames,
|
|
"index": index,
|
|
"footer_start": footer_start,
|
|
}
|
|
|
|
|
|
def _kind_value(name):
|
|
return {"data": BLOCK_DATA, "enc": BLOCK_ENC_HEADER,
|
|
"ref": BLOCK_DEDUP_REF}[name]
|
|
|
|
|
|
def _matching_frames(layout, kind, require_encrypted):
|
|
matches = [frame for frame in layout["frames"]
|
|
if frame["type"] == _kind_value(kind)]
|
|
if require_encrypted:
|
|
matches = [frame for frame in matches
|
|
if frame["flags"] & BLOCK_FLAG_ENCRYPTED]
|
|
return matches
|
|
|
|
|
|
def _same_metadata(left, right):
|
|
fields = ("type", "codec", "flags", "uncompressed_size",
|
|
"compressed_size", "checksum")
|
|
return all(left[field] == right[field] for field in fields)
|
|
|
|
|
|
def _select_equal_length_pair(frames, same_metadata):
|
|
for index, left in enumerate(frames):
|
|
for right in frames[index + 1:]:
|
|
if left["end"] - left["start"] != right["end"] - right["start"]:
|
|
continue
|
|
if same_metadata and not _same_metadata(left, right):
|
|
continue
|
|
return left, right
|
|
qualifier = " with identical metadata" if same_metadata else ""
|
|
raise ArchiveError("no two equal-length frames" + qualifier)
|
|
|
|
|
|
def _write(destination, data):
|
|
pathlib.Path(destination).write_bytes(data)
|
|
|
|
|
|
def command_strip_ait(args):
|
|
data = pathlib.Path(args.source).read_bytes()
|
|
layout = _parse_current_archive(data)
|
|
_write(args.destination, data[:layout["footer_start"] + FOOTER_SIZE])
|
|
|
|
|
|
def command_flip_payload(args):
|
|
data = bytearray(pathlib.Path(args.source).read_bytes())
|
|
layout = _parse_current_archive(data)
|
|
frames = _matching_frames(layout, args.kind, args.require_encrypted)
|
|
if not frames:
|
|
raise ArchiveError("requested frame was not found")
|
|
frame = frames[0]
|
|
if frame["compressed_size"] == 0:
|
|
raise ArchiveError("requested frame has no payload")
|
|
position = frame["payload_start"] + frame["compressed_size"] // 2
|
|
data[position] ^= 0x01
|
|
_write(args.destination, data)
|
|
|
|
|
|
def command_swap_frames(args):
|
|
data = bytearray(pathlib.Path(args.source).read_bytes())
|
|
layout = _parse_current_archive(data)
|
|
frames = _matching_frames(layout, args.kind, args.require_encrypted)
|
|
left, right = _select_equal_length_pair(frames, args.same_metadata)
|
|
left_bytes = bytes(data[left["start"]:left["end"]])
|
|
right_bytes = bytes(data[right["start"]:right["end"]])
|
|
if left_bytes == right_bytes:
|
|
raise ArchiveError("selected frames are byte-identical; swap is a no-op")
|
|
data[left["start"]:left["end"]] = right_bytes
|
|
data[right["start"]:right["end"]] = left_bytes
|
|
_write(args.destination, data)
|
|
|
|
|
|
def command_replay_frame(args):
|
|
data = bytearray(pathlib.Path(args.source).read_bytes())
|
|
layout = _parse_current_archive(data)
|
|
frames = _matching_frames(layout, args.kind, args.require_encrypted)
|
|
source, destination = _select_equal_length_pair(frames,
|
|
args.same_metadata)
|
|
replay = bytes(data[source["start"]:source["end"]])
|
|
if replay == bytes(data[destination["start"]:destination["end"]]):
|
|
raise ArchiveError("selected frames are already byte-identical")
|
|
data[destination["start"]:destination["end"]] = replay
|
|
_write(args.destination, data)
|
|
|
|
|
|
def command_preface_positions(args):
|
|
data = pathlib.Path(args.source).read_bytes()
|
|
layout = _parse_current_archive(data)
|
|
for frame in layout["frames"] + [layout["index"]]:
|
|
for position in range(frame["start"], frame["payload_start"]):
|
|
print(position)
|
|
|
|
|
|
def command_set_frame_type(args):
|
|
data = bytearray(pathlib.Path(args.source).read_bytes())
|
|
layout = _parse_current_archive(data)
|
|
frames = _matching_frames(layout, args.kind, args.require_encrypted)
|
|
if not frames:
|
|
raise ArchiveError("requested frame was not found")
|
|
replacement = {"data": BLOCK_DATA, "comment": BLOCK_COMMENT}[args.type]
|
|
data[frames[0]["start"] + 2] = replacement
|
|
_write(args.destination, data)
|
|
|
|
|
|
def _add_frame_options(parser):
|
|
parser.add_argument("source")
|
|
parser.add_argument("destination")
|
|
parser.add_argument("--kind", choices=("data", "enc", "ref"), required=True)
|
|
parser.add_argument("--require-encrypted", action="store_true")
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
commands = parser.add_subparsers(dest="command")
|
|
|
|
strip_ait = commands.add_parser("strip-ait")
|
|
strip_ait.add_argument("source")
|
|
strip_ait.add_argument("destination")
|
|
strip_ait.set_defaults(function=command_strip_ait)
|
|
|
|
flip = commands.add_parser("flip-payload")
|
|
_add_frame_options(flip)
|
|
flip.set_defaults(function=command_flip_payload)
|
|
|
|
swap = commands.add_parser("swap-frames")
|
|
_add_frame_options(swap)
|
|
swap.add_argument("--same-metadata", action="store_true")
|
|
swap.set_defaults(function=command_swap_frames)
|
|
|
|
replay = commands.add_parser("replay-frame")
|
|
_add_frame_options(replay)
|
|
replay.add_argument("--same-metadata", action="store_true")
|
|
replay.set_defaults(function=command_replay_frame)
|
|
|
|
prefaces = commands.add_parser("preface-positions")
|
|
prefaces.add_argument("source")
|
|
prefaces.set_defaults(function=command_preface_positions)
|
|
|
|
set_type = commands.add_parser("set-frame-type")
|
|
_add_frame_options(set_type)
|
|
set_type.add_argument("--type", choices=("data", "comment"), required=True)
|
|
set_type.set_defaults(function=command_set_frame_type)
|
|
|
|
args = parser.parse_args()
|
|
if not hasattr(args, "function"):
|
|
parser.error("a mutation command is required")
|
|
try:
|
|
args.function(args)
|
|
except (ArchiveError, OSError) as error:
|
|
parser.error(str(error))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|