add_adb_root: wrap ramdisk operations in a class

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

1 files changed, 40 insertions, 23 deletions

diff --git a/images/add_adb_root/add_adb_root.py b/images/add_adb_root/add_adb_root.py
index f8ac4f2..1220059 100755
--- a/images/add_adb_root/add_adb_root.py
+++ b/images/add_adb_root/add_adb_root.py
@@ -29,6 +29,41 @@ class ImageType(enum.Enum):
     zImage = 1
     bootimage = 2
 
+class CompressionType(enum.Enum):
+    none = 0
+    gz = 1
+
+class Ramdisk(object):
+    def __init__(self, path, compression_type):
+        self._path = path
+        self._compression_type = compression_type
+
+    def compress(self):
+        if self._compression_type is not CompressionType.none:
+            sh.gzip("--no-name", self._path)
+            self._path = self._path + ".gz"
+        return self._path
+
+    def uncompress(self):
+        if self._compression_type is not CompressionType.none:
+            sh.gunzip(self._path)
+            self._path = re.sub("\.gz$", "", self._path)
+        return self._path
+
+    def add_adb_root(self):
+        self.uncompress()
+
+        sh.sed("s#ro.adb.secure=1#               #",
+               "-i", self._path)
+
+        sh.sed("s#ro.secure=1#ro.secure=0#",
+               "-i", self._path)
+
+        sh.sed("s#persist.sys.usb.config=none#persist.sys.usb.config=adb #",
+               "-i", self._path)
+
+        return self.compress()
+
 def usage(progname):
     print("{} <path/to/old/recovery.img> <path/to/new/recovery.img>".format(
         progname))
@@ -115,16 +150,6 @@ def check_file(file_path):
 
     return file_infos
 
-def add_adb_to_ramdisk(ramdisk):
-    sh.sed("s#ro.adb.secure=1#               #",
-           "-i", ramdisk)
-
-    sh.sed("s#ro.secure=1#ro.secure=0#",
-           "-i", ramdisk)
-
-    sh.sed("s#persist.sys.usb.config=none#persist.sys.usb.config=adb #",
-           "-i", ramdisk)
-
 def add_adb_to_zImage(input_file, output_file):
     tmpdir = str(sh.mktemp("-d")).replace(os.linesep, "")
 
@@ -178,7 +203,7 @@ def add_adb_to_zImage(input_file, output_file):
 
     print("Ramdisk: {} @ {}".format(ramdisk, str(hex(ramdisk_offset))))
 
-    add_adb_to_ramdisk(ramdisk)
+    ramdisk = Ramdisk(ramdisk, CompressionType.none).add_adb_root()
 
     # ddrescue handles block size automatically
     # and doesn't truncate the file by default
@@ -219,28 +244,20 @@ def add_adb_to_bootimage(input_file, output_file):
     tmpdir = str(sh.mktemp("-d")).replace(os.linesep, "")
 
     kernel = tmpdir + os.sep + "kernel.img"
-    ramdisk = tmpdir + os.sep + "ramdisk.cpio"
-    compressed_ramdisk = ramdisk + ".gz"
+    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, compressed_ramdisk)
-
-    sh.gunzip(compressed_ramdisk)
-
-    # TODO:
-    # file ramdisk.cpio
-
-    add_adb_to_ramdisk(ramdisk)
+    sh.abootimg("-x", input_file, config, kernel, ramdisk)
 
-    sh.gzip("--no-name", ramdisk)
+    ramdisk = Ramdisk(ramdisk, CompressionType.gz).add_adb_root()
 
     sh.mkbootimg("--base", base_address,
                  "--kernel", kernel,
-                 "--ramdisk", compressed_ramdisk,
+                 "--ramdisk", ramdisk,
                  "--cmdline={}".format(cmdline),
                  "-o", output_file)