1# SPDX-FileCopyrightText: 2020-2023 Espressif Systems (Shanghai) CO LTD
2# SPDX-License-Identifier: GPL-2.0-or-later
3# Code was originally licensed under Apache 2.0 before the release of ESP-IDF v5.2
4
5import hashlib
6import os
7import struct
8from typing import List
9
10from esptool.util import div_roundup
11
12
13class UF2Writer(object):
14    # The UF2 format is described here: https://github.com/microsoft/uf2
15    UF2_BLOCK_SIZE = 512
16    # max value of CHUNK_SIZE reduced by optional parts. Currently, MD5_PART only.
17    UF2_DATA_SIZE = 476
18    UF2_MD5_PART_SIZE = 24
19    UF2_FIRST_MAGIC = 0x0A324655
20    UF2_SECOND_MAGIC = 0x9E5D5157
21    UF2_FINAL_MAGIC = 0x0AB16F30
22    UF2_FLAG_FAMILYID_PRESENT = 0x00002000
23    UF2_FLAG_MD5_PRESENT = 0x00004000
24
25    def __init__(
26        self,
27        chip_id: int,
28        output_file: os.PathLike,
29        chunk_size: int,
30        md5_enabled: bool = True,
31    ) -> None:
32        if not md5_enabled:
33            self.UF2_MD5_PART_SIZE = 0
34            self.UF2_FLAG_MD5_PRESENT = 0x00000000
35        self.md5_enabled = md5_enabled
36        self.chip_id = chip_id
37        self.CHUNK_SIZE = (
38            self.UF2_DATA_SIZE - self.UF2_MD5_PART_SIZE
39            if chunk_size is None
40            else chunk_size
41        )
42        self.f = open(output_file, "wb")
43
44    def __enter__(self) -> "UF2Writer":
45        return self
46
47    def __exit__(self, exc_type: str, exc_val: int, exc_tb: List) -> None:
48        if self.f:
49            self.f.close()
50
51    @staticmethod
52    def _to_uint32(num: int) -> bytes:
53        return struct.pack("<I", num)
54
55    def _write_block(
56        self, addr: int, chunk: bytes, len_chunk: int, block_no: int, blocks: int
57    ) -> None:
58        assert len_chunk > 0
59        assert len_chunk <= self.CHUNK_SIZE
60        assert block_no < blocks
61        block = struct.pack(
62            "<IIIIIIII",
63            self.UF2_FIRST_MAGIC,
64            self.UF2_SECOND_MAGIC,
65            self.UF2_FLAG_FAMILYID_PRESENT | self.UF2_FLAG_MD5_PRESENT,
66            addr,
67            len_chunk,
68            block_no,
69            blocks,
70            self.chip_id,
71        )
72        block += chunk
73
74        if self.md5_enabled:
75            md5_part = struct.pack("<II", addr, len_chunk)
76            md5_part += hashlib.md5(chunk).digest()
77            assert len(md5_part) == self.UF2_MD5_PART_SIZE
78
79            block += md5_part
80        block += b"\x00" * (self.UF2_DATA_SIZE - self.UF2_MD5_PART_SIZE - len_chunk)
81        block += self._to_uint32(self.UF2_FINAL_MAGIC)
82        assert len(block) == self.UF2_BLOCK_SIZE
83        self.f.write(block)
84
85    def add_file(self, addr: int, image: bytes) -> None:
86        blocks = div_roundup(len(image), self.CHUNK_SIZE)
87        chunks = [
88            image[i : i + self.CHUNK_SIZE]
89            for i in range(0, len(image), self.CHUNK_SIZE)
90        ]
91        for i, chunk in enumerate(chunks):
92            len_chunk = len(chunk)
93            self._write_block(addr, chunk, len_chunk, i, blocks)
94            addr += len_chunk
95