add_adb_root: wrap bootimage operations in a class

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

1 files changed, 35 insertions, 30 deletions

diff --git a/images/add_adb_root/add_adb_root.py b/images/add_adb_root/add_adb_root.py
index 1220059..389f2d3 100755
--- a/images/add_adb_root/add_adb_root.py
+++ b/images/add_adb_root/add_adb_root.py
@@ -33,6 +33,35 @@ class CompressionType(enum.Enum):
     none = 0
     gz = 1
 
+class Bootimage(object):
+    def __init__(self, path):
+        self._path = path
+        self._file_infos = check_file(self._path)
+
+        # TODO:
+        # tempfile.TemporaryDirectory().name
+        self._tmpdir = str(sh.mktemp("-d")).replace(os.linesep, "")
+
+        self._kernel = self._tmpdir + os.sep + "kernel.img"
+        self._ramdisk = self._tmpdir + os.sep + "ramdisk.cpio.gz"
+        self._config = self._tmpdir + os.sep + "bootimg.cfg"
+
+    def extract(self):
+        sh.abootimg("-x", self._path, self._config, self._kernel, self._ramdisk)
+
+    def recreate(self, output_file_path):
+        # TODO: autodetect cmdline and base_address
+        sh.mkbootimg("--base", self._file_infos['base_address'],
+                     "--kernel", self._kernel,
+                     "--ramdisk", self._ramdisk,
+                     "--cmdline={}".format(self._file_infos['cmdline']),
+                     "-o", output_file_path)
+
+    def add_adb_root(self, output_file_path):
+        self.extract()
+        self._ramdisk = Ramdisk(self._ramdisk, CompressionType.gz).add_adb_root()
+        self.recreate(output_file_path)
+
 class Ramdisk(object):
     def __init__(self, path, compression_type):
         self._path = path
@@ -236,43 +265,19 @@ def add_adb_to_zImage(input_file, output_file):
 
     sh.ddrescue(args)
 
-def add_adb_to_bootimage(input_file, output_file):
-    file_infos = check_file(input_file)
-
-    # TODO:
-    # tempfile.TemporaryDirectory().name
-    tmpdir = str(sh.mktemp("-d")).replace(os.linesep, "")
-
-    kernel = tmpdir + os.sep + "kernel.img"
-    ramdisk = tmpdir + os.sep + "ramdisk.cpio.gz"
-    config = tmpdir + os.sep + "bootimg.cfg"
-
-    # TODO: autodetect cmdline and base_address
-    base_address = file_infos['base_address']
-    cmdline = file_infos['cmdline']
-
-    sh.abootimg("-x", input_file, config, kernel, ramdisk)
-
-    ramdisk = Ramdisk(ramdisk, CompressionType.gz).add_adb_root()
-
-    sh.mkbootimg("--base", base_address,
-                 "--kernel", kernel,
-                 "--ramdisk", ramdisk,
-                 "--cmdline={}".format(cmdline),
-                 "-o", output_file)
-
 if __name__ == "__main__":
     if len(sys.argv) != 3:
         usage(sys.argv[0])
 
-    file_path = sys.argv[1]
-    target_image = sys.argv[2]
+    input_file_path = sys.argv[1]
+    output_file_path = sys.argv[2]
 
-    image_type = identify_image_type(file_path)
+    image_type = identify_image_type(input_file_path)
     if image_type == ImageType.zImage:
-        add_adb_to_zImage(file_path, target_image)
+        add_adb_to_zImage(input_file_path, output_file_path)
     elif image_type == ImageType.bootimage:
-        add_adb_to_bootimage(file_path, target_image)
+        bootimage = Bootimage(input_file_path)
+        bootimage.add_adb_root(output_file_path=output_file_path)
 
     # TODO: clenup:
     # - remove tmpdir