add_adb_root: also move files identification code to the BootImage class

They are not used by the code handling pure zImage KERNEL
and RECOVERY images.

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

1 files changed, 46 insertions, 66 deletions

diff --git a/images/add_adb_root/add_adb_root.py b/images/add_adb_root/add_adb_root.py
index 389f2d3..7fb41ce 100755
--- a/images/add_adb_root/add_adb_root.py
+++ b/images/add_adb_root/add_adb_root.py
@@ -36,7 +36,7 @@ class CompressionType(enum.Enum):
 class Bootimage(object):
     def __init__(self, path):
         self._path = path
-        self._file_infos = check_file(self._path)
+        self._metadata = self._get_metadata()
 
         # TODO:
         # tempfile.TemporaryDirectory().name
@@ -46,22 +46,63 @@ class Bootimage(object):
         self._ramdisk = self._tmpdir + os.sep + "ramdisk.cpio.gz"
         self._config = self._tmpdir + os.sep + "bootimg.cfg"
 
+    def _get_metadata(self):
+        metadata = {}
+        metadata['release'] = None
+
+        # Try to identify the file from the checksums and retrieve all its
+        # metadata with a checksum based lookup
+        try:
+            output = sh.sha512sum(path).split(" ")
+            checksum = output[0]
+            metadata = files_checksums.checksums[checksum]
+        except:
+            pass
+
+        # The checksum based lookup failed so we need to instead scan the
+        # file and retrieve the metadata we can from it
+        output = str(sh.abootimg("-i", self._path)).split(os.linesep)
+
+        for line in output:
+            if line.startswith('  kernel:'):
+                for word in line.split(" "):
+                    if word.startswith("0x"):
+                        kernel_load_addr = int(word, 16)
+                        # TODO: Add support for 64bit targets
+                        metadata['base_address'] = kernel_load_addr & 0xffff0000
+
+            if line.startswith('* cmdline = '):
+                metadata['cmdline'] = line[len('* cmdline = '):]
+
+            if line.startswith('* id ='):
+                metadata['id'] = line[len('* id ='):].replace(" 0x", "")
+
+        assert(metadata['base_address'])
+        assert(metadata['cmdline'])
+        assert(metadata['id'])
+
+        return metadata
+
     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'],
+        sh.mkbootimg("--base", self._metadata['base_address'],
                      "--kernel", self._kernel,
                      "--ramdisk", self._ramdisk,
-                     "--cmdline={}".format(self._file_infos['cmdline']),
+                     "--cmdline={}".format(self._metadata['cmdline']),
                      "-o", output_file_path)
 
+        return output_file_path
+
     def add_adb_root(self, output_file_path):
         self.extract()
-        self._ramdisk = Ramdisk(self._ramdisk, CompressionType.gz).add_adb_root()
+        self._ramdisk = Ramdisk(self._ramdisk,
+                                CompressionType.gz).add_adb_root()
         self.recreate(output_file_path)
 
+        return output_file_path
+
 class Ramdisk(object):
     def __init__(self, path, compression_type):
         self._path = path
@@ -98,51 +139,6 @@ def usage(progname):
         progname))
     sys.exit(1)
 
-def automatically_identify_file(path):
-    file_infos = {}
-
-    # TODO: use file to get that
-    file_infos['format'] = 'boot.img'
-
-    # TODO: detect that and also allow 'boot'
-    file_infos['type'] = 'recovery'
-
-    # TODO: use file hash for that
-    file_infos['release'] = None
-
-    # TODO: is it really relevant?
-    file_infos['target'] = None
-
-    output = str(sh.abootimg("-i", path)).split(os.linesep)
-
-    for line in output:
-        if line.startswith('  kernel:'):
-            for word in line.split(" "):
-                if word.startswith("0x"):
-                    kernel_load_addr = int(word, 16)
-                    # TODO: Add support for 64bit targets
-                    file_infos['base_address'] = kernel_load_addr & 0xffff0000
-
-        if line.startswith('* cmdline = '):
-            file_infos['cmdline'] = line[len('* cmdline = '):]
-
-        if line.startswith('* id ='):
-            file_infos['id'] = line[len('* id ='):].replace(" 0x", "")
-
-    return file_infos
-
-def identify_file(path):
-    output = sh.sha512sum(path).split(" ")
-    checksum = output[0]
-
-    # TODO: handle checksums not found
-    file_infos = None
-    try:
-        file_infos = files_checksums.checksums[checksum]
-    except:
-        pass
-    return file_infos
-
 def identify_image_type(path):
     try:
         output = sh.abootimg("-i", path)
@@ -163,22 +159,6 @@ def identify_image_type(path):
 
     return ImageType.bootimage
 
-def check_file(file_path):
-    file_infos = identify_file(file_path)
-    if file_infos is None:
-        file_infos = automatically_identify_file(file_path)
-
-    if file_infos is None:
-        print("/!\ TODO: Add support for that file by adding new checksums")
-        sys.exit(1)
-
-    if file_infos['type'] != 'recovery':
-        print("/!\ The file is not a recovery")
-        print("/!\ TODO: Add support new file types")
-        sys.exit(1)
-
-    return file_infos
-
 def add_adb_to_zImage(input_file, output_file):
     tmpdir = str(sh.mktemp("-d")).replace(os.linesep, "")