1# Tests for espsecure.py using the pytest framework
2#
3# Assumes openssl binary is in the PATH
4
5import binascii
6import io
7import os
8import os.path
9import subprocess
10import sys
11import tempfile
12from collections import namedtuple
13
14from conftest import need_to_install_package_err
15
16import pytest
17
18try:
19    import esptool
20    import espsecure
21except ImportError:
22    need_to_install_package_err()
23
24TEST_DIR = os.path.abspath(os.path.dirname(__file__))
25
26
27@pytest.mark.host_test
28class EspSecureTestCase:
29    def run_espsecure(self, args):
30        """
31        Run espsecure.py with the specified arguments
32
33        Returns output as a string if there is any,
34        raises an exception if espsecure.py fails
35        """
36        cmd = [sys.executable, "-m", "espsecure"] + args.split(" ")
37        print("\nExecuting {}...".format(" ".join(cmd)))
38
39        try:
40            output = subprocess.check_output(
41                [str(s) for s in cmd], cwd=TEST_DIR, stderr=subprocess.STDOUT
42            )
43            output = output.decode("utf-8")
44            print(output)
45            return output
46        except subprocess.CalledProcessError as e:
47            print(e.output.decode("utf-8"))
48            raise e
49
50    @classmethod
51    def setup_class(self):
52        self.cleanup_files = []  # keep a list of files _open()ed by each test case
53
54    @classmethod
55    def teardown_class(self):
56        for f in self.cleanup_files:
57            f.close()
58
59    def _open(self, image_file):
60        f = open(os.path.join(TEST_DIR, "secure_images", image_file), "rb")
61        self.cleanup_files.append(f)
62        return f
63
64
65class TestESP32SecureBootloader(EspSecureTestCase):
66    def test_digest_bootloader(self):
67        DBArgs = namedtuple(
68            "digest_bootloader_args", ["keyfile", "output", "iv", "image"]
69        )
70
71        try:
72            output_file = tempfile.NamedTemporaryFile(delete=False)
73            output_file.close()
74
75            args = DBArgs(
76                self._open("256bit_key.bin"),
77                output_file.name,
78                self._open("256bit_iv.bin"),
79                self._open("bootloader.bin"),
80            )
81            espsecure.digest_secure_bootloader(args)
82
83            with open(output_file.name, "rb") as of:
84                with self._open("bootloader_digested.bin") as ef:
85                    assert ef.read() == of.read()
86        finally:
87            os.unlink(output_file.name)
88
89    def test_digest_rsa_public_key(self):
90        DigestRSAArgs = namedtuple("digest_rsa_public_key_args", ["keyfile", "output"])
91
92        try:
93            output_file = tempfile.NamedTemporaryFile(delete=False)
94            output_file.close()
95
96            args = DigestRSAArgs(
97                self._open("rsa_secure_boot_signing_key.pem"), output_file.name
98            )
99            espsecure.digest_rsa_public_key(args)
100
101            with open(output_file.name, "rb") as of:
102                with self._open("rsa_public_key_digest.bin") as ef:
103                    assert ef.read() == of.read()
104        finally:
105            os.unlink(output_file.name)
106
107
108class TestSigning(EspSecureTestCase):
109    VerifyArgs = namedtuple(
110        "verify_signature_args", ["version", "hsm", "hsm_config", "keyfile", "datafile"]
111    )
112
113    SignArgs = namedtuple(
114        "sign_data_args",
115        [
116            "version",
117            "keyfile",
118            "output",
119            "append_signatures",
120            "hsm",
121            "hsm_config",
122            "pub_key",
123            "signature",
124            "datafile",
125        ],
126    )
127
128    ExtractKeyArgs = namedtuple(
129        "extract_public_key_args", ["version", "keyfile", "public_keyfile"]
130    )
131
132    GenerateKeyArgs = namedtuple("generate_key_args", ["version", "scheme", "keyfile"])
133
134    def test_key_generation_v1(self):
135        with tempfile.TemporaryDirectory() as keydir:
136            # keyfile cannot exist before generation -> tempfile.NamedTemporaryFile()
137            # cannot be used for keyfile
138            keyfile_name = os.path.join(keydir, "key.pem")
139            self.run_espsecure(f"generate_signing_key --version 1 {keyfile_name}")
140
141    def test_key_generation_v2(self):
142        with tempfile.TemporaryDirectory() as keydir:
143            # keyfile cannot exist before generation -> tempfile.NamedTemporaryFile()
144            # cannot be used for keyfile
145            keyfile_name = os.path.join(keydir, "key.pem")
146            self.run_espsecure(f"generate_signing_key --version 2 {keyfile_name}")
147
148    def _test_sign_v1_data(self, key_name):
149        try:
150            output_file = tempfile.NamedTemporaryFile(delete=False)
151            output_file.close()
152
153            # Note: signing bootloader is not actually needed
154            # for ESP32, it's just a handy file to sign
155            args = self.SignArgs(
156                "1",
157                [self._open(key_name)],
158                output_file.name,
159                False,
160                False,
161                None,
162                None,
163                None,
164                self._open("bootloader.bin"),
165            )
166            espsecure.sign_data(args)
167
168            with open(output_file.name, "rb") as of:
169                with self._open("bootloader_signed.bin") as ef:
170                    assert ef.read() == of.read()
171
172        finally:
173            os.unlink(output_file.name)
174
175    def test_sign_v1_data(self):
176        self._test_sign_v1_data("ecdsa_secure_boot_signing_key.pem")
177
178    def test_sign_v1_data_pkcs8(self):
179        self._test_sign_v1_data("ecdsa_secure_boot_signing_key_pkcs8.pem")
180
181    def test_sign_v1_with_pre_calculated_signature(self):
182        # Sign using pre-calculated signature + Verify
183        signing_pubkey = "ecdsa_secure_boot_signing_pubkey.pem"
184        pre_calculated_signature = "pre_calculated_bootloader_signature.bin"
185
186        try:
187            output_file = tempfile.NamedTemporaryFile(delete=False)
188            args = self.SignArgs(
189                "1",
190                None,
191                output_file.name,
192                False,
193                False,
194                None,
195                [self._open(signing_pubkey)],
196                [self._open(pre_calculated_signature)],
197                self._open("bootloader.bin"),
198            )
199            espsecure.sign_data(args)
200
201            args = self.VerifyArgs(
202                "1", False, None, self._open(signing_pubkey), output_file
203            )
204            espsecure.verify_signature(args)
205        finally:
206            output_file.close()
207            os.unlink(output_file.name)
208
209    def test_sign_v2_data(self):
210        signing_keys = [
211            "rsa_secure_boot_signing_key.pem",
212            "ecdsa192_secure_boot_signing_key.pem",
213            "ecdsa_secure_boot_signing_key.pem",
214            "ecdsa384_secure_boot_signing_key.pem",
215        ]
216        for key in signing_keys:
217            try:
218                output_file = tempfile.NamedTemporaryFile(delete=False)
219                args = self.SignArgs(
220                    "2",
221                    [self._open(key)],
222                    output_file.name,
223                    False,
224                    False,
225                    None,
226                    None,
227                    None,
228                    self._open("bootloader_unsigned_v2.bin"),
229                )
230                espsecure.sign_data(args)
231
232                args = self.VerifyArgs("2", False, None, self._open(key), output_file)
233                espsecure.verify_signature(args)
234            finally:
235                output_file.close()
236                os.unlink(output_file.name)
237
238    def test_sign_v2_multiple_keys(self):
239        # 3 keys + Verify with 3rd key
240        try:
241            output_file = tempfile.NamedTemporaryFile(delete=False)
242            args = self.SignArgs(
243                "2",
244                [
245                    self._open("rsa_secure_boot_signing_key.pem"),
246                    self._open("rsa_secure_boot_signing_key2.pem"),
247                    self._open("rsa_secure_boot_signing_key3.pem"),
248                ],
249                output_file.name,
250                False,
251                False,
252                None,
253                None,
254                None,
255                self._open("bootloader_unsigned_v2.bin"),
256            )
257            espsecure.sign_data(args)
258
259            args = self.VerifyArgs(
260                "2",
261                False,
262                None,
263                self._open("rsa_secure_boot_signing_key3.pem"),
264                output_file,
265            )
266            espsecure.verify_signature(args)
267
268            output_file.seek(0)
269            args = self.VerifyArgs(
270                "2",
271                False,
272                None,
273                self._open("rsa_secure_boot_signing_key2.pem"),
274                output_file,
275            )
276            espsecure.verify_signature(args)
277
278            output_file.seek(0)
279            args = self.VerifyArgs(
280                "2",
281                False,
282                None,
283                self._open("rsa_secure_boot_signing_key.pem"),
284                output_file,
285            )
286            espsecure.verify_signature(args)
287        finally:
288            output_file.close()
289            os.unlink(output_file.name)
290
291    def test_sign_v2_append_signatures(self):
292        # Append signatures + Verify with an appended key
293        # (bootloader_signed_v2.bin already signed with rsa_secure_boot_signing_key.pem)
294        try:
295            output_file = tempfile.NamedTemporaryFile(delete=False)
296            args = self.SignArgs(
297                "2",
298                [
299                    self._open("rsa_secure_boot_signing_key2.pem"),
300                    self._open("rsa_secure_boot_signing_key3.pem"),
301                ],
302                output_file.name,
303                True,
304                False,
305                None,
306                None,
307                None,
308                self._open("bootloader_signed_v2.bin"),
309            )
310            espsecure.sign_data(args)
311
312            args = self.VerifyArgs(
313                "2",
314                False,
315                None,
316                self._open("rsa_secure_boot_signing_key.pem"),
317                output_file,
318            )
319            espsecure.verify_signature(args)
320
321            output_file.seek(0)
322            args = self.VerifyArgs(
323                "2",
324                False,
325                None,
326                self._open("rsa_secure_boot_signing_key2.pem"),
327                output_file,
328            )
329            espsecure.verify_signature(args)
330
331            output_file.seek(0)
332            args = self.VerifyArgs(
333                "2",
334                False,
335                None,
336                self._open("rsa_secure_boot_signing_key3.pem"),
337                output_file,
338            )
339            espsecure.verify_signature(args)
340        finally:
341            output_file.close()
342            os.unlink(output_file.name)
343
344    def test_sign_v2_append_signatures_multiple_steps(self):
345        # similar to previous test, but sign in two invocations
346        try:
347            output_file1 = tempfile.NamedTemporaryFile(delete=False)
348            output_file2 = tempfile.NamedTemporaryFile(delete=False)
349            args = self.SignArgs(
350                "2",
351                [self._open("rsa_secure_boot_signing_key2.pem")],
352                output_file1.name,
353                True,
354                False,
355                None,
356                None,
357                None,
358                self._open("bootloader_signed_v2.bin"),
359            )
360            espsecure.sign_data(args)
361
362            args = self.SignArgs(
363                "2",
364                [self._open("rsa_secure_boot_signing_key3.pem")],
365                output_file2.name,
366                True,
367                False,
368                None,
369                None,
370                None,
371                output_file1,
372            )
373            espsecure.sign_data(args)
374
375            args = self.VerifyArgs(
376                "2",
377                False,
378                None,
379                self._open("rsa_secure_boot_signing_key.pem"),
380                output_file2,
381            )
382            espsecure.verify_signature(args)
383
384            output_file2.seek(0)
385            args = self.VerifyArgs(
386                "2",
387                False,
388                None,
389                self._open("rsa_secure_boot_signing_key2.pem"),
390                output_file2,
391            )
392            espsecure.verify_signature(args)
393
394            output_file2.seek(0)
395            args = self.VerifyArgs(
396                "2",
397                False,
398                None,
399                self._open("rsa_secure_boot_signing_key3.pem"),
400                output_file2,
401            )
402            espsecure.verify_signature(args)
403        finally:
404            output_file1.close()
405            os.unlink(output_file1.name)
406            output_file2.close()
407            os.unlink(output_file2.name)
408
409    def test_sign_v2_with_pre_calculated_signature(self):
410        # Sign using pre-calculated signature + Verify
411        signing_keys = [
412            "rsa_secure_boot_signing_pubkey.pem",
413            "ecdsa192_secure_boot_signing_pubkey.pem",
414            "ecdsa_secure_boot_signing_pubkey.pem",
415            "ecdsa384_secure_boot_signing_pubkey.pem",
416        ]
417        pre_calculated_signatures = [
418            "pre_calculated_bootloader_signature_rsa.bin",
419            "pre_calculated_bootloader_signature_ecdsa192.bin",
420            "pre_calculated_bootloader_signature_ecdsa256.bin",
421            "pre_calculated_bootloader_signature_ecdsa384.bin",
422        ]
423        for pub_key, signature in zip(signing_keys, pre_calculated_signatures):
424            try:
425                output_file = tempfile.NamedTemporaryFile(delete=False)
426                args = self.SignArgs(
427                    "2",
428                    None,
429                    output_file.name,
430                    False,
431                    False,
432                    None,
433                    [self._open(pub_key)],
434                    [self._open(signature)],
435                    self._open("bootloader_unsigned_v2.bin"),
436                )
437                espsecure.sign_data(args)
438
439                args = self.VerifyArgs(
440                    "2", False, None, self._open(pub_key), output_file
441                )
442                espsecure.verify_signature(args)
443            finally:
444                output_file.close()
445                os.unlink(output_file.name)
446
447    def test_sign_v2_with_multiple_pre_calculated_signatures(self):
448        # Sign using multiple pre-calculated signatures + Verify
449        signing_pubkeys = [
450            "rsa_secure_boot_signing_pubkey.pem",
451            "rsa_secure_boot_signing_pubkey.pem",
452            "rsa_secure_boot_signing_pubkey.pem",
453        ]
454        pre_calculated_signatures = [
455            "pre_calculated_bootloader_signature_rsa.bin",
456            "pre_calculated_bootloader_signature_rsa.bin",
457            "pre_calculated_bootloader_signature_rsa.bin",
458        ]
459        try:
460            output_file = tempfile.NamedTemporaryFile(delete=False)
461            args = self.SignArgs(
462                "2",
463                None,
464                output_file.name,
465                False,
466                False,
467                None,
468                [self._open(pub_key) for pub_key in signing_pubkeys],
469                [self._open(signature) for signature in pre_calculated_signatures],
470                self._open("bootloader_unsigned_v2.bin"),
471            )
472            espsecure.sign_data(args)
473
474            args = self.VerifyArgs(
475                "2", False, None, self._open(signing_pubkeys[0]), output_file
476            )
477            espsecure.verify_signature(args)
478        finally:
479            output_file.close()
480            os.unlink(output_file.name)
481
482    def test_verify_signature_signing_key(self):
483        # correct key v1
484        args = self.VerifyArgs(
485            "1",
486            False,
487            None,
488            self._open("ecdsa_secure_boot_signing_key.pem"),
489            self._open("bootloader_signed.bin"),
490        )
491        espsecure.verify_signature(args)
492
493        # correct key v2
494        args = self.VerifyArgs(
495            "2",
496            False,
497            None,
498            self._open("rsa_secure_boot_signing_key.pem"),
499            self._open("bootloader_signed_v2.bin"),
500        )
501        espsecure.verify_signature(args)
502
503        # correct key v2 (ecdsa384)
504        args = self.VerifyArgs(
505            "2",
506            False,
507            None,
508            self._open("ecdsa384_secure_boot_signing_key.pem"),
509            self._open("bootloader_signed_v2_ecdsa384.bin"),
510        )
511        espsecure.verify_signature(args)
512
513        # correct key v2 (ecdsa256)
514        args = self.VerifyArgs(
515            "2",
516            False,
517            None,
518            self._open("ecdsa_secure_boot_signing_key.pem"),
519            self._open("bootloader_signed_v2_ecdsa256.bin"),
520        )
521        espsecure.verify_signature(args)
522
523        # correct key v2 (ecdsa192)
524        args = self.VerifyArgs(
525            "2",
526            False,
527            None,
528            self._open("ecdsa192_secure_boot_signing_key.pem"),
529            self._open("bootloader_signed_v2_ecdsa192.bin"),
530        )
531        espsecure.verify_signature(args)
532
533        # wrong key v1
534        args = self.VerifyArgs(
535            "1",
536            False,
537            None,
538            self._open("ecdsa_secure_boot_signing_key2.pem"),
539            self._open("bootloader_signed.bin"),
540        )
541        with pytest.raises(esptool.FatalError) as cm:
542            espsecure.verify_signature(args)
543        assert "Signature is not valid" in str(cm.value)
544
545        # wrong key v2
546        args = self.VerifyArgs(
547            "2",
548            False,
549            None,
550            self._open("rsa_secure_boot_signing_key2.pem"),
551            self._open("bootloader_signed_v2.bin"),
552        )
553        with pytest.raises(esptool.FatalError) as cm:
554            espsecure.verify_signature(args)
555        assert "Signature could not be verified with the provided key." in str(cm.value)
556
557        # right key, wrong scheme (ecdsa256, v2)
558        args = self.VerifyArgs(
559            "2",
560            False,
561            None,
562            self._open("ecdsa_secure_boot_signing_key.pem"),
563            self._open("bootloader_signed.bin"),
564        )
565        with pytest.raises(esptool.FatalError) as cm:
566            espsecure.verify_signature(args)
567        assert "Invalid datafile" in str(cm.value)
568
569        # wrong key v2 (ecdsa384)
570        args = self.VerifyArgs(
571            "2",
572            False,
573            None,
574            self._open("ecdsa384_secure_boot_signing_key2.pem"),
575            self._open("bootloader_signed_v2_ecdsa384.bin"),
576        )
577        with pytest.raises(esptool.FatalError) as cm:
578            espsecure.verify_signature(args)
579        assert "Signature could not be verified with the provided key." in str(cm.value)
580
581        # wrong key v2 (ecdsa256)
582        args = self.VerifyArgs(
583            "2",
584            False,
585            None,
586            self._open("ecdsa_secure_boot_signing_key2.pem"),
587            self._open("bootloader_signed_v2_ecdsa256.bin"),
588        )
589        with pytest.raises(esptool.FatalError) as cm:
590            espsecure.verify_signature(args)
591        assert "Signature could not be verified with the provided key." in str(cm.value)
592
593        # wrong key v2 (ecdsa192)
594        args = self.VerifyArgs(
595            "2",
596            False,
597            None,
598            self._open("ecdsa192_secure_boot_signing_key2.pem"),
599            self._open("bootloader_signed_v2_ecdsa192.bin"),
600        )
601        with pytest.raises(esptool.FatalError) as cm:
602            espsecure.verify_signature(args)
603        assert "Signature could not be verified with the provided key." in str(cm.value)
604
605        # multi-signed wrong key v2
606        args = self.VerifyArgs(
607            "2",
608            False,
609            None,
610            self._open("rsa_secure_boot_signing_key4.pem"),
611            self._open("bootloader_multi_signed_v2.bin"),
612        )
613        with pytest.raises(esptool.FatalError) as cm:
614            espsecure.verify_signature(args)
615        assert "Signature could not be verified with the provided key." in str(cm.value)
616
617    def test_verify_signature_public_key(self):
618        # correct key v1
619        args = self.VerifyArgs(
620            "1",
621            False,
622            None,
623            self._open("ecdsa_secure_boot_signing_pubkey.pem"),
624            self._open("bootloader_signed.bin"),
625        )
626        espsecure.verify_signature(args)
627
628        # correct key v2
629        args = self.VerifyArgs(
630            "2",
631            False,
632            None,
633            self._open("rsa_secure_boot_signing_pubkey.pem"),
634            self._open("bootloader_signed_v2.bin"),
635        )
636        espsecure.verify_signature(args)
637
638        # correct key v2 (ecdsa384)
639        args = self.VerifyArgs(
640            "2",
641            False,
642            None,
643            self._open("ecdsa384_secure_boot_signing_pubkey.pem"),
644            self._open("bootloader_signed_v2_ecdsa384.bin"),
645        )
646        espsecure.verify_signature(args)
647
648        # correct key v2 (ecdsa256)
649        args = self.VerifyArgs(
650            "2",
651            False,
652            None,
653            self._open("ecdsa_secure_boot_signing_pubkey.pem"),
654            self._open("bootloader_signed_v2_ecdsa256.bin"),
655        )
656        espsecure.verify_signature(args)
657
658        # correct key v2 (ecdsa192)
659        args = self.VerifyArgs(
660            "2",
661            False,
662            None,
663            self._open("ecdsa192_secure_boot_signing_pubkey.pem"),
664            self._open("bootloader_signed_v2_ecdsa192.bin"),
665        )
666        espsecure.verify_signature(args)
667
668        # wrong key v1
669        args = self.VerifyArgs(
670            "1",
671            False,
672            None,
673            self._open("ecdsa_secure_boot_signing_pubkey2.pem"),
674            self._open("bootloader_signed.bin"),
675        )
676        with pytest.raises(esptool.FatalError) as cm:
677            espsecure.verify_signature(args)
678        assert "Signature is not valid" in str(cm.value)
679
680        # wrong key v2
681        args = self.VerifyArgs(
682            "2",
683            False,
684            None,
685            self._open("rsa_secure_boot_signing_pubkey2.pem"),
686            self._open("bootloader_signed_v2.bin"),
687        )
688        with pytest.raises(esptool.FatalError) as cm:
689            espsecure.verify_signature(args)
690        assert "Signature could not be verified with the provided key." in str(cm.value)
691
692        # wrong key v2 (ecdsa384)
693        args = self.VerifyArgs(
694            "2",
695            False,
696            None,
697            self._open("ecdsa384_secure_boot_signing_pubkey2.pem"),
698            self._open("bootloader_signed_v2_ecdsa384.bin"),
699        )
700        with pytest.raises(esptool.FatalError) as cm:
701            espsecure.verify_signature(args)
702        assert "Signature could not be verified with the provided key." in str(cm.value)
703
704        # wrong key v2 (ecdsa256)
705        args = self.VerifyArgs(
706            "2",
707            False,
708            None,
709            self._open("ecdsa_secure_boot_signing_pubkey2.pem"),
710            self._open("bootloader_signed_v2_ecdsa256.bin"),
711        )
712        with pytest.raises(esptool.FatalError) as cm:
713            espsecure.verify_signature(args)
714        assert "Signature could not be verified with the provided key." in str(cm.value)
715
716        # wrong key v2 (ecdsa192)
717        args = self.VerifyArgs(
718            "2",
719            False,
720            None,
721            self._open("ecdsa192_secure_boot_signing_pubkey2.pem"),
722            self._open("bootloader_signed_v2_ecdsa192.bin"),
723        )
724        with pytest.raises(esptool.FatalError) as cm:
725            espsecure.verify_signature(args)
726        assert "Signature could not be verified with the provided key." in str(cm.value)
727
728        # multi-signed wrong key v2
729        args = self.VerifyArgs(
730            "2",
731            False,
732            None,
733            self._open("rsa_secure_boot_signing_pubkey4.pem"),
734            self._open("bootloader_multi_signed_v2.bin"),
735        )
736        with pytest.raises(esptool.FatalError) as cm:
737            espsecure.verify_signature(args)
738        assert "Signature could not be verified with the provided key." in str(cm.value)
739
740    def test_extract_binary_public_key(self):
741        with tempfile.NamedTemporaryFile() as pub_keyfile, tempfile.NamedTemporaryFile() as pub_keyfile2:  # noqa E501
742            args = self.ExtractKeyArgs(
743                "1", self._open("ecdsa_secure_boot_signing_key.pem"), pub_keyfile
744            )
745            espsecure.extract_public_key(args)
746
747            args = self.ExtractKeyArgs(
748                "1", self._open("ecdsa_secure_boot_signing_key2.pem"), pub_keyfile2
749            )
750            espsecure.extract_public_key(args)
751
752            pub_keyfile.seek(0)
753            pub_keyfile2.seek(0)
754
755            # use correct extracted public key to verify
756            args = self.VerifyArgs(
757                "1", False, None, pub_keyfile, self._open("bootloader_signed.bin")
758            )
759            espsecure.verify_signature(args)
760
761            # use wrong extracted public key to try and verify
762            args = self.VerifyArgs(
763                "1", False, None, pub_keyfile2, self._open("bootloader_signed.bin")
764            )
765            with pytest.raises(esptool.FatalError) as cm:
766                espsecure.verify_signature(args)
767            assert "Signature is not valid" in str(cm.value)
768
769    def test_generate_and_extract_key_v2(self):
770        with tempfile.TemporaryDirectory() as keydir:
771            # keyfile cannot exist before generation -> tempfile.NamedTemporaryFile()
772            # cannot be used for keyfile
773            keyfile_name = os.path.join(keydir, "key.pem")
774
775            # We need to manually delete the keyfile as we are iterating over
776            # different schemes with the same keyfile so instead of using addCleanup,
777            # we remove it using os.remove at the end of each pass
778            for scheme in ["rsa3072", "ecdsa192", "ecdsa256", "ecdsa384"]:
779                args = self.GenerateKeyArgs("2", scheme, keyfile_name)
780                espsecure.generate_signing_key(args)
781
782                with tempfile.NamedTemporaryFile() as pub_keyfile, open(
783                    keyfile_name, "rb"
784                ) as keyfile:
785                    args = self.ExtractKeyArgs("2", keyfile, pub_keyfile)
786                    espsecure.extract_public_key(args)
787                os.remove(keyfile_name)
788
789
790class TestFlashEncryption(EspSecureTestCase):
791    EncryptArgs = namedtuple(
792        "encrypt_flash_data_args",
793        [
794            "keyfile",
795            "output",
796            "address",
797            "flash_crypt_conf",
798            "aes_xts",
799            "plaintext_file",
800        ],
801    )
802
803    DecryptArgs = namedtuple(
804        "decrypt_flash_data_args",
805        [
806            "keyfile",
807            "output",
808            "address",
809            "flash_crypt_conf",
810            "aes_xts",
811            "encrypted_file",
812        ],
813    )
814
815    def _test_encrypt_decrypt(
816        self,
817        input_plaintext,
818        expected_ciphertext,
819        key_path,
820        offset,
821        flash_crypt_conf=0xF,
822        aes_xts=None,
823    ):
824        original_plaintext = self._open(input_plaintext)
825        keyfile = self._open(key_path)
826        ciphertext = io.BytesIO()
827
828        args = self.EncryptArgs(
829            keyfile, ciphertext, offset, flash_crypt_conf, aes_xts, original_plaintext
830        )
831        espsecure.encrypt_flash_data(args)
832
833        original_plaintext.seek(0)
834        assert original_plaintext.read() != ciphertext.getvalue()
835        with self._open(expected_ciphertext) as f:
836            assert f.read() == ciphertext.getvalue()
837
838        ciphertext.seek(0)
839        keyfile.seek(0)
840        plaintext = io.BytesIO()
841        args = self.DecryptArgs(
842            keyfile, plaintext, offset, flash_crypt_conf, aes_xts, ciphertext
843        )
844        espsecure.decrypt_flash_data(args)
845
846        original_plaintext.seek(0)
847        assert original_plaintext.read() == plaintext.getvalue()
848
849
850class TestESP32FlashEncryption(TestFlashEncryption):
851    def test_encrypt_decrypt_bootloader(self):
852        self._test_encrypt_decrypt(
853            "bootloader.bin", "bootloader-encrypted.bin", "256bit_key.bin", 0x1000, 0xF
854        )
855
856    def test_encrypt_decrypt_app(self):
857        self._test_encrypt_decrypt(
858            "hello-world-signed.bin",
859            "hello-world-signed-encrypted.bin",
860            "ef-flashencryption-key.bin",
861            0x20000,
862            0xF,
863        )
864
865    def test_encrypt_decrypt_non_default_conf(self):
866        """Try some non-default (non-recommended) flash_crypt_conf settings"""
867        for conf in [0x0, 0x3, 0x9, 0xC]:
868            self._test_encrypt_decrypt(
869                "bootloader.bin",
870                f"bootloader-encrypted-conf{conf:x}.bin",
871                "256bit_key.bin",
872                0x1000,
873                conf,
874            )
875
876
877class TestAesXtsFlashEncryption(TestFlashEncryption):
878    def test_encrypt_decrypt_bootloader(self):
879        self._test_encrypt_decrypt(
880            "bootloader.bin",
881            "bootloader-encrypted-aes-xts.bin",
882            "256bit_key.bin",
883            0x1000,
884            aes_xts=True,
885        )
886
887    def test_encrypt_decrypt_app(self):
888        self._test_encrypt_decrypt(
889            "hello-world-signed.bin",
890            "hello-world-signed-encrypted-aes-xts.bin",
891            "ef-flashencryption-key.bin",
892            0x20000,
893            aes_xts=True,
894        )
895
896    def test_encrypt_decrypt_app_512_bit_key(self):
897        self._test_encrypt_decrypt(
898            "hello-world-signed.bin",
899            "hello-world-signed-encrypted-aes-xts-256.bin",
900            "512bit_key.bin",
901            0x10000,
902            aes_xts=True,
903        )
904
905    def test_padding(self):
906        # Random 2048 bits hex string
907        plaintext = binascii.unhexlify(
908            "c33b7c49f12a969a9bb45af5f660b73f"
909            "3b372685012da570df1cf99d1a82eabb"
910            "fdf6aa16b9675bd8a2f95e871513e175"
911            "3bc89f57986ecfb2707a3d3b59a46968"
912            "5e6609d2e9c21d4b2310571175e6e3de"
913            "2656ee22243f557b925ef39ff782ab56"
914            "f821e6859ee852000daae7c03a7c77ce"
915            "58744f15fbdf0ad4ae6e964aedd6316a"
916            "cf0e36935eef895cd14a60fe682fb971"
917            "eb239eae38b770bdf969017c9decfd91"
918            "b7c60329fb0c896684f0e7415f99dec1"
919            "da0572fac360a3e6d7219973a7de07e5"
920            "33b5abfdf5917ed5bfe54d660a6f5047"
921            "32fdb8d07259bfcdc67da87293857c11"
922            "427b2bae5f00da4a4b2b00b588ff5109"
923            "4c41f07f02f680f8826841b43da3f25b"
924        )
925
926        plaintext_file = io.BytesIO(plaintext)
927        ciphertext_full_block = io.BytesIO()
928
929        keyfile = self._open("256bit_key.bin")
930        address = 0x1000
931
932        encrypt_args_padded = self.EncryptArgs(
933            keyfile, ciphertext_full_block, address, None, "aes_xts", plaintext_file
934        )
935
936        espsecure.encrypt_flash_data(encrypt_args_padded)
937
938        # Test with different number of bytes per encryption call
939        # Final ciphertext should still be the same if padding is done correctly
940        bytes_per_encrypt = [16, 32, 64, 128]
941
942        for b in bytes_per_encrypt:
943            ciphertext = io.BytesIO()
944            num_enc_calls = len(plaintext) // b
945
946            for i in range(0, num_enc_calls):
947                keyfile.seek(0)
948                offset = b * i
949
950                # encrypt the whole plaintext a substring of b bytes at a time
951                plaintext_sub = io.BytesIO(plaintext[offset : offset + b])
952
953                encrypt_args = self.EncryptArgs(
954                    keyfile,
955                    ciphertext,
956                    address + offset,
957                    None,
958                    "aes_xts",
959                    plaintext_sub,
960                )
961
962                espsecure.encrypt_flash_data(encrypt_args)
963
964            assert ciphertext_full_block.getvalue() == ciphertext.getvalue()
965
966
967class TestDigest(EspSecureTestCase):
968    def test_digest_private_key(self):
969        with tempfile.NamedTemporaryFile() as f:
970            outfile_name = f.name
971
972        self.run_espsecure(
973            "digest_private_key "
974            "--keyfile secure_images/ecdsa_secure_boot_signing_key.pem "
975            f"{outfile_name}"
976        )
977
978        with open(outfile_name, "rb") as f:
979            assert f.read() == binascii.unhexlify(
980                "7b7b53708fc89d5e0b2df2571fb8f9d778f61a422ff1101a22159c4b34aad0aa"
981            )
982
983    def test_digest_private_key_with_invalid_output(self, capsys):
984        fname = "secure_images/ecdsa_secure_boot_signing_key.pem"
985
986        with pytest.raises(subprocess.CalledProcessError):
987            self.run_espsecure(f"digest_private_key --keyfile {fname} {fname}")
988        output = capsys.readouterr().out
989        assert "should not be the same!" in output
990