src/snagrecover/firmware/rk_fw.py: Add support for sending rockchip bin files

Add support for reading and sending rockchip binary files, made with
boot_merger.

Signed-off-by: Arnaud Patard <[email protected]>

Author: apatard

src/snagrecover/firmware/rk_fw.py

@@ -0,0 +1,223 @@
+#!/usr/bin/python3
+# Copyright 2025 Collabora Ltd.
+#
+# SPDX-License-Identifier: GPL-2.0+
+#
+# Author: Arnaud Patard <[email protected]>
+#
+# Notes:
+# to unpack / descramble encrypted parts, the rc4 key is inside u-boot's code.
+# Some information used here are coming from rkdeveloptool, which is GPL-2.0+
+
+import logging
+import struct
+import time
+from cryptography.hazmat.primitives.ciphers import Cipher
+from cryptography.hazmat.decrepit.ciphers.algorithms import ARC4
+from crccheck import crc
+
+logger = logging.getLogger("snagrecover")
+from snagrecover.protocols import rockchip
+
+BOOTTAG = b"BOOT"
+LDRTAG = b"LDR "
+TAG_LIST = [ BOOTTAG, LDRTAG ]
+BOOTENTRYSIZE = 57
+BOOTHEADERENTRYSIZE = 6
+BOOTHEADERSIZE = 102
+BOOTHEADERTIMESIZE = 7
+RC4_KEY = bytearray([124, 78, 3, 4, 85, 5, 9, 7, 45, 44, 123, 56, 23, 13, 23, 17])
+
+class BootEntryError(Exception):
+	pass
+
+class BootEntry():
+	def __init__(self, entry):
+		if len(entry) != BOOTENTRYSIZE:
+			raise BootEntryError(f"Invalid boot entry size {len(entry)}")
+		self.entry = entry
+
+	def unpack(self):
+		(self.size, self.type) = struct.unpack("<BI", self.entry[:5])
+		self.name = self.entry[5:45].decode('utf-16le')
+		(self.data_offset, self.data_size, self.data_delay) = struct.unpack("<III", self.entry[45:57])
+
+	def __str__(self):
+		return f"Entry {self.name} (type: {self.type}, size: {self.size}, data offset: {self.data_offset}, data size: {self.data_size}, delay: {self.data_delay})"
+
+
+class BootHeaderEntryError(Exception):
+	pass
+
+class BootHeaderEntry():
+	def __init__(self, name, entry):
+		self.name = name
+		if len(entry) != BOOTHEADERENTRYSIZE:
+			raise BootHeaderEntryError(f"Invalid boot header entry size {len(entry)}")
+		self.entry = entry
+
+	def unpack(self):
+		(self.count, self.offset, self.size) = struct.unpack("<BIB", self.entry)
+		if self.size != BOOTENTRYSIZE:
+			raise BootHeaderError("Invalid boot entry size in the header definition")
+
+	def __str__(self):
+		return f"Header entry {self.name}: ({self.count}, {self.offset})"
+
+class BootHeaderError(Exception):
+	pass
+
+class BootHeader():
+	def __init__(self, header):
+		if len(header) != BOOTHEADERSIZE:
+			raise BootHeaderError(f"Invalid boot header size {len(header)}")
+		self.header = header
+
+	def unpack(self):
+		offset = 0
+		self.tag = self.header[offset:offset+4]
+		offset += 4
+		(self.size, version, self.merge_version) = struct.unpack("<HII", self.header[offset:offset+10])
+		if self.size != BOOTHEADERSIZE:
+			raise BootHeaderError("Invalid header size in the header definition")
+		# not sure how to exactly parse version/merge_version
+		self.maj_ver = version >> 8
+		self.min_ver = version & 0xff
+		offset += 10
+		self.releasetime = self.header[offset:offset+BOOTHEADERTIMESIZE]
+		offset += BOOTHEADERTIMESIZE
+		self.chip = self.header[offset:offset+4][::-1]
+		offset += 4
+		self.entry471 = BootHeaderEntry("471", self.header[offset:offset+BOOTHEADERENTRYSIZE])
+		self.entry471.unpack()
+		offset += BOOTHEADERENTRYSIZE
+		self.entry472 = BootHeaderEntry("472", self.header[offset:offset+BOOTHEADERENTRYSIZE])
+		self.entry472.unpack()
+		offset += BOOTHEADERENTRYSIZE
+		self.loader = BootHeaderEntry("loader", self.header[offset:offset+BOOTHEADERENTRYSIZE])
+		self.loader.unpack()
+		offset += BOOTHEADERENTRYSIZE
+		# rc4 = 1: disable rc4 on 471/472. Not sure about rc4 on loader.
+		(self.sign, self.rc4) = struct.unpack("<BB", self.header[offset:offset+2])
+		if self.rc4:
+			self.rc4 = False
+		else:
+			self.rc4 = True
+		if self.sign == 'S':
+			self.sign = True
+		else:
+			self.sign = False
+
+	def __str__(self):
+		return f"{self.tag},{self.maj_ver}.{self.min_ver},{self.merge_version},{self.chip},{self.entry471},{self.entry472},{self.loader},sign:{self.sign},enc:{self.rc4}"
+
+class LoaderFileError(Exception):
+	pass
+
+class RkCrc32(crc.Crc32Base):
+	"""CRC-32/ROCKCHIP
+	"""
+	_names = ('CRC-32/ROCKCHIP')
+	_width = 32
+	_poly = 0x04c10db7
+	_initvalue = 0x00000000
+	_reflect_input = False
+	_reflect_output = False
+	_xor_output = 0
+
+class LoaderFile():
+	def __init__(self, blob):
+		self.blob = blob
+		buf = self.blob[0:BOOTHEADERSIZE]
+		offset = BOOTHEADERSIZE
+		self.header = BootHeader(buf)
+		self.header.unpack()
+		if self.header.tag not in TAG_LIST:
+			raise LoaderFileError(f"Invalid tag {self.header.tag}")
+
+		offset = self.header.entry471.offset
+		self.entry471 = []
+		for _i in range(self.header.entry471.count):
+			self.entry471.append(BootEntry(self.blob[offset:offset+BOOTENTRYSIZE]))
+			offset += BOOTENTRYSIZE
+
+		offset = self.header.entry472.offset
+		self.entry472 = []
+		for _i in range(self.header.entry472.count):
+			self.entry472.append(BootEntry(self.blob[offset:offset+BOOTENTRYSIZE]))
+			offset += BOOTENTRYSIZE
+
+		offset = self.header.loader.offset
+		self.loader = []
+		for _i in range(self.header.loader.count):
+			self.loader.append(BootEntry(self.blob[offset:offset+BOOTENTRYSIZE]))
+			offset += BOOTENTRYSIZE
+		crc32 = self.blob[-4:]
+		calc_crc32 = RkCrc32.calc(self.blob[:-4])
+		(self.crc32,) = struct.unpack("<I", crc32)
+		assert self.crc32 == calc_crc32
+
+	def entry_data(self, name, idx = 0):
+		entry = None
+		if name == "471":
+			entry = self.entry471
+		elif name == "472":
+			entry = self.entry472
+		elif name == "loader":
+			entry = self.loader
+		else:
+			raise LoaderFileError(f"Invalid name {name}")
+
+		if idx > len(entry):
+			raise LoaderFileError(f"Invalid index {idx}. Only has {len(entry)} entries.")
+		e = entry[idx]
+		e.unpack()
+		logger.debug(f"{e}")
+		return (self.blob[e.data_offset:e.data_offset+e.data_size], e.data_delay)
+
+	def __str__(self):
+		return f"{self.header} crc: {self.crc32:02x}"
+
+def rc4_encrypt(fw_blob):
+
+	# Round to 4096 block size
+	blob_len = len(fw_blob)
+	padded_len = (blob_len+4095)//4096 * 4096
+	fw_blob = bytearray(fw_blob)
+	fw_blob += bytearray([0]*(padded_len - blob_len))
+	a = ARC4(RC4_KEY)
+	c = Cipher(a, mode=None)
+	d = c.encryptor()
+	obuf = bytearray()
+	for i in range(padded_len):
+		obuf += d.update(fw_blob[i*512:(i+1)*512])
+	return obuf
+
+def rockchip_run(dev, fw_name, fw_blob):
+	rom = rockchip.RochipBootRom(dev)
+
+	if fw_name == 'code471':
+		logger.info("Downloading code471...")
+		blob = rc4_encrypt(fw_blob)
+		rom.write_blob(blob, 0x471)
+	elif fw_name == 'code472':
+		logger.info("Downloading code472...")
+		blob = rc4_encrypt(fw_blob)
+		rom.write_blob(blob, 0x472)
+	else:
+		fw = LoaderFile(fw_blob)
+		logger.info(f"{fw}")
+		for i in range(fw.header.entry471.count):
+			logger.info(f"Downloading entry 471 {i}...")
+			(data, delay) = fw.entry_data("471", i)
+			rom.write_blob(data, 0x471)
+			logger.info(f"Sleeping {delay}ms")
+			time.sleep(delay / 1000)
+			logger.info("Done")
+		for i in range(fw.header.entry472.count):
+			logger.info(f"Downloading entry 472 {i}...")
+			(data, delay) = fw.entry_data("472", i)
+			rom.write_blob(data, 0x472)
+			logger.info(f"Sleeping {delay}ms")
+			time.sleep(delay / 1000)
+			logger.info("Done")