add_adb_root: wrap zImage operations in a class

Signed-off-by: Denis 'GNUtoo' Carikli <GNUtoo@cyberdimension.org>

1 files changed, 99 insertions, 87 deletions

diff --git a/images/add_adb_root/add_adb_root.py b/images/add_adb_root/add_adb_root.py
index 3bfc185..8dc6d34 100755
--- a/images/add_adb_root/add_adb_root.py
+++ b/images/add_adb_root/add_adb_root.py
@@ -104,6 +104,103 @@ class Bootimage(object):
 
         return output_file_path
 
+class zImage(object):
+    def __init__(self, path):
+        self._path = path
+        self._ramdisk = None
+        self._ramdisk_offset = None
+        self._kernel = None
+        self._kernel_offset = None
+
+    def extract(self):
+        tmpdir = tempfile.mkdtemp()
+
+        args = ["-e", "-C", tmpdir, self._path]
+        print(["binwalk"] + args)
+        output = sh.binwalk(args)
+
+        # Example: /tmp/tmp.dkbDvuu7PL/_recovery-i9100.img.extracted
+        binwalk_dir = tmpdir + os.sep \
+            + "_" + os.path.basename(self._path) + ".extracted"
+        print("binwalk_dir: {}".format(binwalk_dir))
+
+        files = os.listdir(binwalk_dir)
+        for f in files:
+            if f.endswith(".7z"):
+                self._kernel = binwalk_dir + os.sep + f[0:-3]
+
+        # example: in
+        # /tmp/tmp.dkbDvuu7PL/_recovery-i9100.img.extracted/1E74
+        # we want 0x1E74 as int
+        self._kernel_offset = int("0x" + os.path.basename(self._kernel), 16)
+
+        print("Uncompressed Image: {} @ {}".format(
+            self._kernel, str(hex(self._kernel_offset))))
+
+        # We want the ramdisk cpio file
+        args = ["-e", "-C", binwalk_dir, self._kernel]
+        print(["binwalk"] + args)
+        output = sh.binwalk(args)
+
+        files = os.listdir(binwalk_dir)
+        ramdisk_dir = binwalk_dir \
+            + os.sep + "_" \
+            + os.path.basename(self._kernel) \
+            + ".extracted"
+
+        files = os.listdir(ramdisk_dir)
+        for f in files:
+            if f.endswith(".cpio"):
+                self._ramdisk = ramdisk_dir + os.sep + f
+
+        # example: in
+        # /tmp/tmp.*/_recovery-i9100.img.extracted/_1E74.extracted/32C9C.cpio
+        # we want 0x32C9C as int
+        self._ramdisk_offset = \
+            int("0x" + os.path.basename(self._ramdisk)[0:-5], 16)
+
+        print("Ramdisk: {} @ {}".format(self._ramdisk,
+                                        str(hex(self._ramdisk_offset))))
+
+    def recreate(self, output_file_path):
+        # ddrescue handles block size automatically
+        # and doesn't truncate the file by default
+        args = [ self._ramdisk,
+                 self._kernel,
+                 "-o",
+                 str(self._ramdisk_offset)]
+        print(["ddrescue"] + args)
+        sh.ddrescue(args)
+
+        # Issues:
+        # no size_append like in scripts/Makefile.lib
+        # Size too big
+        sh.lzma("-9", self._kernel)
+        self._kernel = self._kernel + ".lzma"
+        # TODO: Add the size of the file to the zImage
+
+        if os.path.exists(output_file_path):
+            sh.unlink(output_file_path)
+
+        sh.cp(self._path, output_file_path)
+
+        args = [self._kernel,
+                output_file_path,
+                "-o",
+                str(self._kernel_offset)]
+
+        print(["ddrescue"] + args)
+
+        sh.ddrescue(args)
+
+    def add_adb_root(self, output_file_path):
+        self.extract()
+        self._ramdisk = Ramdisk(self._ramdisk,
+                                CompressionType.none).add_adb_root()
+        self.recreate(output_file_path)
+
+        return output_file_path
+
 class Ramdisk(object):
     def __init__(self, path, compression_type):
         self._path = path
