From defc9e0a91fd6975d30e555796cfac465aa99b46 Mon Sep 17 00:00:00 2001
From: Sungmin Han <sungminhan@telechips.com>
Date: Mon, 13 Jan 2025 15:21:40 +0900
Subject: [PATCH] sign_encrypt.py: fix an error in the verify command with
 '--enc_key'.

Fix a bug where the verify command requires '--enc_key' option for
encrypted TA, but an error occurs when the option is used.

Signed-off-by: Sungmin Han <sungminhan@telechips.com>
Acked-by: Jerome Forissier <jerome.forissier@linaro.org>
---
 scripts/sign_encrypt.py | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/scripts/sign_encrypt.py b/scripts/sign_encrypt.py
index 23cfe833f84..ad47a417b69 100755
--- a/scripts/sign_encrypt.py
+++ b/scripts/sign_encrypt.py
@@ -266,6 +266,7 @@ def assign_default_value(parsed, attr, func):
     arg_add_uuid(parser_verify)
     arg_add_in(parser_verify)
     arg_add_key(parser_verify)
+    arg_add_enc_key(parser_verify)
 
     parser_display = subparsers.add_parser(
         'display', prog=parser.prog + ' display',
@@ -505,9 +506,9 @@ def parse(self):
                 offs += EHDR_SIZE
                 [enc_algo, flags, nonce_len,
                  tag_len] = struct.unpack('<IIHH', self.ehdr)
-                if enc_value not in enc_tee_alg.values():
+                if enc_algo not in enc_tee_alg.values():
                     raise Exception('Unrecognized encrypt algorithm: 0x{:08x}'
-                                    .format(enc_value))
+                                    .format(enc_algo))
                 if nonce_len != 12:
                     raise Exception("Unexpected nonce len: {}"
                                     .format(nonce_len))
@@ -516,8 +517,10 @@ def parse(self):
 
                 if tag_len != 16:
                     raise Exception("Unexpected tag len: {}".format(tag_len))
-                self.tag = self.inf[-tag_len:]
-                self.ciphertext = self.inf[offs:-tag_len]
+                self.tag = self.inf[offs:offs + tag_len]
+                offs += tag_len
+
+                self.ciphertext = self.inf[offs:]
                 if len(self.ciphertext) != img_size:
                     raise Exception("Unexpected ciphertext size: ",
                                     "got {}, expected {}"
@@ -718,11 +721,11 @@ def display_ta():
             else:
                 raise Exception("Unsupported image type: {}".format(img_type))
 
-    def decrypt_ta(enc_key):
+    def decrypt_ta(self, enc_key):
         from cryptography.hazmat.primitives.ciphers.aead import AESGCM
 
         cipher = AESGCM(bytes.fromhex(enc_key))
-        self.img = cipher.decrypt(self.nonce, self.ciphertext, None)
+        self.img = cipher.decrypt(self.nonce, self.ciphertext + self.tag, None)
 
     def __get_padding(self):
         from cryptography.hazmat.primitives.asymmetric import padding
@@ -912,7 +915,7 @@ def command_verify(args):
                                             next_uuid))
             if hasattr(image, 'ciphertext'):
                 if args.enc_key is None:
-                    logger.error('--enc_key needed to decrypt TA')
+                    logger.error('--enc-key needed to decrypt TA')
                     sys.exit(1)
                 image.decrypt_ta(args.enc_key)
             image.verify_signature()