From 838e8eca9c285b8382e04f02ce8480d0510235c6 Mon Sep 17 00:00:00 2001
From: Michael Zimmermann <sigmaepsilon92@gmail.com>
Date: Mon, 9 Dec 2024 17:29:58 +0100
Subject: [PATCH] feat: add `image state-write`

---
 smpmgr/image_management.py | 29 ++++++++++++++++++++++++++++-
 1 file changed, 28 insertions(+), 1 deletion(-)

diff --git a/smpmgr/image_management.py b/smpmgr/image_management.py
index 76e49e0..7b6a917 100644
--- a/smpmgr/image_management.py
+++ b/smpmgr/image_management.py
@@ -20,7 +20,7 @@
 from smpclient import SMPClient
 from smpclient.generics import error, success
 from smpclient.mcuboot import ImageInfo
-from smpclient.requests.image_management import ImageStatesRead
+from smpclient.requests.image_management import ImageStatesRead, ImageStatesWrite
 from typing_extensions import Annotated
 
 from smpmgr.common import Options, connect_with_spinner, get_smpclient, smp_request
@@ -56,6 +56,33 @@ async def f() -> None:
     asyncio.run(f())
 
 
+@app.command()
+def state_write(
+    ctx: typer.Context,
+    hash: Annotated[str, typer.Argument(help="SHA256 hash of the image header and body.")],
+    confirm: Annotated[bool, typer.Argument(help="Confirm the image given by hash.")],
+) -> None:
+    """Request to write the state of FW images on the SMP Server."""
+
+    options = cast(Options, ctx.obj)
+    smpclient = get_smpclient(options)
+    hash_bytes = bytes.fromhex(hash)
+
+    async def f() -> None:
+        await connect_with_spinner(smpclient, options.timeout)
+
+        r = await smp_request(smpclient, options, ImageStatesWrite(hash=hash_bytes, confirm=confirm), "Waiting for image state write...")  # type: ignore # noqa
+
+        if error(r):
+            print(r)
+        elif success(r):
+            pass
+        else:
+            raise Exception("Unreachable")
+
+    asyncio.run(f())
+
+
 async def upload_with_progress_bar(
     smpclient: SMPClient, file: typer.FileBinaryRead | BufferedReader, slot: int = 0
 ) -> None: