From c1c9bcd49f6cb159de53bbdabb238e87981424ea Mon Sep 17 00:00:00 2001
From: Vijai Kumar K <Vijaikumar_Kanagarajan@mentor.com>
Date: Sun, 2 Jan 2022 16:19:04 +0530
Subject: [PATCH] mtda: add support for xz compressed images

Based on bz2 logic.

Closes: #121

Signed-off-by: Vijai Kumar K <Vijaikumar_Kanagarajan@mentor.com>
---
 mtda/client.py         |  2 ++
 mtda/constants.py      |  1 +
 mtda/storage/writer.py | 29 +++++++++++++++++++++++++++++
 3 files changed, 32 insertions(+)

diff --git a/mtda/client.py b/mtda/client.py
index 47516771..c85a715d 100644
--- a/mtda/client.py
+++ b/mtda/client.py
@@ -245,6 +245,8 @@ def storage_write_image(self, path, callback=None):
                 compression = CONSTS.IMAGE.GZ.value
             elif path.endswith(".zst"):
                 compression = CONSTS.IMAGE.ZST.value
+            elif path.endswith(".xz"):
+                compression = CONSTS.IMAGE.XZ.value
             else:
                 compression = CONSTS.IMAGE.RAW.value
             self._impl.storage_compression(compression, self._session)
diff --git a/mtda/constants.py b/mtda/constants.py
index dcd3b663..ba40998c 100644
--- a/mtda/constants.py
+++ b/mtda/constants.py
@@ -29,6 +29,7 @@ class IMAGE(Enum):
     BZ2 = 1
     GZ = 2
     ZST = 3
+    XZ = 4
 
 
 class MDNS:
diff --git a/mtda/storage/writer.py b/mtda/storage/writer.py
index 6f752719..d1d92213 100644
--- a/mtda/storage/writer.py
+++ b/mtda/storage/writer.py
@@ -15,6 +15,7 @@
 import mtda.constants as CONSTS
 import zlib
 import zstandard as zstd
+import lzma
 
 
 class AsyncImageWriter(queue.Queue):
@@ -55,6 +56,8 @@ def compression(self, compression):
             self._write = self.write_gz
         elif compression == CONSTS.IMAGE.ZST:
             self._write = self.write_zst
+        elif compression == CONSTS.IMAGE.XZ:
+            self._write = self.write_xz
         else:
             raise ValueError("unsupported image compression!")
         self._compression = compression
@@ -213,6 +216,32 @@ def write_zst(self, data):
         self.mtda.debug(3, "storage.writer.write_zst(): %s" % str(result))
         return result
 
+    def write_xz(self, data):
+        self.mtda.debug(3, "storage.writer.write_xz()")
+
+        # Create a xz decompressor when called for the first time
+        if self._zdec is None:
+            self._zdec = lzma.LZMADecompressor()
+
+        try:
+            cont = True
+            result = None
+            while cont is True:
+                uncompressed = self._zdec.decompress(data, self._blksz)
+                result = self.storage.write(uncompressed)
+                self._written += result if result is not None else 0
+                cont = self._zdec.needs_input is False
+                data = b''
+        except EOFError:
+            result = 0
+        except OSError as e:
+            self.mtda.debug(1, "storage.writer.write_xz(): "
+                               "%s" % str(e.args[0]))
+            self._failed = True
+
+        self.mtda.debug(3, "storage.writer.write_xz(): %s" % str(result))
+        return result
+
     @property
     def writing(self):
         return self._writing