1import hashlib
2import os
3import os.path
4import re
5import struct
6import subprocess
7import sys
8
9from conftest import need_to_install_package_err
10
11from elftools.elf.elffile import ELFFile
12
13import pytest
14
15TEST_DIR = os.path.join(os.path.abspath(os.path.dirname(__file__)), "elf2image")
16
17try:
18    import esptool
19except ImportError:
20    need_to_install_package_err()
21
22
23def try_delete(path):
24    try:
25        os.remove(path)
26    except OSError:
27        pass
28
29
30def segment_matches_section(segment, section):
31    """segment is an ImageSegment from an esptool binary.
32    section is an elftools ELF section
33
34    Returns True if they match
35    """
36    sh_size = (section.header.sh_size + 0x3) & ~3  # pad length of ELF sections
37    return section.header.sh_addr == segment.addr and sh_size == len(segment.data)
38
39
40@pytest.mark.host_test
41class BaseTestCase:
42    @classmethod
43    def setup_class(self):
44        # Save the current working directory to be restored later
45        self.stored_dir = os.getcwd()
46        os.chdir(TEST_DIR)
47
48    @classmethod
49    def teardown_class(self):
50        # Restore the stored working directory
51        os.chdir(self.stored_dir)
52
53    def assertEqualHex(self, expected, actual, message=None):
54        try:
55            expected = hex(expected)
56        except TypeError:  # if expected is character
57            expected = hex(ord(expected))
58        try:
59            actual = hex(actual)
60        except TypeError:  # if actual is character
61            actual = hex(ord(actual))
62        assert expected == actual, message
63
64    def assertImageDoesNotContainSection(self, image, elf, section_name):
65        """
66        Assert an esptool binary image object does not
67        contain the data for a particular ELF section.
68        """
69        with open(elf, "rb") as f:
70            e = ELFFile(f)
71            section = e.get_section_by_name(section_name)
72            assert section, f"{section_name} should be in the ELF"
73            sh_addr = section.header.sh_addr
74            data = section.data()
75            # no section should start at the same address as the ELF section.
76            for seg in sorted(image.segments, key=lambda s: s.addr):
77                print(
78                    f"comparing seg {seg.addr:#x} sec {sh_addr:#x} len {len(data):#x}"
79                )
80                assert (
81                    seg.addr != sh_addr
82                ), f"{section_name} should not be in the binary image"
83
84    def assertImageContainsSection(self, image, elf, section_name):
85        """
86        Assert an esptool binary image object contains
87        the data for a particular ELF section.
88        """
89        with open(elf, "rb") as f:
90            e = ELFFile(f)
91            section = e.get_section_by_name(section_name)
92            assert section, f"{section_name} should be in the ELF"
93            sh_addr = section.header.sh_addr
94            data = section.data()
95            # section contents may be smeared across multiple image segments,
96            # so look through each segment and remove it from ELF section 'data'
97            # as we find it in the image segments. When we're done 'data' should
98            # all be accounted for
99            for seg in sorted(image.segments, key=lambda s: s.addr):
100                print(
101                    f"comparing seg {seg.addr:#x} sec {sh_addr:#x} len {len(data):#x}"
102                )
103                if seg.addr == sh_addr:
104                    overlap_len = min(len(seg.data), len(data))
105                    assert (
106                        data[:overlap_len] == seg.data[:overlap_len]
107                    ), f"ELF '{section_name}' section has mis-matching bin image data"
108                    sh_addr += overlap_len
109                    data = data[overlap_len:]
110
111            # no bytes in 'data' should be left unmatched
112            assert len(data) == 0, (
113                f"ELF {elf} section '{section_name}' has no encompassing"
114                f" segment(s) in bin image (image segments: {image.segments})"
115            )
116
117    def assertImageInfo(self, binpath, chip="esp8266", assert_sha=False):
118        """
119        Run esptool.py image_info on a binary file,
120        assert no red flags about contents.
121        """
122        cmd = [sys.executable, "-m", "esptool", "--chip", chip, "image_info", binpath]
123        try:
124            output = subprocess.check_output(cmd)
125            output = output.decode("utf-8")
126            print(output)
127        except subprocess.CalledProcessError as e:
128            print(e.output)
129            raise
130        assert re.search(
131            r"Checksum: [a-fA-F0-9]{2} \(valid\)", output
132        ), "Checksum calculation should be valid"
133        if assert_sha:
134            assert re.search(
135                r"Validation Hash: [a-fA-F0-9]{64} \(valid\)", output
136            ), "SHA256 should be valid"
137        assert (
138            "warning" not in output.lower()
139        ), "Should be no warnings in image_info output"
140
141    def run_elf2image(self, chip, elf_path, version=None, extra_args=[]):
142        """Run elf2image on elf_path"""
143        cmd = [sys.executable, "-m", "esptool", "--chip", chip, "elf2image"]
144        if version is not None:
145            cmd += ["--version", str(version)]
146        cmd += [elf_path] + extra_args
147        print("\nExecuting {}".format(" ".join(cmd)))
148        try:
149            output = subprocess.check_output(cmd)
150            output = output.decode("utf-8")
151            print(output)
152            assert (
153                "warning" not in output.lower()
154            ), "elf2image should not output warnings"
155        except subprocess.CalledProcessError as e:
156            print(e.output)
157            raise
158
159
160class TestESP8266V1Image(BaseTestCase):
161    ELF = "esp8266-nonosssdk20-iotdemo.elf"
162    BIN_LOAD = "esp8266-nonosssdk20-iotdemo.elf-0x00000.bin"
163    BIN_IROM = "esp8266-nonosssdk20-iotdemo.elf-0x10000.bin"
164
165    @classmethod
166    def setup_class(self):
167        super(TestESP8266V1Image, self).setup_class()
168        self.run_elf2image(self, "esp8266", self.ELF, 1)
169
170    @classmethod
171    def teardown_class(self):
172        super(TestESP8266V1Image, self).teardown_class()
173        try_delete(self.BIN_LOAD)
174        try_delete(self.BIN_IROM)
175
176    def test_irom_bin(self):
177        with open(self.ELF, "rb") as f:
178            e = ELFFile(f)
179            irom_section = e.get_section_by_name(".irom0.text")
180            assert (
181                irom_section.header.sh_size == os.stat(self.BIN_IROM).st_size
182            ), "IROM raw binary file should be same length as .irom0.text section"
183
184    def test_loaded_sections(self):
185        image = esptool.bin_image.LoadFirmwareImage("esp8266", self.BIN_LOAD)
186        # Adjacent sections are now merged, len(image.segments) should
187        # equal 2 (instead of 3).
188        assert len(image.segments) == 2
189        self.assertImageContainsSection(image, self.ELF, ".data")
190        self.assertImageContainsSection(image, self.ELF, ".text")
191        # Section .rodata is merged in the binary with the previous one,
192        # so it won't be found in the binary image.
193        self.assertImageDoesNotContainSection(image, self.ELF, ".rodata")
194
195
196class TestESP8266V12SectionHeaderNotAtEnd(BaseTestCase):
197    """Ref https://github.com/espressif/esptool/issues/197 -
198    this ELF image has the section header not at the end of the file"""
199
200    ELF = "esp8266-nonossdkv12-example.elf"
201    BIN_LOAD = ELF + "-0x00000.bin"
202    BIN_IROM = ELF + "-0x40000.bin"
203
204    @classmethod
205    def teardown_class(self):
206        try_delete(self.BIN_LOAD)
207        try_delete(self.BIN_IROM)
208
209    def test_elf_section_header_not_at_end(self):
210        self.run_elf2image("esp8266", self.ELF)
211        image = esptool.bin_image.LoadFirmwareImage("esp8266", self.BIN_LOAD)
212        assert len(image.segments) == 3
213        self.assertImageContainsSection(image, self.ELF, ".data")
214        self.assertImageContainsSection(image, self.ELF, ".text")
215        self.assertImageContainsSection(image, self.ELF, ".rodata")
216
217
218class TestESP8266V2Image(BaseTestCase):
219    def _test_elf2image(self, elfpath, binpath, mergedsections=[]):
220        try:
221            self.run_elf2image("esp8266", elfpath, 2)
222            image = esptool.bin_image.LoadFirmwareImage("esp8266", binpath)
223            print("In test_elf2image", len(image.segments))
224            assert 4 - len(mergedsections) == len(image.segments)
225            sections = [".data", ".text", ".rodata"]
226            # Remove the merged sections from the `sections` list
227            sections = [sec for sec in sections if sec not in mergedsections]
228            for sec in sections:
229                self.assertImageContainsSection(image, elfpath, sec)
230            for sec in mergedsections:
231                self.assertImageDoesNotContainSection(image, elfpath, sec)
232
233            irom_segment = image.segments[0]
234            assert irom_segment.addr == 0, "IROM segment 'load address' should be zero"
235            with open(elfpath, "rb") as f:
236                e = ELFFile(f)
237                sh_size = (
238                    e.get_section_by_name(".irom0.text").header.sh_size + 15
239                ) & ~15
240                assert len(irom_segment.data) == sh_size, (
241                    f"irom segment ({len(irom_segment.data):#x}) should be same size "
242                    f"(16 padded) as .irom0.text section ({sh_size:#x})"
243                )
244
245            # check V2 CRC (for ESP8266 SDK bootloader)
246            with open(binpath, "rb") as f:
247                f.seek(-4, os.SEEK_END)
248                image_len = f.tell()
249                crc_stored = struct.unpack("<I", f.read(4))[0]
250                f.seek(0)
251                crc_calc = esptool.bin_image.esp8266_crc32(f.read(image_len))
252                assert crc_stored == crc_calc
253
254            # test imageinfo doesn't fail
255            self.assertImageInfo(binpath)
256
257        finally:
258            try_delete(binpath)
259
260    def test_nonossdkimage(self):
261        ELF = "esp8266-nonossdkv20-at-v2.elf"
262        BIN = "esp8266-nonossdkv20-at-v2-0x01000.bin"
263        self._test_elf2image(ELF, BIN)
264
265    def test_espopenrtosimage(self):
266        ELF = "esp8266-openrtos-blink-v2.elf"
267        BIN = "esp8266-openrtos-blink-v2-0x02000.bin"
268        # .rodata section is merged with the previous one: .data
269        self._test_elf2image(ELF, BIN, [".rodata"])
270
271
272class TestESP32Image(BaseTestCase):
273    def _test_elf2image(self, elfpath, binpath, extra_args=[]):
274        try:
275            self.run_elf2image("esp32", elfpath, extra_args=extra_args)
276            image = esptool.bin_image.LoadFirmwareImage("esp32", binpath)
277            self.assertImageInfo(
278                binpath,
279                "esp32",
280                True if "--ram-only-header" not in extra_args else False,
281            )
282            return image
283        finally:
284            try_delete(binpath)
285
286    def test_bootloader(self):
287        ELF = "esp32-bootloader.elf"
288        BIN = "esp32-bootloader.bin"
289        image = self._test_elf2image(ELF, BIN)
290        assert len(image.segments) == 3
291        for section in [".iram1.text", ".iram_pool_1.text", ".dram0.rodata"]:
292            self.assertImageContainsSection(image, ELF, section)
293
294    def test_app_template(self):
295        ELF = "esp32-app-template.elf"
296        BIN = "esp32-app-template.bin"
297        image = self._test_elf2image(ELF, BIN)
298        # Adjacent sections are now merged, len(image.segments) should
299        # equal 5 (instead of 6).
300        assert len(image.segments) == 5
301        # the other segment is a padding or merged segment
302        for section in [
303            ".iram0.vectors",
304            ".dram0.data",
305            ".flash.rodata",
306            ".flash.text",
307        ]:
308            self.assertImageContainsSection(image, ELF, section)
309        # check that merged sections are not in the binary image
310        for mergedsection in [".iram0.text"]:
311            self.assertImageDoesNotContainSection(image, ELF, mergedsection)
312
313    def test_too_many_sections(self, capsys):
314        ELF = "esp32-too-many-sections.elf"
315        BIN = "esp32-too-many-sections.bin"
316        with pytest.raises(subprocess.CalledProcessError):
317            self._test_elf2image(ELF, BIN)
318        output = capsys.readouterr().out
319        assert "max 16" in output
320        assert "linker script" in output
321
322    def test_use_segments(self):
323        ELF = "esp32-zephyr.elf"
324        BIN = "esp32-zephyr.bin"
325        # default behaviour uses ELF sections,
326        # this ELF will produce 8 segments in the bin
327        image = self._test_elf2image(ELF, BIN)
328        # Adjacent sections are now merged, len(image.segments) should
329        # equal 5 (instead of 8).
330        assert len(image.segments) == 5
331
332        # --use_segments uses ELF segments(phdrs), produces just 2 segments in the bin
333        image = self._test_elf2image(ELF, BIN, ["--use_segments"])
334        assert len(image.segments) == 2
335
336    def test_ram_only_header(self):
337        ELF = "esp32-app-template.elf"
338        BIN = "esp32-app-template.bin"
339        # --ram-only-header produces just 2 visible segments in the bin
340        image = self._test_elf2image(ELF, BIN, ["--ram-only-header"])
341        assert len(image.segments) == 2
342
343
344class TestESP8266FlashHeader(BaseTestCase):
345    def test_2mb(self):
346        ELF = "esp8266-nonossdkv20-at-v2.elf"
347        BIN = "esp8266-nonossdkv20-at-v2-0x01000.bin"
348        try:
349            self.run_elf2image(
350                "esp8266",
351                ELF,
352                version=2,
353                extra_args=["--flash_size", "2MB", "--flash_mode", "dio"],
354            )
355            with open(BIN, "rb") as f:
356                header = f.read(4)
357                print(f"header {header}")
358                self.assertEqualHex(0xEA, header[0])
359                self.assertEqualHex(0x02, header[2])
360                self.assertEqualHex(0x30, header[3])
361        finally:
362            try_delete(BIN)
363
364
365class TestESP32FlashHeader(BaseTestCase):
366    def test_16mb(self):
367        ELF = "esp32-app-template.elf"
368        BIN = "esp32-app-template.bin"
369        try:
370            self.run_elf2image(
371                "esp32",
372                ELF,
373                extra_args=[
374                    "--flash_size",
375                    "16MB",
376                    "--flash_mode",
377                    "dio",
378                    "--min-rev",
379                    "1",
380                ],
381            )
382            with open(BIN, "rb") as f:
383                header = f.read(24)
384                self.assertEqualHex(0xE9, header[0])
385                self.assertEqualHex(0x02, header[2])
386                self.assertEqualHex(0x40, header[3])
387                self.assertEqualHex(0x01, header[14])  # chip revision
388        finally:
389            try_delete(BIN)
390
391
392class TestELFSHA256(BaseTestCase):
393    ELF = "esp32-app-cust-ver-info.elf"
394    SHA_OFFS = 0xB0  # absolute offset of the SHA in the .bin file
395    BIN = "esp32-app-cust-ver-info.bin"
396
397    """
398    esp32-app-cust-ver-info.elf was built with the following application version info:
399
400    const __attribute__((section(".rodata_desc"))) esp_app_desc_t esp_app_desc = {
401        .magic_word = 0xffffffff,
402        .secure_version = 0xffffffff,
403        .reserv1 = {0xffffffff, 0xffffffff},
404        .version = "_______________________________",
405        .project_name = "-------------------------------",
406        .time = "xxxxxxxxxxxxxxx",
407        .date = "yyyyyyyyyyyyyyy",
408        .idf_ver = "zzzzzzzzzzzzzzzzzzzzzzzzzzzzzzz",
409        .app_elf_sha256 =
410        {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
411        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0},
412        .reserv2 = {0xffffffff, 0xffffffff, 0xffffffff, 0xffffffff, 0xffffffff,
413                    0xffffffff, 0xffffffff, 0xffffffff, 0xffffffff, 0xffffffff,
414                    0xffffffff, 0xffffffff, 0xffffffff, 0xffffffff, 0xffffffff,
415                    0xffffffff, 0xffffffff, 0xffffffff, 0xffffffff, 0xffffffff},
416    };
417
418    This leaves zeroes only for the fields of SHA-256 and the test will fail
419    if the placement of zeroes are tested at the wrong place.
420
421    00000000: e907 0020 780f 0840 ee00 0000 0000 0000  ... x..@........
422    00000010: 0000 0000 0000 0001 2000 403f 605a 0000  ........ .@?`Z..
423    00000020: ffff ffff ffff ffff ffff ffff ffff ffff  ................
424    00000030: 5f5f 5f5f 5f5f 5f5f 5f5f 5f5f 5f5f 5f5f  ________________
425    00000040: 5f5f 5f5f 5f5f 5f5f 5f5f 5f5f 5f5f 5f00  _______________.
426    00000050: 2d2d 2d2d 2d2d 2d2d 2d2d 2d2d 2d2d 2d2d  ----------------
427    00000060: 2d2d 2d2d 2d2d 2d2d 2d2d 2d2d 2d2d 2d00  ---------------.
428    00000070: 7878 7878 7878 7878 7878 7878 7878 7800  xxxxxxxxxxxxxxx.
429    00000080: 7979 7979 7979 7979 7979 7979 7979 7900  yyyyyyyyyyyyyyy.
430    00000090: 7a7a 7a7a 7a7a 7a7a 7a7a 7a7a 7a7a 7a7a  zzzzzzzzzzzzzzzz
431    000000a0: 7a7a 7a7a 7a7a 7a7a 7a7a 7a7a 7a7a 7a00  zzzzzzzzzzzzzzz.
432    000000b0: 0000 0000 0000 0000 0000 0000 0000 0000  ................    SHA-256 here
433    000000c0: 0000 0000 0000 0000 0000 0000 0000 0000  ................
434    000000d0: ffff ffff ffff ffff ffff ffff ffff ffff  ................
435    000000e0: ffff ffff ffff ffff ffff ffff ffff ffff  ................
436    000000f0: ffff ffff ffff ffff ffff ffff ffff ffff  ................
437    00000100: ffff ffff ffff ffff ffff ffff ffff ffff  ................
438    00000110: ffff ffff ffff ffff ffff ffff ffff ffff  ................
439    00000120: 6370 755f 7374 6172 7400 0000 1b5b 303b  cpu_start....[0;
440
441    """
442
443    def test_binary_patched(self):
444        try:
445            self.run_elf2image(
446                "esp32",
447                self.ELF,
448                extra_args=["--elf-sha256-offset", f"{self.SHA_OFFS:#x}"],
449            )
450            image = esptool.bin_image.LoadFirmwareImage("esp32", self.BIN)
451            rodata_segment = image.segments[0]
452            bin_sha256 = rodata_segment.data[
453                self.SHA_OFFS - 0x20 : self.SHA_OFFS - 0x20 + 32
454            ]  # subtract 0x20 byte header here
455
456            with open(self.ELF, "rb") as f:
457                elf_computed_sha256 = hashlib.sha256(f.read()).digest()
458
459            with open(self.BIN, "rb") as f:
460                f.seek(self.SHA_OFFS)
461                bin_sha256_raw = f.read(len(elf_computed_sha256))
462
463            assert elf_computed_sha256 == bin_sha256
464            assert elf_computed_sha256 == bin_sha256_raw
465        finally:
466            try_delete(self.BIN)
467
468    def test_no_overwrite_data(self, capsys):
469        with pytest.raises(subprocess.CalledProcessError):
470            self.run_elf2image(
471                "esp32",
472                "esp32-bootloader.elf",
473                extra_args=["--elf-sha256-offset", "0xb0"],
474            )
475        output = capsys.readouterr().out
476        assert "SHA256" in output
477        assert "zero" in output
478
479
480class TestHashAppend(BaseTestCase):
481    ELF = "esp32-bootloader.elf"
482    BIN = "esp32-bootloader.bin"
483
484    # 15th byte of the extended header after the 8-byte image header
485    HASH_APPEND_OFFSET = 15 + 8
486
487    @classmethod
488    def teardown_class(self):
489        try_delete(self.BIN)
490
491    def test_hash_append(self):
492        self.run_elf2image(
493            "esp32",
494            self.ELF,
495            extra_args=["-o", self.BIN],
496        )
497        with open(self.BIN, "rb") as f:
498            bin_with_hash = f.read()
499
500        assert bin_with_hash[self.HASH_APPEND_OFFSET] == 1
501
502        # drop the last 32 bytes (SHA256 digest)
503        expected_bin_without_hash = bytearray(bin_with_hash[:-32])
504        # disable the hash append byte in the file header
505        expected_bin_without_hash[self.HASH_APPEND_OFFSET] = 0
506
507        try_delete(self.BIN)
508        self.run_elf2image(
509            "esp32",
510            self.ELF,
511            extra_args=["--dont-append-digest", "-o", self.BIN],
512        )
513
514        with open(self.BIN, "rb") as f:
515            bin_without_hash = f.read()
516
517        assert bin_without_hash[self.HASH_APPEND_OFFSET] == 0
518        assert bytes(expected_bin_without_hash) == bin_without_hash
519