Validate firmware image before patching

combine.py

@@ -3,57 +3,103 @@
 
 from firmware_addresses import DEFAULT_VERSION, FirmwareAddressError, get_firmware_addresses
 
-print("Enter target system: A for Rev. A, B for Rev. B, or AB for both")
-target = input("|)> ")
 
+FIRMWARE_HEADER_SIZE = 32
+LENGTH_OFFSET = 8
+HASH_OFFSET = 24
+HASH_SIZE = 8
 
-#TODO: add ptr patching to REV. A once i figure out how to develop cfw for it
 
+def read_firmware_image(path):
+	with open(path, "rb") as f:
+		fw_header = bytearray(f.read(FIRMWARE_HEADER_SIZE))
+		fw = f.read()
 
+	if len(fw_header) != FIRMWARE_HEADER_SIZE:
+		raise ValueError(f"{path}: expected 32-byte header, got {len(fw_header)} bytes")
 
-if "A" in target.upper():
-	with open("pdfw-a", "rb") as f:
-		fw_header = bytearray(f.read(32))
-		fw = f.read()
-
+	expected_length = int.from_bytes(fw_header[LENGTH_OFFSET:LENGTH_OFFSET + 4], byteorder="little")
+	if expected_length != len(fw):
+		raise ValueError(f"{path}: header body length {expected_length} does not match actual body length {len(fw)}")
+
+	expected_hash = bytes(fw_header[HASH_OFFSET:HASH_OFFSET + HASH_SIZE])
+	actual_hash = md5(fw).digest()[:HASH_SIZE]
+	if expected_hash != actual_hash:
+		raise ValueError(f"{path}: header body hash does not match actual body hash")
+
+	return fw_header, fw
+
+
+def read_loader_binary(path):
+	try:
+		with open(path, "rb") as f:
+			return f.read()
+	except FileNotFoundError:
+		raise FileNotFoundError(f"required loader input not found: {path}") from None
+	except OSError as exc:
+		raise OSError(f"required loader input is not readable: {path}") from exc
+
+
+def patch_rev_b_pointer(fw, ptr_addr, update_func_addr):
+	if ptr_addr < 0 or ptr_addr + 4 > len(fw):
+		raise ValueError(f"Rev. B patch address 0x{ptr_addr:x} is outside firmware body length {len(fw)}")
+	struct.pack_into("<I", fw, ptr_addr, update_func_addr)
+
+
+def finalize_firmware_header(fw_header, fw):
+	fw_header[LENGTH_OFFSET:LENGTH_OFFSET + 4] = len(fw).to_bytes(4, byteorder="little")
+	fw_header[HASH_OFFSET:HASH_OFFSET + HASH_SIZE] = md5(fw).digest()[:HASH_SIZE]
+
+
+def pad_firmware_body(fw):
 	while bool(len(fw) & 3):
 		fw += b"\x00"
-
-	with open("build/loader-a.bin", "rb") as f:
-		fw += f.read()
-
-	fw_header[8:12] = len(fw).to_bytes(4, byteorder="little")
-	fw_header[24:] = md5(fw).digest()[:8]
-
+	return fw
+
+
+def patch_rev_a():
+	fw_header, fw = read_firmware_image("pdfw-a")
+	fw = pad_firmware_body(fw)
+	fw += read_loader_binary("build/loader-a.bin")
+	finalize_firmware_header(fw_header, fw)
+
 	with open("build/pdfw-a-patched", "wb") as f:
 		f.write(fw_header)
 		f.write(fw)
 
-if "B" in target.upper():
+
+def patch_rev_b():
 	addresses = get_firmware_addresses(version=DEFAULT_VERSION, revision="B")
 	try:
 		ptr_addr = addresses.require_patch_pointer_offset()
 		loader_entry_address = addresses.require_loader_entry_address()
 	except FirmwareAddressError as exc:
 		raise SystemExit(f"error: {exc}") from None
 
-	with open("pdfw-b", "rb") as f:
-		fw_header = bytearray(f.read(32))
-		fw = f.read()
+	fw_header, fw = read_firmware_image("pdfw-b")
+	fw = pad_firmware_body(fw)
+	fw += read_loader_binary("build/loader-b.bin")
 
-
-	while bool(len(fw) & 3):
-		fw += b"\x00"
-
-	with open("build/loader-b.bin", "rb") as f:
-		fw += f.read()
-
 	fw = bytearray(fw)
