From a84c3a8eda6404d61cb4c22c45b76df8c94c9331 Mon Sep 17 00:00:00 2001
From: Damian Krolik <damian.krolik@nordicsemi.no>
Date: Thu, 13 Jan 2022 19:28:51 +0100
Subject: [PATCH] [OTA] Add script for generating OTA image

Add script for encoding spec-compliant OTA image header
and concatenating a set of payload files.
---
 src/app/make_ota_image.py | 171 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 171 insertions(+)
 create mode 100755 src/app/make_ota_image.py

diff --git a/src/app/make_ota_image.py b/src/app/make_ota_image.py
new file mode 100755
index 00000000000000..828f36dd908891
--- /dev/null
+++ b/src/app/make_ota_image.py
@@ -0,0 +1,171 @@
+#!/usr/bin/env python3
+
+#
+#    Copyright (c) 2022 Project CHIP Authors
+#    All rights reserved.
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+#
+
+from chip.tlv import TLVWriter, uint
+import argparse
+import hashlib
+import os
+import struct
+import sys
+
+sys.path.append(os.path.join(
+    os.path.dirname(__file__), '../controller/python'))
+
+HEADER_MAGIC = 0x1BEEF11E
+
+HEADER_TAG_VENDOR_ID = 0
+HEADER_TAG_PRODUCT_ID = 1
+HEADER_TAG_VERSION = 2
+HEADER_TAG_VERSION_STRING = 3
+HEADER_TAG_PAYLOAD_SIZE = 4
+HEADER_TAG_MIN_VERSION = 5
+HEADER_TAG_MAX_VERSION = 6
+HEADER_TAG_RELEASE_NOTES_URL = 7
+HEADER_TAG_DIGEST_TYPE = 8
+HEADER_TAG_DIGEST = 9
+
+DIGEST_ALGORITHM_ID = dict(
+    sha256=1,
+    sha256_128=2,
+    sha256_120=3,
+    sha256_96=4,
+    sha256_64=5,
+    sha256_32=6,
+    sha384=7,
+    sha512=8,
+    sha3_224=9,
+    sha3_256=10,
+    sha3_384=11,
+    sha3_512=12,
+)
+DIGEST_ALL_ALGORITHMS = hashlib.algorithms_available.intersection(
+    DIGEST_ALGORITHM_ID.keys())
+
+PAYLOAD_BUFFER_SIZE = 10 * 1024
+
+
+def generate_payload_summary(args):
+    """
+    Calculate total size and hash of all concatenated input payload files
+    """
+
+    total_size = 0
+    digest = hashlib.new(args.digest_algorithm)
+
+    for path in args.input_files:
+        with open(path, 'rb') as file:
+            while chunk := file.read(PAYLOAD_BUFFER_SIZE):
+                total_size += len(chunk)
+                digest.update(chunk)
+
+    return total_size, digest.digest()
+
+
+def generate_header_tlv(args, payload_size, payload_digest):
+    """
+    Generate anonymous TLV structure with fields describing the OTA image contents
+    """
+
+    fields = {
+        HEADER_TAG_VENDOR_ID: uint(args.vendor_id),
+        HEADER_TAG_PRODUCT_ID: uint(args.product_id),
+        HEADER_TAG_VERSION: uint(args.version),
+        HEADER_TAG_VERSION_STRING: args.version_str,
+        HEADER_TAG_PAYLOAD_SIZE: uint(payload_size),
+        HEADER_TAG_DIGEST_TYPE: uint(DIGEST_ALGORITHM_ID[args.digest_algorithm]),
+        HEADER_TAG_DIGEST: payload_digest,
+    }
+
+    if args.min_version:
+        fields.update({HEADER_TAG_MIN_VERSION: uint(args.min_version)})
+
+    if args.max_version:
+        fields.update({HEADER_TAG_MAX_VERSION: uint(args.max_version)})
+
+    if args.release_notes:
+        fields.append({HEADER_TAG_RELEASE_NOTES_URL: args.release_notes})
+
+    writer = TLVWriter()
+    writer.put(None, fields)
+
+    return writer.encoding
+
+
+def generate_header(header_tlv, payload_size):
+    """
+    Generate OTA image header
+    """
+
+    fixed_fields_format = '<IQI'
+    fixed_fields = struct.pack(fixed_fields_format,
+                               HEADER_MAGIC,
+                               struct.calcsize(fixed_fields_format) +
+                               len(header_tlv) + payload_size,
+                               len(header_tlv))
+
+    return fixed_fields + header_tlv
+
+
+def write_output(args, header):
+    """
+    Write OTA image file consisting of header and concatenated payload files
+    """
+
+    with open(args.output_file, 'wb') as out_file:
+        out_file.write(header)
+
+        for path in args.input_files:
+            with open(path, 'rb') as file:
+                while chunk := file.read(PAYLOAD_BUFFER_SIZE):
+                    out_file.write(chunk)
+
+
+def main():
+    def any_base_int(s): return int(s, 0)
+
+    parser = argparse.ArgumentParser(
+        description='Matter OTA image creation tool', fromfile_prefix_chars='@')
+    parser.add_argument('-v', '--vendor-id', type=any_base_int,
+                        required=True, help='Vendor ID')
+    parser.add_argument('-p', '--product-id', type=any_base_int,
+                        required=True, help='Product ID')
+    parser.add_argument('-vn', '--version', type=any_base_int,
+                        required=True, help='Software version (numeric)')
+    parser.add_argument('-vs', '--version-str', required=True,
+                        help='Software version (string)')
+    parser.add_argument('-da', '--digest-algorithm', choices=DIGEST_ALL_ALGORITHMS,
+                        required=True, help='Digest algorithm')
+    parser.add_argument('-mi', '--min-version', type=any_base_int,
+                        help='Minimum software version that can be updated to this image')
+    parser.add_argument('-ma', '--max-version', type=any_base_int,
+                        help='Maximum software version that can be updated to this image')
+    parser.add_argument('-rn', '--release-notes', help='Release note URL')
+    parser.add_argument('input_files', nargs='+',
+                        help='Path to input image payload file')
+    parser.add_argument('output_file', help='Path to output image file')
+    args = parser.parse_args()
+
+    payload_size, payload_digest = generate_payload_summary(args)
+    header_tlv = generate_header_tlv(args, payload_size, payload_digest)
+    header = generate_header(header_tlv, payload_size)
+    write_output(args, header)
+
+
+if __name__ == "__main__":
+    main()