Add option to copy to an Android Sparse Image

README.md

@@ -323,7 +323,7 @@ python3 -m venv .venv
 Next install the project in editable mode with development dependencies:
 
 ```bash
-pip install -m .[dev]
+pip install -e .[dev]
 ```
 
 Finally, to run tests use `unittest`:

src/bmaptool/BmapCopy.py

@@ -59,6 +59,7 @@
 import os
 import re
 import stat
+import struct
 import sys
 import hashlib
 import logging
@@ -792,6 +793,189 @@ def sync(self):
                 )
 
 
+SPARSE_HEADER_MAGIC = 0xed26ff3a
+CHUNK_TYPE_RAW = 0xCAC1
+CHUNK_TYPE_FILL = 0xCAC2
+CHUNK_TYPE_DONT_CARE = 0xCAC3
+CHUNK_TYPE_CRC32 = 0xCAC4
+
+MAJOR_VERSION = 1
+MINOR_VERSION = 0
+FILE_HEADER_SIZE = 28 # always for version 1.0
+CHUNK_HEADER_SIZE = 12 # always for version 1.0
+CHECKSUM_DONT_CARE = 0
+CHUNK_HEADER_RESERVED1 = 0
+
+
+class BmapAndroidSparseImageCopy(BmapCopy):
+    """
+    Copies the source image to an Android Sparse Image file
+
+    See: https://android.googlesource.com/platform/system/core/+/refs/heads/main/libsparse/sparse_format.h
+    """
+    def __init__(self, image, dest, bmap=None, image_size=None):
+        super().__init__(image, dest, bmap, image_size)
+
+        self._sparse_image_chunk_cnt = 0
+
+        # Use bigger batch sizes to avoid so many raw chunks in the output
+        # Sparse Image file
+        self._batch_bytes = 100 * 2**20
+        self._batch_blocks = self._batch_bytes // self.block_size
+
+
+    def _write_chunk(self, header, data=None):
+        assert(len(header) == CHUNK_HEADER_SIZE)
+
+        try:
+            self._f_dest.write(header)
+
+            if data is not None:
+                self._f_dest.write(data)
+        except IOError as err:
+            raise Error(
+                "error while writing blocks %d-%d of '%s': %s"
+                % (start, end, self._dest_path, err)
+            )
+
+        self._sparse_image_chunk_cnt += 1
+
+
+    def _write_raw_chunk(self, buf):
+        assert(len(buf) % self.block_size == 0)
+
+        block_count = len(buf) // self.block_size
+
+        chunk_header = struct.pack("<2H2I",
+            CHUNK_TYPE_RAW,
+            CHUNK_HEADER_RESERVED1,
+            block_count,
+            CHUNK_HEADER_SIZE + len(buf),
+            )
+        self._write_chunk(chunk_header, buf)
+
+
+    def _write_dont_care_chunk(self, block_count):
+        chunk_header = struct.pack("<2H2I",
+            CHUNK_TYPE_DONT_CARE,
+            CHUNK_HEADER_RESERVED1,
+            block_count,
+            CHUNK_HEADER_SIZE,
+            )
+        self._write_chunk(chunk_header)
+
+
+    def copy(self, sync=True, verify=True):
+        # Create the queue for block batches and start the reader thread, which
+        # will read the image in batches and put the results to '_batch_queue'.
+        self._batch_queue = queue.Queue(self._batch_queue_len)
+        thread.start_new_thread(self._get_data, (verify,))
+
+        blocks_written = 0
+        bytes_written = 0
+        fsync_last = 0
+
+        self._sparse_image_chunk_cnt = 0
+        self._progress_started = False
+        self._progress_index = 0
+        self._progress_time = datetime.datetime.now()
+
+        # skip past the file header. we'll seek back and write it later once we
+        # know all the information that goes in it
+        self._f_dest.seek(FILE_HEADER_SIZE)
+
+        # Read the image in '_batch_blocks' chunks and write them to the
+        # destination file
+        prev_end = -1
+        while True:
+            batch = self._batch_queue.get()
+            if batch is None:
+                # No more data, the image is written
+                break
+            elif batch[0] == "error":
+                # The reader thread encountered an error and passed us the
+                # exception.
+                exc_info = batch[1]
+                raise exc_info[1]
+
+            (start, end, buf) = batch[1:4]
+            batch_block_count = end - start + 1
+
+            assert len(buf) <= (batch_block_count) * self.block_size
+            assert len(buf) > (end - start) * self.block_size
+
+            # Synchronize the destination file if we reached the watermark
+            if self._dest_fsync_watermark:
+                if blocks_written >= fsync_last + self._dest_fsync_watermark:
+                    fsync_last = blocks_written
+                    self.sync()
+
+            blocks_since_previous_chunk = start - prev_end - 1
+
+            if blocks_since_previous_chunk < 0:
+                raise Error("Out-of-order bmap blocks are not supported")
+
+            if blocks_since_previous_chunk > 0:
+                self._write_dont_care_chunk(blocks_since_previous_chunk)
+
+            self._write_raw_chunk(buf)
+
+            self._batch_queue.task_done()
+            blocks_written += batch_block_count
+            bytes_written += len(buf)
+
+            prev_end = end
+
+            self._update_progress(blocks_written)
+
+        if prev_end < self.blocks_cnt:
+            self._write_dont_care_chunk(self.blocks_cnt - prev_end - 1)
+
+        if not self.image_size:
+            # The image size was unknown up until now, set it
+            self._set_image_size(bytes_written)
+
+        # This is just a sanity check - we should have written exactly
+        # 'mapped_cnt' blocks.
+        if blocks_written != self.mapped_cnt:
+            raise Error(
+                "wrote %u blocks from image '%s' to '%s', but should "
+                "have %u - bmap file '%s' does not belong to this "
+                "image"
+                % (
+                    blocks_written,
+                    self._image_path,
+                    self._dest_path,
+                    self.mapped_cnt,
+                    self._bmap_path,
+                )
+            )
+
+        file_header = struct.pack("<I4H4I",
+            SPARSE_HEADER_MAGIC,
+            MAJOR_VERSION,
+            MINOR_VERSION,
+            FILE_HEADER_SIZE,
+            CHUNK_HEADER_SIZE,
+            self.block_size,
+            self.blocks_cnt,
+            self._sparse_image_chunk_cnt,
+            CHECKSUM_DONT_CARE
+        )
+        assert(len(file_header) == FILE_HEADER_SIZE)
+
+        self._f_dest.seek(0)
+        self._f_dest.write(file_header)
+
+        try:
+            self._f_dest.flush()
+        except IOError as err:
+            raise Error("cannot flush '%s': %s" % (self._dest_path, err))
+
+        if sync:
+            self.sync()
+
+
 class BmapBdevCopy(BmapCopy):
     """
     This class is a specialized version of 'BmapCopy' which copies the image to

src/bmaptool/CLI.py

@@ -526,7 +526,12 @@ def copy_command(args):
             writer = BmapCopy.BmapBdevCopy(image_obj, dest_obj, bmap_obj, image_size)
         else:
             dest_str = "file '%s'" % os.path.basename(args.dest)
-            writer = BmapCopy.BmapCopy(image_obj, dest_obj, bmap_obj, image_size)
+            if args.format == "simg":
+                writer = BmapCopy.BmapAndroidSparseImageCopy(image_obj, dest_obj, bmap_obj, image_size)
+            elif args.format == "raw":
+                writer = BmapCopy.BmapCopy(image_obj, dest_obj, bmap_obj, image_size)
+            else:
+                error_out(f"Unsupported destination format: {args.format}")
     except BmapCopy.Error as err:
         error_out(err)
 
@@ -714,6 +719,10 @@ def parse_arguments():
     # The --bmap option
     text = "the block map file for the image"
     parser_copy.add_argument("--bmap", help=text)
+    parser_copy.add_argument("--format",
+                             help="format of the destination file",
+                             choices=['raw', 'simg'],
+                             default="raw")
 
     # The --nobmap option
     text = "allow copying without a bmap file"