scripts: imgtool: compression

Adds LZMA2 compression to imgtool.
Python lzma library is unable to compress with proper parameters while using
"ALONE" container, therefore 2 header bytes are calculated and added
to payload by imgtool.

Signed-off-by: Mateusz Michalek <mateusz.michalek@nordicsemi.no>
diff --git a/scripts/imgtool/image.py b/scripts/imgtool/image.py
index 1f9149b..690f0b7 100644
--- a/scripts/imgtool/image.py
+++ b/scripts/imgtool/image.py
@@ -20,7 +20,15 @@
 Image signing and management.
 """
 
+from . import version as versmod
+from .boot_record import create_sw_component_data
+import click
+import copy
+from enum import Enum
+import array
+from intelhex import IntelHex
 import hashlib
+import array
 import os.path
 import struct
 from enum import Enum
@@ -60,6 +68,8 @@
         'NON_BOOTABLE':          0x0000010,
         'RAM_LOAD':              0x0000020,
         'ROM_FIXED':             0x0000100,
+        'COMPRESSED_LZMA1':      0x0000200,
+        'COMPRESSED_LZMA2':      0x0000400,
 }
 
 TLV_VALUES = {
@@ -80,6 +90,9 @@
         'DEPENDENCY': 0x40,
         'SEC_CNT': 0x50,
         'BOOT_RECORD': 0x60,
+        'DECOMP_SIZE': 0x70,
+        'DECOMP_SHA': 0x71,
+        'DECOMP_SIGNATURE': 0x72,
 }
 
 TLV_SIZE = 4
@@ -238,6 +251,9 @@
         if load_addr and rom_fixed:
             raise click.UsageError("Can not set rom_fixed and load_addr at the same time")
 
+        self.image_hash = None
+        self.image_size = None
+        self.signature = None
         self.version = version or versmod.decode_version("0")
         self.header_size = header_size
         self.pad_header = pad_header
@@ -253,6 +269,7 @@
         self.rom_fixed = rom_fixed
         self.erased_val = 0xff if erased_val is None else int(erased_val, 0)
         self.payload = []
+        self.infile_data = []
         self.enckey = None
         self.save_enctlv = save_enctlv
         self.enctlv_len = 0
@@ -307,13 +324,31 @@
         try:
             if ext == INTEL_HEX_EXT:
                 ih = IntelHex(path)
-                self.payload = ih.tobinarray()
+                self.infile_data = ih.tobinarray()
+                self.payload = copy.copy(self.infile_data)
                 self.base_addr = ih.minaddr()
             else:
                 with open(path, 'rb') as f:
-                    self.payload = f.read()
+                    self.infile_data = f.read()
+                    self.payload = copy.copy(self.infile_data)
         except FileNotFoundError:
             raise click.UsageError("Input file not found")
+        self.image_size = len(self.payload)
+
+        # Add the image header if needed.
+        if self.pad_header and self.header_size > 0:
+            if self.base_addr:
+                # Adjust base_addr for new header
+                self.base_addr -= self.header_size
+            self.payload = bytes([self.erased_val] * self.header_size) + \
+                self.payload
+
+        self.check_header()
+
+    def load_compressed(self, data, compression_header):
+        """Load an image from buffer"""
+        self.payload = compression_header + data
+        self.image_size = len(self.payload)
 
         # Add the image header if needed.
         if self.pad_header and self.header_size > 0:
@@ -408,7 +443,8 @@
         return cipherkey, ciphermac, pubk
 
     def create(self, key, public_key_format, enckey, dependencies=None,
-               sw_type=None, custom_tlvs=None, encrypt_keylen=128, clear=False,
+               sw_type=None, custom_tlvs=None, compression_tlvs=None,
+               compression_type=None, encrypt_keylen=128, clear=False,
                fixed_sig=None, pub_key=None, vector_to_sign=None, user_sha='auto'):
         self.enckey = enckey
 
@@ -471,6 +507,9 @@
             dependencies_num = len(dependencies[DEP_IMAGES_KEY])
             protected_tlv_size += (dependencies_num * 16)
 
+        if compression_tlvs is not None:
+            for value in compression_tlvs.values():
+                protected_tlv_size += TLV_SIZE + len(value)
         if custom_tlvs is not None:
             for value in custom_tlvs.values():
                 protected_tlv_size += TLV_SIZE + len(value)
@@ -492,11 +531,15 @@
                 else:
                     self.payload.extend(pad)
 
+        compression_flags = 0x0
+        if compression_tlvs is not None:
+            if compression_type == "lzma2":
+                compression_flags = IMAGE_F['COMPRESSED_LZMA2']
         # This adds the header to the payload as well
         if encrypt_keylen == 256:
-            self.add_header(enckey, protected_tlv_size, 256)
+            self.add_header(enckey, protected_tlv_size, compression_flags, 256)
         else:
-            self.add_header(enckey, protected_tlv_size)
+            self.add_header(enckey, protected_tlv_size, compression_flags)
 
         prot_tlv = TLV(self.endian, TLV_PROT_INFO_MAGIC)
 
@@ -526,6 +569,9 @@
                     )
                     prot_tlv.add('DEPENDENCY', payload)
 
+            if compression_tlvs is not None:
+                for tag, value in compression_tlvs.items():
+                    prot_tlv.add(tag, value)
             if custom_tlvs is not None:
                 for tag, value in custom_tlvs.items():
                     prot_tlv.add(tag, value)
@@ -544,6 +590,7 @@
         digest = sha.digest()
         message = digest;
         tlv.add(hash_tlv, digest)
+        self.image_hash = digest
 
         if vector_to_sign == 'payload':
             # Stop amending data to the image
@@ -623,10 +670,16 @@
 
         self.check_trailer()
 
+    def get_struct_endian(self):
+        return STRUCT_ENDIAN_DICT[self.endian]
+
     def get_signature(self):
         return self.signature
 
-    def add_header(self, enckey, protected_tlv_size, aes_length=128):
+    def get_infile_data(self):
+        return self.infile_data
+
+    def add_header(self, enckey, protected_tlv_size, compression_flags, aes_length=128):
         """Install the image header."""
 
         flags = 0
@@ -664,7 +717,7 @@
                              protected_tlv_size,  # TLV Info header +
                                                   # Protected TLVs
                              len(self.payload) - self.header_size,  # ImageSz
-                             flags,
+                             flags | compression_flags,
                              self.version.major,
                              self.version.minor or 0,
                              self.version.revision or 0,