140 lines
4.4 KiB
Python
140 lines
4.4 KiB
Python
# Copyright (c) Kuba Szczodrzyński 2022-05-27.
|
|
|
|
from math import ceil
|
|
from typing import Dict
|
|
|
|
from models import Flags, Tag
|
|
|
|
from tools.util.intbin import align_up, intto8, inttole24, inttole32, letoint
|
|
from tools.util.models import Family
|
|
from tools.util.platform import get_family
|
|
|
|
|
|
class Block:
|
|
flags: Flags
|
|
|
|
address: int = 0
|
|
length: int = 0
|
|
|
|
block_seq: int = 0
|
|
block_count: int = 0
|
|
|
|
file_size: int = 0
|
|
family: Family
|
|
|
|
data: bytes = None
|
|
md5_data: bytes = None
|
|
tags: Dict[Tag, bytes] = {}
|
|
|
|
def __init__(self, family: Family = None) -> None:
|
|
self.flags = Flags()
|
|
self.family = family
|
|
self.flags.has_family_id = not not self.family
|
|
|
|
def encode(self) -> bytes:
|
|
self.flags.has_tags = not not self.tags
|
|
# UF2 magic 1 and 2
|
|
data = b"\x55\x46\x32\x0A\x57\x51\x5D\x9E"
|
|
# encode integer variables
|
|
data += inttole32(self.flags.encode())
|
|
data += inttole32(self.address)
|
|
data += inttole32(self.length)
|
|
data += inttole32(self.block_seq)
|
|
data += inttole32(self.block_count)
|
|
if self.flags.file_container:
|
|
data += inttole32(self.file_size)
|
|
elif self.flags.has_family_id:
|
|
data += inttole32(self.family.id)
|
|
else:
|
|
data += b"\x00\x00\x00\x00"
|
|
if not self.data:
|
|
self.data = b""
|
|
# append tags
|
|
tags = b""
|
|
if self.flags.has_tags:
|
|
for k, v in self.tags.items():
|
|
tag_size = 4 + len(v)
|
|
tags += intto8(tag_size)
|
|
tags += inttole24(k.value)
|
|
tags += v
|
|
tag_size %= 4
|
|
if tag_size:
|
|
tags += b"\x00" * (4 - tag_size)
|
|
# append block data with padding
|
|
data += self.data
|
|
data += tags
|
|
data += b"\x00" * (476 - len(self.data) - len(tags))
|
|
data += b"\x30\x6F\xB1\x0A" # magic 3
|
|
return data
|
|
|
|
def decode(self, data: bytes) -> bool:
|
|
# check block size
|
|
if len(data) != 512:
|
|
print(f"Invalid block size ({len(data)})")
|
|
return False
|
|
# check Magic 1
|
|
if letoint(data[0:4]) != 0x0A324655:
|
|
print(f"Invalid Magic 1 ({data[0:4]})")
|
|
return False
|
|
# check Magic 2
|
|
if letoint(data[4:8]) != 0x9E5D5157:
|
|
print(f"Invalid Magic 2 ({data[4:8]})")
|
|
return False
|
|
# check Magic 3
|
|
if letoint(data[508:512]) != 0x0AB16F30:
|
|
print(f"Invalid Magic 13({data[508:512]})")
|
|
return False
|
|
|
|
self.flags.decode(letoint(data[8:12]))
|
|
self.address = letoint(data[12:16])
|
|
self.length = letoint(data[16:20])
|
|
self.block_seq = letoint(data[20:24])
|
|
self.block_count = letoint(data[24:28])
|
|
if self.flags.file_container:
|
|
self.file_size = letoint(data[28:32])
|
|
if self.flags.has_family_id:
|
|
self.family = get_family(id=letoint(data[28:32]))
|
|
|
|
if self.flags.has_md5:
|
|
self.md5_data = data[484:508] # last 24 bytes of data[]
|
|
|
|
# decode tags
|
|
self.tags = {}
|
|
if self.flags.has_tags:
|
|
tags = data[32 + self.length :]
|
|
i = 0
|
|
while i < len(tags):
|
|
length = tags[i]
|
|
if not length:
|
|
break
|
|
tag_type = letoint(tags[i + 1 : i + 4])
|
|
tag_data = tags[i + 4 : i + length]
|
|
self.tags[Tag(tag_type)] = tag_data
|
|
i += length
|
|
i = int(ceil(i / 4) * 4)
|
|
|
|
self.data = data[32 : 32 + self.length]
|
|
return True
|
|
|
|
@staticmethod
|
|
def get_tags_length(tags: Dict[Tag, bytes]) -> int:
|
|
out = 0
|
|
# add tag headers
|
|
out += 4 * len(tags)
|
|
# add all tag lengths, padded to 4 bytes
|
|
out += sum(align_up(l, 4) for l in map(len, tags.values()))
|
|
# add final 0x00 tag
|
|
out += 4
|
|
return out
|
|
|
|
def __str__(self) -> str:
|
|
flags = self.flags
|
|
address = hex(self.address)
|
|
length = hex(self.length)
|
|
block_seq = self.block_seq
|
|
block_count = self.block_count
|
|
file_size = self.file_size
|
|
family = self.family.short_name
|
|
tags = [(k.name, v) for k, v in self.tags.items()]
|
|
return f"Block[{block_seq}/{block_count}](flags={flags}, address={address}, length={length}, file_size={file_size}, family={family}, tags={tags})"
|