UF2 (USB Flashing Format): Lua parsing library

-- This is a generated file! Please edit source .ksy file and use kaitai-struct-compiler to rebuild
--
-- This file is compatible with Lua 5.3

local class = require("class")
require("kaitaistruct")
local enum = require("enum")
local stringstream = require("string_stream")
local utils = require("utils")
local str_decode = require("string_decode")

-- 
-- UF2 is a file format, developed by Microsoft for PXT (also known as
-- Microsoft MakeCode), that is particularly suitable for flashing
-- microcontrollers over MSC (Mass Storage Class; aka removable flash drive).
-- 
-- The list of family IDs is stored in a separate JSON file:
-- 
-- <https://github.com/microsoft/uf2/blob/master/utils/uf2families.json>
-- 
-- This JSON file is regularly updated. The `family_id` enum should be kept in
-- sync with it. For simplicity and consistency, it's strongly recommended to
-- auto-generate this enum from `uf2families.json` directly. This was last done
-- using the following commands:
-- 
-- ```bash
-- $ curl -fsSLO https://github.com/microsoft/uf2/raw/90e9741f217f5a40c98ba74d663e408041037578/utils/uf2families.json
-- $ jq -r '
--   .[]
--   | "    \(.id | ascii_downcase):\n"
--   + "      id: \(.short_name | ascii_downcase | gsub("-"; "_"))\n"
--   + "      doc: \(.description | tojson)"
-- ' uf2families.json
-- ```
-- 
-- Test files, picked to cover as many different shapes as possible:
-- 
-- * <https://micropython.org/download/RPI_PICO/> - a typical case: all blocks
--   have the same family ID and `num_blocks`.
-- * <https://micropython.org/download/RPI_PICO2/> - these .uf2 files are
--   actually two UF2 files concatenated. The first UF2 file is the standalone
--   `family_id::rp2xxx_absolute` block that `picotool` prepends (see the
--   `is_rp2350_e10_block` value instance in the `block` type).
-- * <https://circuitpython.org/downloads> - the builds for SAMD boards (for
--   example
--   [Feather M0 Express](https://circuitpython.org/board/feather_m0_express/))
--   set no flags at all, so the field at offset 28 is read as `file_size` rather
--   than `family_id`.
-- * <https://github.com/raspberrypi/pico-sdk-prebuilts/releases> -
--   [Universal UF2](https://github.com/raspberrypi/pico-examples/blob/c81c855ffdedc825975a40ba357723a71358ddf0/universal/README.md#universal-binary-vs-universal-uf2)
--   files, which are again two UF2 files concatenated. What's interesting about
--   these is that the last block of the first file (with the family ID
--   `family_id::rp2040`) is moved to the end of the second file, which was
--   intended for `family_id::rp2xxx_absolute` (see
--   <https://github.com/raspberrypi/pico-examples/blob/c81c855ffdedc825975a40ba357723a71358ddf0/universal/CMakeLists.txt#L159-L169>).
-- * <https://github.com/microsoft/pxt-microsoft-boot-sequence/releases> - some
--   files contain blocks with the `flags.is_file_container` flag set (for
--   example
--   <https://github.com/microsoft/pxt-microsoft-boot-sequence/releases/download/v0.0.4/arcade-p0.uf2>
--   is a pure file container for
--   `file_name == "Projects/microsoft-boot-sequence.elf"`). Generated by
--   MakeCode, the only known producer of file container UF2s.
-- * <https://github.com/umi-eng/uftwo/tree/35bccf75b4f81c43f088696a8c4a9912f1f4104e/uftwo/tests> -
--   synthetic test files for features that real firmware doesn't seem to use
--   (e.g. MD5 checksums).
-- See also: Source (https://github.com/microsoft/uf2/blob/90e9741f217f5a40c98ba74d663e408041037578/README.md)
Uf2 = class.class(KaitaiStruct)

Uf2.ExtensionTagType = enum.Enum {
  end = 0,
  page_size = 780791,
  description = 6622621,
  rp2_ignore_block = 10049507,
  version = 10471356,
  sha2_checksum = 11824560,
  device_type_id = 13149993,
}

Uf2.FamilyId = enum.Enum {
  stm32l4 = 16738585,
  stm32l5 = 69471199,
  stm32f411xc = 114362747,
  m0sense = 299792458,
  atmega32 = 374814231,
  saml21 = 407992330,
  nrf52 = 458716255,
  esp32 = 475996592,
  stm32l1 = 505365293,
  stm32l0 = 539900561,
  stm32wl = 558239728,
  rtl8710b = 585160444,
  lpc55 = 716994540,
  esp32c2 = 730387100,
  stm32f411xe = 767756741,
  stm32g0 = 806311475,
  esp32s31 = 822212545,
  gd32f350 = 835856582,
  esp32h2 = 858203894,
  rtl8720d = 863621090,
  esp32p4 = 1026592404,
  maixplay_u4 = 1265126769,
  stm32g4 = 1282483210,
  stm32h5 = 1318001757,
  csk4 = 1332399698,
  mimxrt10xx = 1337120189,
  xr809 = 1374225320,
  stm32f7 = 1404571392,
  esp32c6 = 1410195298,
  samd51 = 1427194976,
  stm32f4 = 1467308631,
  fx2 = 1511523995,
  stm32f2 = 1561987630,
  stm32f1 = 1591873650,
  nrf52833 = 1646171002,
  stm32f0 = 1685595318,
  bk7231u = 1733968048,
  samd21 = 1760373640,
  ch32v = 1771791084,
  bk7251 = 1786956866,
  stm32f3 = 1803837832,
  stm32f407 = 1829315322,
  stm32h7 = 1840668802,
  csk6 = 1853049000,
  nrf52832xxab = 1869948536,
  stm32wb = 1892771411,
  nrf52832xxaa = 1920081230,
  max32690 = 1947226634,
  esp32c61 = 2010665156,
  bk7231n = 2067722800,
  ra4m1 = 2078840685,
  py32f071_uvk5_v3 = 2105173743,
  esp8266 = 2125160941,
  kl32l2 = 2139350931,
  nrf52820 = 2181929567,
  stm32f407vg = 2410701054,
  max78002 = 2446589208,
  rza1lu = 2501329455,
  gd32vf103 = 2599435827,
  esp32h4 = 2651564682,
  rtl8710a = 2684343619,
  at32f415 = 2697558926,
  nrf52840 = 2913282112,
  esp32h21 = 3067936943,
  esp32s2 = 3218951918,
  esp32s3 = 3296614247,
  esp32c3 = 3559628908,
  max32650 = 3594487346,
  bl602 = 3725750455,
  rtl8720c = 3767498084,
  rp2040 = 3834380118,
  rp2xxx_absolute = 3834380119,
  rp2xxx_data = 3834380120,
  rp2350_arm_s = 3834380121,
  rp2350_riscv = 3834380122,
  rp2350_arm_ns = 3834380123,
  max32666 = 4039314801,
  esp32c5 = 4145808195,
}

function Uf2:_init(io, parent, root)
  KaitaiStruct._init(self, io)
  self._parent = parent
  self._root = root or self
  self:_read()
end

function Uf2:_read()
  self.first_block = Uf2.Block(self._io, self, self._root)
  self.blocks = {}
  for i = 0, self.first_block.num_blocks - 1 - 1 do
    self.blocks[i + 1] = Uf2.Block(self._io, self, self._root)
  end
end


Uf2.Block = class.class(KaitaiStruct)

function Uf2.Block:_init(io, parent, root)
  KaitaiStruct._init(self, io)
  self._parent = parent
  self._root = root
  self:_read()
end

function Uf2.Block:_read()
  self.magic = self._io:read_bytes(4)
  if not(self.magic == "\085\070\050\010") then
    error("not equal, expected " .. "\085\070\050\010" .. ", but got " .. self.magic)
  end
  self.second_magic = self._io:read_bytes(4)
  if not(self.second_magic == "\087\081\093\158") then
    error("not equal, expected " .. "\087\081\093\158" .. ", but got " .. self.second_magic)
  end
  self.flags = Uf2.Flags(self._io, self, self._root)
  self.target_address = self._io:read_u4le()
  local _ = self.target_address
  if not(_ % 4 == 0) then
    error("ValidationExprError")
  end
  self.len_payload = self._io:read_u4le()
  local _ = self.len_payload
  if not(_ % 4 == 0) then
    error("ValidationExprError")
  end
  self.block_number = self._io:read_u4le()
  self.num_blocks_raw = self._io:read_u4le()
  if not(self.num_blocks_raw >= self.block_number + 1) then
    error("ValidationLessThanError")
  end
  if not(self.flags.has_family_id) then
    self.file_size = self._io:read_u4le()
  end
  if self.flags.has_family_id then
    self.family_id = Uf2.FamilyId(self._io:read_u4le())
  end
  self._raw_data = self._io:read_bytes(476)
  local _io = KaitaiStream(stringstream(self._raw_data))
  self.data = Uf2.BlockData(_io, self, self._root)
  self.final_magic = self._io:read_bytes(4)
  if not(self.final_magic == "\048\111\177\010") then
    error("not equal, expected " .. "\048\111\177\010" .. ", but got " .. self.final_magic)
  end
end

-- 
-- Determines whether this is a block that `picotool` prepends to RP2350
-- flash images as a workaround for erratum RP2350-E10 (i.e. a hardware
-- bug in the A2 version of the RP2350 boot ROM).
-- 
-- Such a block is always written on its own, but its `num_blocks_raw` is
-- set to 2. If we trusted this value, we would attempt to read one block
-- too many (which would most likely fail). Therefore, we must correct it
-- to 1 before using it.
-- 
-- The conditions for detecting this block come from the
-- [`check_abs_block()`](https://github.com/raspberrypi/picotool/blob/6f6458d792b93685a11423b244a585eaa99eafcf/elf2uf2/elf2uf2.cpp#L147)
-- function in `picotool`. However, there are some differences:
-- 
-- 1. In Kaitai Struct, we cannot easily check whether all 256 payload
--    bytes are set to `0xef`, so we only check the first and last bytes.
-- 2. There are .uf2 files in the wild where `flags.has_extension_tags`
--    is true, but there are actually no extension tags (the first and
--    only tag has a size of 0, which is just a terminator), so our
--    condition allows for this case. You can download an example of such
--    a .uf2 file here:
--    <https://github.com/neednotapply/DC32-cfw/releases/tag/1.69.13.37>
-- 
-- It's worth noting that we cannot require the presence of the
-- `extension_tag_type::rp2_ignore_block`
-- (`UF2_EXTENSION_RP2_IGNORE_BLOCK`) tag because the UF2 files generated
-- by `picotool` prior to
-- <https://github.com/raspberrypi/picotool/commit/78c9bd121b09399823b67ee7ea89003ca0d3315f>
-- don't have it. Therefore, we check whether this tag is present only if
-- `flags.has_extension_tags` is set.
-- 
-- Test .uf2 files with this special block can be downloaded from
-- <https://micropython.org/download/RPI_PICO2/>. Note that all the .uf2
-- files there are actually two UF2 (sub)files concatenated, and this
-- Kaitai Struct implementation parses only one at a time (so in order to
-- parse both, you need something like the helper spec `uf2_files.ksy`
-- from
-- <https://github.com/kaitai-io/kaitai_struct_formats/pull/542#discussion_r3906386820>).
-- v1.24.x releases predate the `extension_tag_type::rp2_ignore_block`
-- tag, while releases v1.25.0 and later include it.
-- See also: Git tag "2.3.0" (https://github.com/raspberrypi/picotool/blob/6f6458d792b93685a11423b244a585eaa99eafcf/elf2uf2/elf2uf2.cpp#L147)
-- See also: Source (https://github.com/raspberrypi/picotool/commit/78c9bd121b09399823b67ee7ea89003ca0d3315f)
Uf2.Block.property.is_rp2350_e10_block = {}
function Uf2.Block.property.is_rp2350_e10_block:get()
  if self._m_is_rp2350_e10_block ~= nil then
    return self._m_is_rp2350_e10_block
  end

  self._m_is_rp2350_e10_block =  (( ((self.flags.value == 8192) or (self.flags.value == 40960)) ) and (self.family_id == Uf2.FamilyId.rp2xxx_absolute) and (self.num_blocks_raw == 2) and (self.block_number == 0) and (self.len_payload == 256) and (string.byte(self.data.payload, 1) == 239) and (string.byte(self.data.payload, #self.data.payload) == 239) and ( ((not(self.flags.has_extension_tags)) or (self.data.extension_tags[0 + 1].len_tag == 0) or ( ((self.data.extension_tags[0 + 1].len_tag == 4) and (self.data.extension_tags[0 + 1].tag_type == Uf2.ExtensionTagType.rp2_ignore_block)) )) )) 
  return self._m_is_rp2350_e10_block
end

Uf2.Block.property.num_blocks = {}
function Uf2.Block.property.num_blocks:get()
  if self._m_num_blocks ~= nil then
    return self._m_num_blocks
  end

  self._m_num_blocks = utils.box_unwrap((self.is_rp2350_e10_block) and utils.box_wrap(1) or (self.num_blocks_raw))
  return self._m_num_blocks
end

-- 
-- Address in flash where `data.payload` should be written, or an offset
-- in the file specified by `data.file_name` if `flags.is_file_container`
-- is set.
-- 
-- The [official
-- spec](https://github.com/microsoft/uf2/blob/90e9741f217f5a40c98ba74d663e408041037578/README.md#payload-sizes)
-- says:
-- 
-- > In any event, payload size and target address should always be
-- > 4-byte aligned.
-- 
-- The [official
-- spec](https://github.com/microsoft/uf2/blob/90e9741f217f5a40c98ba74d663e408041037578/README.md#payload-sizes)
-- says:
-- 
-- > In any event, payload size and target address should always be
-- > 4-byte aligned.
-- 
-- Number of blocks that make up the UF2 file to which this block
-- belongs. Every block of a file has the same value. It is at least 1
-- and every `block_number` in the file must be less than this value
-- (which is validated by this Kaitai Struct implementation).
-- 
-- Size of the file this block belongs to, but only if
-- `flags.is_file_container` is true. Otherwise, the official spec allows
-- this field to be set to anything - though in practice, it's always
-- zero.

Uf2.BlockData = class.class(KaitaiStruct)

function Uf2.BlockData:_init(io, parent, root)
  KaitaiStruct._init(self, io)
  self._parent = parent
  self._root = root
  self:_read()
end

function Uf2.BlockData:_read()
  self.payload = self._io:read_bytes(self._parent.len_payload)
  if self._parent.flags.is_file_container then
    self.file_name = str_decode.decode(self._io:read_bytes_term(0, false, true, true), "UTF-8")
  end
  if self._parent.flags.has_extension_tags then
    self.extension_tags = {}
    local i = 0
    while true do
      local _ = Uf2.ExtensionTag(self._io, self, self._root)
      self.extension_tags[i + 1] = _
      if _.len_tag == 0 then
        break
      end
      i = i + 1
    end
  end
end

-- 
-- Describes a region that doesn't need to be flashed again if the
-- checksum matches.
-- 
-- The [official
-- spec](https://github.com/microsoft/uf2/blob/90e9741f217f5a40c98ba74d663e408041037578/README.md#md5-checksum)
-- says that "This is currently only used on ESP32", but no real-world
-- firmware sample has been found (at least none of the .uf2 files from
-- <https://github.com/adafruit/tinyuf2/releases>,
-- [MicroPython](https://micropython.org/download/) or
-- [CircuitPython](https://circuitpython.org/downloads) contain it). It
-- was found only in some synthetic test files, e.g.
-- <https://github.com/umi-eng/uftwo/blob/35bccf75b4f81c43f088696a8c4a9912f1f4104e/uftwo/tests/checksum_256.uf2>.
Uf2.BlockData.property.md5_checksum = {}
function Uf2.BlockData.property.md5_checksum:get()
  if self._m_md5_checksum ~= nil then
    return self._m_md5_checksum
  end

  if self._parent.flags.has_md5_checksum then
    local _pos = self._io:pos()
    self._io:seek(self._io:size() - 24)
    self._m_md5_checksum = Uf2.Md5Checksum(self._io, self, self._root)
    self._io:seek(_pos)
  end
  return self._m_md5_checksum
end

-- 
-- The bytes to be written to `_parent.target_address`, which is either
-- an address in flash, or an offset in the file specified by `file_name`
-- if `_parent.flags.is_file_container` is set.

-- 
-- See also: Source (https://github.com/microsoft/uf2/blob/90e9741f217f5a40c98ba74d663e408041037578/README.md#extension-tags)
Uf2.ExtensionTag = class.class(KaitaiStruct)

function Uf2.ExtensionTag:_init(io, parent, root)
  KaitaiStruct._init(self, io)
  self._parent = parent
  self._root = root
  self:_read()
end

function Uf2.ExtensionTag:_read()
  self.len_tag = self._io:read_u1()
  local _ = self.len_tag
  if not( ((_ == 0) or (_ >= self.min_len_tag)) ) then
    error("ValidationExprError")
  end
  self.tag_type = Uf2.ExtensionTagType(self._io:read_bits_int_le(24))
  self._io:align_to_byte()
  if self.len_tag ~= 0 then
    self.value = self._io:read_bytes(self.len_value)
  end
  self.padding = self._io:read_bytes(-(self.len_tag) % 4)
end

Uf2.ExtensionTag.property.len_value = {}
function Uf2.ExtensionTag.property.len_value:get()
  if self._m_len_value ~= nil then
    return self._m_len_value
  end

  self._m_len_value = utils.box_unwrap((self.len_tag >= self.min_len_tag) and utils.box_wrap(self.len_tag - self.min_len_tag) or (0))
  return self._m_len_value
end

Uf2.ExtensionTag.property.min_len_tag = {}
function Uf2.ExtensionTag.property.min_len_tag:get()
  if self._m_min_len_tag ~= nil then
    return self._m_min_len_tag
  end

  self._m_min_len_tag = 1 + 3
  return self._m_min_len_tag
end

-- 
-- Total size of the tag in bytes, including this byte and `tag_type`, so
-- at least 4. The exception is the last tag which terminates the list -
-- it specifies a total size of 0.
-- 
-- Tags are 4-byte aligned, so a tag whose size is not a multiple
-- of 4 is followed by padding.

-- 
-- See also: Source (https://github.com/microsoft/uf2/blob/90e9741f217f5a40c98ba74d663e408041037578/uf2.h#L43-L47)
-- See also: Git tag "2.3.0" (https://github.com/raspberrypi/pico-sdk/blob/98a542c1a62fb549ffb5d66a3e5892b06276b670/src/common/boot_uf2_headers/include/boot/uf2.h#L23-L27)
Uf2.Flags = class.class(KaitaiStruct)

function Uf2.Flags:_init(io, parent, root)
  KaitaiStruct._init(self, io)
  self._parent = parent
  self._root = root
  self:_read()
end

function Uf2.Flags:_read()
  self.value = self._io:read_u4le()
  local _ = self.value
  if not( ((_ & ~61441 == 0) and (not( ((self.is_file_container) and (self.has_extension_tags)) ))) ) then
    error("ValidationExprError")
  end
end

-- 
-- Indicates whether extension tags are present after the payload.
Uf2.Flags.property.has_extension_tags = {}
function Uf2.Flags.property.has_extension_tags:get()
  if self._m_has_extension_tags ~= nil then
    return self._m_has_extension_tags
  end

  self._m_has_extension_tags = self.value & 32768 ~= 0
  return self._m_has_extension_tags
end

-- 
-- The field at offset 28 in the block is `family_id` instead of
-- `file_size`.
Uf2.Flags.property.has_family_id = {}
function Uf2.Flags.property.has_family_id:get()
  if self._m_has_family_id ~= nil then
    return self._m_has_family_id
  end

  self._m_has_family_id = self.value & 8192 ~= 0
  return self._m_has_family_id
end

-- 
-- Indicates whether `md5_checksum` is present at the end of `data`.
Uf2.Flags.property.has_md5_checksum = {}
function Uf2.Flags.property.has_md5_checksum:get()
  if self._m_has_md5_checksum ~= nil then
    return self._m_has_md5_checksum
  end

  self._m_has_md5_checksum = self.value & 16384 ~= 0
  return self._m_has_md5_checksum
end

-- 
-- When set, the UF2 format is used as a container for regular files
-- (akin to a TAR file, or ZIP archive without compression).
-- 
-- `target_address` is the offset in the file where the payload is to be
-- written, and `file_size` is the size of that file. The name of the
-- destination file is stored in `data.file_name`.
Uf2.Flags.property.is_file_container = {}
function Uf2.Flags.property.is_file_container:get()
  if self._m_is_file_container ~= nil then
    return self._m_is_file_container
  end

  self._m_is_file_container = self.value & 4096 ~= 0
  return self._m_is_file_container
end

-- 
-- Indicates that this block should be skipped when writing the device
-- flash. It can be used to store data that does not fit on the device,
-- typically embedded source code or debug info.
Uf2.Flags.property.not_main_flash = {}
function Uf2.Flags.property.not_main_flash:get()
  if self._m_not_main_flash ~= nil then
    return self._m_not_main_flash
  end

  self._m_not_main_flash = self.value & 1 ~= 0
  return self._m_not_main_flash
end

-- 
-- Only the five bits that we cover in value instances below are defined,
-- and no other bit may be set. The [official
-- spec](https://github.com/microsoft/uf2/blob/90e9741f217f5a40c98ba74d663e408041037578/README.md#flags)
-- says "Currently, there are five flags defined". If any other flags are
-- added in the future, this .ksy spec will need to be updated.
-- 
-- The `is_file_container` and `has_extension_tags` flags disagree about
-- what follows the payload in `data` (a file name or a list of extension
-- tags), so this Kaitai Struct implementation treats them as mutually
-- exclusive and will reject a block that sets both. The official spec
-- does not specify how this should be handled, but logically there's a
-- conflict, so we've decided to strictly treat it as an error.

-- 
-- See also: Source (https://github.com/microsoft/uf2/blob/90e9741f217f5a40c98ba74d663e408041037578/README.md#md5-checksum)
Uf2.Md5Checksum = class.class(KaitaiStruct)

function Uf2.Md5Checksum:_init(io, parent, root)
  KaitaiStruct._init(self, io)
  self._parent = parent
  self._root = root
  self:_read()
end

function Uf2.Md5Checksum:_read()
  self.start_address = self._io:read_u4le()
  self.len_region = self._io:read_u4le()
  self.md5 = self._io:read_bytes(16)
end