@@ -160,92 +257,6 @@ def identify_image_type(path):
 
     return ImageType.bootimage
 
-def add_adb_to_zImage(input_file, output_file):
-    tmpdir = tempfile.mkdtemp()
-
-    uncompressed_Image = None
-    ramdisk = None
-
-    args = ["-e", "-C", tmpdir, input_file]
-    print(["binwalk"] + args)
-    output = sh.binwalk(args)
-
-    # Example: /tmp/tmp.dkbDvuu7PL/_recovery-i9100.img.extracted
-    binwalk_dir = tmpdir + os.sep \
-        + "_" + os.path.basename(input_file) + ".extracted"
-    print("binwalk_dir: {}".format(binwalk_dir))
-
-    files = os.listdir(binwalk_dir)
-    for f in files:
-        if f.endswith(".7z"):
-            uncompressed_Image = binwalk_dir + os.sep + f[0:-3]
-
-    # example: in
-    # /tmp/tmp.dkbDvuu7PL/_recovery-i9100.img.extracted/1E74
-    # we want 0x1E74 as int
-    uncompressed_Image_offset = int("0x" + os.path.basename(uncompressed_Image),
-                                    16)
-
-    print("Uncompressed Image: {} @ {}".format(
-        uncompressed_Image,
-        str(hex(uncompressed_Image_offset))))
-
-    # We want the ramdisk cpio file
-    args = ["-e", "-C", binwalk_dir, uncompressed_Image]
-    print(["binwalk"] + args)
-    output = sh.binwalk(args)
-
-    files = os.listdir(binwalk_dir)
-    ramdisk_dir = binwalk_dir \
-        + os.sep + "_" \
-        + os.path.basename(uncompressed_Image) \
-        + ".extracted"
-
-    files = os.listdir(ramdisk_dir)
-    for f in files:
-        if f.endswith(".cpio"):
-            ramdisk = ramdisk_dir + os.sep + f
-
-    # example: in
-    # /tmp/tmp.[...]/_recovery-i9100.img.extracted/_1E74.extracted/32C9C.cpio
-    # we want 0x32C9C as int
-    ramdisk_offset = int("0x" + os.path.basename(ramdisk)[0:-5], 16)
-
-    print("Ramdisk: {} @ {}".format(ramdisk, str(hex(ramdisk_offset))))
-
-    ramdisk = Ramdisk(ramdisk, CompressionType.none).add_adb_root()
-
-    # ddrescue handles block size automatically
-    # and doesn't truncate the file by default
-    args = [ ramdisk,
-             uncompressed_Image,
-             "-o",
-             str(ramdisk_offset)]
-    print(["ddrescue"] + args)
-    sh.ddrescue(args)
-
-    # Issues:
-    # no size_append like in scripts/Makefile.lib
-    # Size too big
-    sh.lzma("-9", uncompressed_Image)
-    compressed_Image = uncompressed_Image + ".lzma"
-
-    # TODO: Add the size of the file to the zImage
-
-    if os.path.exists(output_file):
-        sh.unlink(output_file)
-
-    sh.cp(input_file, output_file)
-
-    args = [compressed_Image,
-            output_file ,
-            "-o",
-            str(uncompressed_Image_offset)]
-
-    print(["ddrescue"] + args)
-
-    sh.ddrescue(args)
-
 if __name__ == "__main__":
     if len(sys.argv) != 3:
         usage(sys.argv[0])
@@ -255,7 +266,8 @@ if __name__ == "__main__":
 
     image_type = identify_image_type(input_file_path)
     if image_type == ImageType.zImage:
-        add_adb_to_zImage(input_file_path, output_file_path)
+        zimage = zImage(input_file_path)
+        zimage.add_adb_root(output_file_path=output_file_path)
     elif image_type == ImageType.bootimage:
         bootimage = Bootimage(input_file_path)
         bootimage.add_adb_root(output_file_path=output_file_path)