-	struct.pack_into("<I", fw, ptr_addr, loader_entry_address)
+	patch_rev_b_pointer(fw, ptr_addr, loader_entry_address)
+	finalize_firmware_header(fw_header, fw)
 
-	fw_header[8:12] = len(fw).to_bytes(4, byteorder="little")
-	fw_header[24:] = md5(fw).digest()[:8]
-
 	with open("build/pdfw-b-patched", "wb") as f:
 		f.write(fw_header)
 		f.write(fw)
+
+
+def main():
+	print("Enter target system: A for Rev. A, B for Rev. B, or AB for both")
+	target = input("|)> ")
+
+	#TODO: add ptr patching to REV. A once i figure out how to develop cfw for it
+	if "A" in target.upper():
+		patch_rev_a()
+
+	if "B" in target.upper():
+		patch_rev_b()
+
+
+if __name__ == "__main__":
+	main()

tests/test_combine.py

@@ -0,0 +1,93 @@
+import hashlib
+import importlib
+import os
+import tempfile
+import unittest
+from pathlib import Path
+
+
+combine = importlib.import_module("combine")
+
+
+def make_firmware(body):
+	header = bytearray(32)
+	header[8:12] = len(body).to_bytes(4, byteorder="little")
+	header[24:32] = hashlib.md5(body).digest()[:8]
+	return bytes(header) + body
+
+
+class FirmwareValidationTests(unittest.TestCase):
+	def test_read_firmware_image_returns_valid_header_and_body(self):
+		body = b"firmware-body"
+		with tempfile.NamedTemporaryFile(delete=False) as tmp:
+			tmp.write(make_firmware(body))
+			path = tmp.name
+		try:
+			header, read_body = combine.read_firmware_image(path)
+		finally:
+			os.unlink(path)
+
+		self.assertEqual(len(header), 32)
+		self.assertEqual(read_body, body)
+
+	def test_read_firmware_image_rejects_bad_body_length(self):
+		body = b"firmware-body"
+		image = bytearray(make_firmware(body))
+		image[8:12] = (len(body) + 1).to_bytes(4, byteorder="little")
+		with tempfile.NamedTemporaryFile(delete=False) as tmp:
+			tmp.write(image)
+			path = tmp.name
+		try:
+			with self.assertRaisesRegex(ValueError, "length"):
+				combine.read_firmware_image(path)
+		finally:
+			os.unlink(path)
+
+	def test_read_firmware_image_rejects_bad_body_hash(self):
+		body = b"firmware-body"
+		image = bytearray(make_firmware(body))
+		image[24:32] = b"\x00" * 8
+		with tempfile.NamedTemporaryFile(delete=False) as tmp:
+			tmp.write(image)
+			path = tmp.name
+		try:
+			with self.assertRaisesRegex(ValueError, "hash"):
+				combine.read_firmware_image(path)
+		finally:
+			os.unlink(path)
+
+	def test_read_firmware_image_rejects_short_header(self):
+		with tempfile.NamedTemporaryFile(delete=False) as tmp:
+			tmp.write(b"short")
+			path = tmp.name
+		try:
+			with self.assertRaisesRegex(ValueError, "32-byte header"):
+				combine.read_firmware_image(path)
+		finally:
+			os.unlink(path)
+
+
+class LoaderInputTests(unittest.TestCase):
+	def test_read_loader_binary_returns_existing_readable_file(self):
+		with tempfile.NamedTemporaryFile(delete=False) as tmp:
+			tmp.write(b"loader")
+			path = tmp.name
+		try:
+			self.assertEqual(combine.read_loader_binary(path), b"loader")
+		finally:
+			os.unlink(path)
+
+	def test_read_loader_binary_rejects_missing_file(self):
+		path = Path(tempfile.gettempdir()) / "missing-playbrew-loader.bin"
+		with self.assertRaises(FileNotFoundError):
+			combine.read_loader_binary(path)
+
+
+class PatchBoundsTests(unittest.TestCase):
+	def test_patch_rev_b_pointer_rejects_out_of_bounds_offset(self):
+		with self.assertRaisesRegex(ValueError, "outside firmware body"):
+			combine.patch_rev_b_pointer(bytearray(3), ptr_addr=0, update_func_addr=1)
+
+
+if __name__ == "__main__":
+	unittest.main()