imgtool: update to new protected TLV format

Signed-off-by: Fabio Utzig <utzig@apache.org>
diff --git a/scripts/imgtool/image.py b/scripts/imgtool/image.py
index 34bed7b..d159e7a 100644
--- a/scripts/imgtool/image.py
+++ b/scripts/imgtool/image.py
@@ -62,6 +62,7 @@
 TLV_SIZE = 4
 TLV_INFO_SIZE = 4
 TLV_INFO_MAGIC = 0x6907
+TLV_PROT_INFO_MAGIC = 0x6908
 
 boot_magic = bytes([
     0x77, 0xc2, 0x95, 0xf3,
@@ -82,20 +83,28 @@
 
 
 class TLV():
-    def __init__(self, endian):
+    def __init__(self, endian, magic=TLV_INFO_MAGIC):
+        self.magic = magic
         self.buf = bytearray()
         self.endian = endian
 
+    def __len__(self):
+        return TLV_INFO_SIZE + len(self.buf)
+
     def add(self, kind, payload):
-        """Add a TLV record.  Kind should be a string found in TLV_VALUES above."""
+        """
+        Add a TLV record.  Kind should be a string found in TLV_VALUES above.
+        """
         e = STRUCT_ENDIAN_DICT[self.endian]
         buf = struct.pack(e + 'BBH', TLV_VALUES[kind], 0, len(payload))
         self.buf += buf
         self.buf += payload
 
     def get(self):
+        if len(self.buf) == 0:
+            return bytes()
         e = STRUCT_ENDIAN_DICT[self.endian]
-        header = struct.pack(e + 'HH', TLV_INFO_MAGIC, TLV_INFO_SIZE + len(self.buf))
+        header = struct.pack(e + 'HH', self.magic, len(self))
         return header + bytes(self.buf)
 
 
@@ -200,10 +209,15 @@
             dependencies_num = len(dependencies[DEP_IMAGES_KEY])
             protected_tlv_size = (dependencies_num * 16) + TLV_INFO_SIZE
 
+        # At this point the image is already on the payload, this adds
+        # the header to the payload as well
         self.add_header(enckey, protected_tlv_size)
 
-        tlv = TLV(self.endian)
+        prot_tlv = TLV(self.endian, TLV_PROT_INFO_MAGIC)
 
+        # Protected TLVs must be added first, because they are also included
+        # in the hash calculation
+        protected_tlv_off = None
         if protected_tlv_size != 0:
             for i in range(dependencies_num):
                 e = STRUCT_ENDIAN_DICT[self.endian]
@@ -215,23 +229,12 @@
                                 dependencies[DEP_VERSIONS_KEY][i].revision,
                                 dependencies[DEP_VERSIONS_KEY][i].build
                                 )
-                tlv.add('DEPENDENCY', payload)
-            # Full TLV size needs to be calculated in advance, because the
-            # header will be protected as well
-            tlv_header_size = 4
-            payload_digest_size = 32
-            keyhash_size = 32
-            cipherkey_size = 32
+                prot_tlv.add('DEPENDENCY', payload)
 
-            full_size = TLV_INFO_SIZE + len(tlv.buf) + tlv_header_size \
-                        + payload_digest_size
-            if key is not None:
-                full_size += tlv_header_size + keyhash_size \
-                             + tlv_header_size + key.sig_len()
-            if enckey is not None:
-                full_size += tlv_header_size + cipherkey_size
-            tlv_header = struct.pack(e + 'HH', TLV_INFO_MAGIC, full_size)
-            self.payload += tlv_header + bytes(tlv.buf)
+            protected_tlv_off = len(self.payload)
+            self.payload += prot_tlv.get()
+
+        tlv = TLV(self.endian)
 
         # Note that ecdsa wants to do the hashing itself, which means
         # we get to hash it twice.
@@ -257,6 +260,11 @@
                 sig = key.sign_digest(digest)
             tlv.add(key.sig_tlv(), sig)
 
+        # At this point the image was hashed + signed, we can remove the
+        # protected TLVs from the payload (will be re-added later)
+        if protected_tlv_off is not None:
+            self.payload = self.payload[:protected_tlv_off]
+
         if enckey is not None:
             plainkey = os.urandom(16)
             cipherkey = enckey._get_public().encrypt(
@@ -271,10 +279,11 @@
                             backend=default_backend())
             encryptor = cipher.encryptor()
             img = bytes(self.payload[self.header_size:])
-            self.payload[self.header_size:] = encryptor.update(img) + \
-                                              encryptor.finalize()
+            self.payload[self.header_size:] = \
+                encryptor.update(img) + encryptor.finalize()
 
-        self.payload += tlv.get()[protected_tlv_size:]
+        self.payload += prot_tlv.get()
+        self.payload += tlv.get()
 
     def add_header(self, enckey, protected_tlv_size):
         """Install the image header."""
@@ -300,9 +309,9 @@
                 IMAGE_MAGIC,
                 self.load_addr,
                 self.header_size,
-                protected_tlv_size,  # TLV Info header + Dependency TLVs
-                len(self.payload) - self.header_size, # ImageSz
-                flags, # Flags
+                protected_tlv_size,  # TLV Info header + Protected TLVs
+                len(self.payload) - self.header_size,  # ImageSz
+                flags,
                 self.version.major,
                 self.version.minor or 0,
                 self.version.revision or 0,