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/main.py b/scripts/imgtool/main.py
index 848fd31..03bb565 100755
--- a/scripts/imgtool/main.py
+++ b/scripts/imgtool/main.py
@@ -22,6 +22,10 @@
 import getpass
 import imgtool.keys as keys
 import sys
+import struct
+import os
+import lzma
+import hashlib
 import base64
 from imgtool import image, imgtool_version
 from imgtool.version import decode_version
@@ -29,6 +33,13 @@
 from .keys import (
     RSAUsageError, ECDSAUsageError, Ed25519UsageError, X25519UsageError)
 
+comp_default_dictsize=131072
+comp_default_pb=2
+comp_default_lc=3
+comp_default_lp=1
+comp_default_preset=9
+
+
 MIN_PYTHON_VERSION = (3, 6)
 if sys.version_info < MIN_PYTHON_VERSION:
     sys.exit("Python %s.%s or newer is required by imgtool."
@@ -300,6 +311,14 @@
         dependencies[image.DEP_VERSIONS_KEY] = versions
         return dependencies
 
+def create_lzma2_header(dictsize, pb, lc, lp):
+    header = bytearray()
+    for i in range(0, 40):
+        if dictsize <= ((2 | ((i) & 1)) << int((i) / 2 + 11)):
+            header.append(i)
+            break
+    header.append( ( pb * 5 + lp) * 9 + lc)
+    return header
 
 class BasedIntParamType(click.ParamType):
     name = 'integer'
@@ -343,6 +362,11 @@
               type=click.Choice(['128', '256']),
               help='When encrypting the image using AES, select a 128 bit or '
                    '256 bit key len.')
+@click.option('--compression', default='disabled',
+              type=click.Choice(['disabled', 'lzma2']),
+              help='Enable image compression using specified type. '
+                   'Will fall back without image compression automatically '
+                   'if the compression increases the image size.')
 @click.option('-c', '--clear', required=False, is_flag=True, default=False,
               help='Output a non-encrypted image with encryption capabilities,'
                    'so it can be installed in the primary slot, and encrypted '
@@ -414,10 +438,11 @@
                .hex extension, otherwise binary format is used''')
 def sign(key, public_key_format, align, version, pad_sig, header_size,
          pad_header, slot_size, pad, confirm, max_sectors, overwrite_only,
-         endian, encrypt_keylen, encrypt, infile, outfile, dependencies,
-         load_addr, hex_addr, erased_val, save_enctlv, security_counter,
-         boot_record, custom_tlv, rom_fixed, max_align, clear, fix_sig,
-         fix_sig_pubkey, sig_out, user_sha, vector_to_sign, non_bootable):
+         endian, encrypt_keylen, encrypt, compression, infile, outfile,
+         dependencies, load_addr, hex_addr, erased_val, save_enctlv,
+         security_counter, boot_record, custom_tlv, rom_fixed, max_align,
+         clear, fix_sig, fix_sig_pubkey, sig_out, user_sha, vector_to_sign,
+         non_bootable):
 
     if confirm:
         # Confirmed but non-padded images don't make much sense, because
@@ -431,6 +456,7 @@
                       erased_val=erased_val, save_enctlv=save_enctlv,
                       security_counter=security_counter, max_align=max_align,
                       non_bootable=non_bootable)
+    compression_tlvs = {}
     img.load(infile)
     key = load_key(key) if key else None
     enckey = load_key(encrypt) if encrypt else None
@@ -484,10 +510,49 @@
         }
 
     img.create(key, public_key_format, enckey, dependencies, boot_record,
-               custom_tlvs, int(encrypt_keylen), clear, baked_signature,
-               pub_key, vector_to_sign, user_sha)
-    img.save(outfile, hex_addr)
+               custom_tlvs, compression_tlvs, int(encrypt_keylen), clear,
+               baked_signature, pub_key, vector_to_sign, user_sha)
 
+    if compression == "lzma2":
+        compressed_img = image.Image(version=decode_version(version),
+                  header_size=header_size, pad_header=pad_header,
+                  pad=pad, confirm=confirm, align=int(align),
+                  slot_size=slot_size, max_sectors=max_sectors,
+                  overwrite_only=overwrite_only, endian=endian,
+                  load_addr=load_addr, rom_fixed=rom_fixed,
+                  erased_val=erased_val, save_enctlv=save_enctlv,
+                  security_counter=security_counter, max_align=max_align)
+        compression_filters = [
+            {"id": lzma.FILTER_LZMA2, "preset": comp_default_preset,
+                "dict_size": comp_default_dictsize, "lp": comp_default_lp,
+                "lc": comp_default_lc}
+        ]
+        compressed_data = lzma.compress(img.get_infile_data(),filters=compression_filters,
+            format=lzma.FORMAT_RAW)
+        uncompressed_size = len(img.get_infile_data())
+        compressed_size = len(compressed_data)
+        print(f"compressed image size: {compressed_size} bytes")
+        print(f"original image size: {uncompressed_size} bytes")
+        compression_tlvs["DECOMP_SIZE"] = struct.pack(
+            img.get_struct_endian() + 'L', img.image_size)
+        compression_tlvs["DECOMP_SHA"] = img.image_hash
+        compression_tlvs_size = len(compression_tlvs["DECOMP_SIZE"])
+        compression_tlvs_size += len(compression_tlvs["DECOMP_SHA"])
+        if img.get_signature():
+            compression_tlvs["DECOMP_SIGNATURE"] = img.get_signature()
+            compression_tlvs_size += len(compression_tlvs["DECOMP_SIGNATURE"])
+        if (compressed_size + compression_tlvs_size) < uncompressed_size:
+            compression_header = create_lzma2_header(
+                dictsize = comp_default_dictsize, pb = comp_default_pb,
+                lc = comp_default_lc, lp = comp_default_lp)
+            compressed_img.load_compressed(compressed_data, compression_header)
+            compressed_img.base_addr = img.base_addr
+            compressed_img.create(key, public_key_format, enckey,
+               dependencies, boot_record, custom_tlvs, compression_tlvs,
+               compression, int(encrypt_keylen), clear, baked_signature,
+               pub_key, vector_to_sign)
+            img = compressed_img
+    img.save(outfile, hex_addr)
     if sig_out is not None:
         new_signature = img.get_signature()
         save_signature(sig_out, new_signature)