From 125b537aa7ca0d76c8af3bf071803de7e3056a5c Mon Sep 17 00:00:00 2001
From: puddly <32534428+puddly@users.noreply.github.com>
Date: Tue, 14 Feb 2023 12:20:19 -0500
Subject: [PATCH] OTA index file generation (#28)

* Create the `ota generate-index` command

* Move the root logger object into its own variable

* Document in README
---
 .pre-commit-config.yaml |  2 +-
 README.md               | 19 +++++++++++++
 zigpy_cli/cli.py        |  5 ++--
 zigpy_cli/ota.py        | 60 +++++++++++++++++++++++++++++++++++++++++
 4 files changed, 83 insertions(+), 3 deletions(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index f09d469..2250487 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -5,7 +5,7 @@ repos:
       - id: black
         args: ["--safe", "--quiet"]
   - repo: https://github.com/charliermarsh/ruff-pre-commit
-    rev: 'v0.0.244'
+    rev: 'v0.0.246'
     hooks:
       - id: ruff
         args: ["--fix"]
\ No newline at end of file
diff --git a/README.md b/README.md
index 6b80c47..3ee956e 100644
--- a/README.md
+++ b/README.md
@@ -155,6 +155,25 @@ $ zigpy ota dump-firmware 10047227-1.2-TRADFRI-cv-cct-unified-2.3.050.ota.ota.si
 Ember Version:    6.3.1.1
 ```
 
+## Generate OTA index files
+
+Create a JSON index for a given directory of firmwares:
+
+```console
+$ zigpy ota generate-index --ota-url-root="https://example.org/fw" path/to/firmwares/**/*.ota
+2023-02-14 12:02:03.532 ubuntu zigpy_cli.ota INFO Parsing path/to/firmwares/fw/test.ota
+2023-02-14 12:02:03.533 ubuntu zigpy_cli.ota INFO Writing path/to/firmwares/fw/test.ota
+[
+    {
+        "binary_url": "https://example.org/fw/test.ota",
+        "file_version": 1762356,
+        "image_type": 1234,
+        "manufacturer_id": 5678,
+        "changelog": "",
+        "checksum": "sha3-256:1ddaa649eb920dea9e5f002fe0d1443cc903ac0c1b26e7ad2c97b928edec2786"
+    },
+...
+```
 
 # PCAP
 ## Re-calculate the FCS on a packet capture
diff --git a/zigpy_cli/cli.py b/zigpy_cli/cli.py
index 0b3489c..694f647 100644
--- a/zigpy_cli/cli.py
+++ b/zigpy_cli/cli.py
@@ -10,6 +10,7 @@
 from zigpy_cli.const import LOG_LEVELS
 
 LOGGER = logging.getLogger(__name__)
+ROOT_LOGGER = logging.getLogger()
 
 
 def click_coroutine(cmd):
@@ -31,7 +32,7 @@ def cli(verbose):
     level_styles["trace"] = level_styles["spam"]
 
     LOGGER.setLevel(log_level)
-    logging.getLogger().setLevel(log_level)
+    ROOT_LOGGER.setLevel(log_level)
 
     coloredlogs.install(
         fmt=(
@@ -42,5 +43,5 @@ def cli(verbose):
         ),
         level=log_level,
         level_styles=level_styles,
-        logger=logging.getLogger(),
+        logger=ROOT_LOGGER,
     )
diff --git a/zigpy_cli/ota.py b/zigpy_cli/ota.py
index cad227d..3a7bcdb 100644
--- a/zigpy_cli/ota.py
+++ b/zigpy_cli/ota.py
@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import hashlib
+import json
 import logging
 import pathlib
 
@@ -65,3 +67,61 @@ def dump_firmware(input, output):
                 break
         else:
             LOGGER.warning("Image has no UPGRADE_IMAGE subelements")
+
+
+@ota.command()
+@click.pass_context
+@click.option("--ota-url-root", type=str, default=None)
+@click.option("--output", type=click.File("w"), default="-")
+@click.argument("files", nargs=-1, type=pathlib.Path)
+def generate_index(ctx, ota_url_root, output, files):
+    if ctx.parent.parent.params["verbose"] == 0:
+        cli.callback(verbose=1)
+
+    ota_metadata = []
+
+    for f in files:
+        if not f.is_file():
+            continue
+
+        LOGGER.info("Parsing %s", f)
+        contents = f.read_bytes()
+
+        try:
+            image, rest = parse_ota_image(contents)
+        except Exception as e:
+            LOGGER.error("Failed to parse: %s", e)
+            continue
+
+        if rest:
+            LOGGER.error("Image has trailing data: %r", rest)
+            continue
+
+        try:
+            validate_ota_image(image)
+        except Exception as e:
+            LOGGER.error("Image is invalid: %s", e)
+
+        if ota_url_root is not None:
+            url = f"{ota_url_root.rstrip('/')}/{f.name}"
+        else:
+            url = None
+
+        metadata = {
+            "binary_url": url,
+            "file_version": image.header.file_version,
+            "image_type": image.header.image_type,
+            "manufacturer_id": image.header.manufacturer_id,
+            "changelog": "",
+            "checksum": f"sha3-256:{hashlib.sha3_256(contents).hexdigest()}",
+        }
+
+        if image.header.hardware_versions_present:
+            metadata["min_hardware_version"] = image.header.minimum_hardware_version
+            metadata["max_hardware_version"] = image.header.maximum_hardware_version
+
+        LOGGER.info("Writing %s", f)
+        ota_metadata.append(metadata)
+
+    json.dump(ota_metadata, output, indent=4)
+    output.write("\n")