Files
libretiny/tools/uf2ota/dump.py
2022-06-02 22:22:23 +02:00

101 lines
2.9 KiB
Python

# Copyright (c) Kuba Szczodrzyński 2022-05-28.
from io import BytesIO, FileIO
from os import makedirs
from os.path import join
from typing import Dict, Tuple
from models import Opcode, Tag
from uf2 import UF2
from tools.util.intbin import inttole32, letoint, letosint
fs: Dict[str, Tuple[int, FileIO]] = {}
output_dir = ""
output_basename = ""
part1 = ""
part2 = ""
def write(part: str, offs: int, data: bytes):
global fs
if part not in fs or fs[part][0] != offs:
path = join(output_dir, output_basename + part + f"_0x{offs:x}.bin")
f = open(path, "wb")
if part in fs:
fs[part][1].close()
else:
f = fs[part][1]
fs[part] = (offs + f.write(data), f)
def update_parts(tags: Dict[Tag, bytes]):
global part1, part2
if Tag.LT_PART_1 in tags:
part1 = tags[Tag.LT_PART_1].decode()
part1 = ("1_" + part1) if part1 else None
if Tag.LT_PART_2 in tags:
part2 = tags[Tag.LT_PART_2].decode()
part2 = ("2_" + part2) if part2 else None
def uf2_dump(uf2: UF2, outdir: str):
global output_dir, output_basename
makedirs(outdir, exist_ok=True)
if Tag.LT_VERSION not in uf2.tags:
raise RuntimeError("Can only dump LibreTuya firmware images")
output_dir = outdir
output_basename = "_".join(
filter(
None,
[
uf2.tags.get(Tag.FIRMWARE, b"").decode(),
uf2.tags.get(Tag.VERSION, b"").decode(),
"lt" + uf2.tags[Tag.LT_VERSION].decode(),
uf2.tags.get(Tag.BOARD, b"").decode(),
],
)
)
output_basename += "_"
update_parts(uf2.tags)
for block in uf2.data:
# update target partition info
update_parts(block.tags)
# skip empty blocks
if not block.length:
continue
data1 = block.data if part1 else None
data2 = block.data if part2 else None
if Tag.LT_BINPATCH in block.tags:
# type 5, 6
data2 = bytearray(data2)
tag = block.tags[Tag.LT_BINPATCH]
binpatch = BytesIO(tag)
while binpatch.tell() < len(tag):
opcode = Opcode(binpatch.read(1)[0])
length = binpatch.read(1)[0]
data = binpatch.read(length)
if opcode == Opcode.DIFF32:
value = letosint(data[0:4])
for offs in data[4:]:
chunk = data2[offs : offs + 4]
chunk = letoint(chunk)
chunk += value
chunk = inttole32(chunk)
data2[offs : offs + 4] = chunk
data2 = bytes(data2)
if data1:
# types 1, 3, 4
write(part1, block.address, data1)
if data2:
# types 2, 3, 4
write(part2, block.address, data2)