Make code more readable

extract-dtb.py

@@ -14,7 +14,6 @@
 
 You should have received a copy of the GNU General Public License
 along with extract-dtb.  If not, see <http://www.gnu.org/licenses/>.
-
 """
 
 import argparse
@@ -32,65 +31,85 @@ def dump_file(filename, content):
 
 
 def split(args):
-    pos_dtb = []
-    content = None
+    """ Reads a file and looks for DTB_HEADER occurrences (beginning of each DTB)
+        Then extract each one. If possible, use the device model as filename.
+    """
+    positions = []
 
     with open(args.filename, "rb") as fp:
         content = fp.read()
 
     dtb_next = content.find(DTB_HEADER)
     while dtb_next != -1:
-        pos_dtb.append(dtb_next)
+        positions.append(dtb_next)
         dtb_next = content.find(DTB_HEADER, dtb_next + 1)
 
-    if len(pos_dtb) == 0:
+    if len(positions) == 0:
         print("No appended dtbs found")
         return
 
     if args.extract:
         os.makedirs(args.output_dir, exist_ok=True)
-        last_pos = 0
-        for n, pos in enumerate(pos_dtb, 0):
-            dtb_filename = "dtbdump_{0:0>2}.dtb".format(n) if n != 0 else "kernel"
-            dump_file(os.path.join(args.output_dir, dtb_filename), content[last_pos:pos])
-            if n != 0:
-                dtb_filename_new = "dtbdump_{0:0>2}_{1}.dtb".format(
-                    n,
-                    next(strings(os.path.join(args.output_dir, dtb_filename))).replace(" ", "_"))
-                os.rename(os.path.join(args.output_dir, dtb_filename),
-                          os.path.join(args.output_dir, dtb_filename_new))
-                dtb_filename = dtb_filename_new
-            print("Dumped {0}, start={1} end={2}".format(dtb_filename, last_pos, pos))
-            last_pos = pos
+        begin_pos = 0
+        for n, pos in enumerate(positions, 0):
+            dtb_filename = get_dtb_filename(n)
+            filepath = os.path.join(args.output_dir, dtb_filename)
+            dump_file(filepath, content[begin_pos:pos])
+            if n > 0:
+                dtb_name = get_dtb_model(filepath)
+                if dtb_name:
+                    dtb_filename_new = get_dtb_filename(n, dtb_name)
+                    os.rename(filepath,
+                              os.path.join(args.output_dir, dtb_filename_new))
+                    dtb_filename = dtb_filename_new
+            print("Dumped {0}, start={1} end={2}"
+                  .format(dtb_filename, begin_pos, pos))
+            begin_pos = pos
 
         # Last chunk
-        dtb_filename = "dtbdump_{0:0>2}.dtb".format(n + 1)
-        dump_file(os.path.join(args.output_dir, dtb_filename), content[last_pos:])
-        dtb_filename_new = "dtbdump_{0:0>2}_{1}.dtb".format(
-            n + 1,
-            next(strings(os.path.join(args.output_dir, dtb_filename))).replace(" ", "_"))
-        os.rename(os.path.join(args.output_dir, dtb_filename),
-                  os.path.join(args.output_dir, dtb_filename_new))
-        print("Dumped {0}, start={1} end={2}".format(dtb_filename_new, last_pos, len(content)))
+        dtb_filename = get_dtb_filename(n + 1)
+        filepath = os.path.join(args.output_dir, dtb_filename)
+        dump_file(filepath, content[begin_pos:])
+        dtb_name = get_dtb_model(filepath)
+        if dtb_name:
+            dtb_filename_new = get_dtb_filename(n + 1, dtb_name)
+            os.rename(os.path.join(filepath),
+                      os.path.join(args.output_dir, dtb_filename_new))
+            dtb_filename = dtb_filename_new
+        print("Dumped {0}, start={1} end={2}"
+              .format(dtb_filename, begin_pos, len(content)))
         print("Extracted {0} appended dtbs + kernel to {1}"
-              .format(len(pos_dtb), args.output_dir))
+              .format(len(positions), args.output_dir))
     else:
-        print("Found {0} appended dtbs".format(len(pos_dtb)))
+        print("Found {0} appended dtbs".format(len(positions)))
+
+
+def get_dtb_filename(n, suffix=""):
+    if n == 0:
+        return "00_kernel"
+    n = str(n).zfill(2)
+    basename = "{0}_dtbdump".format(n)
+    if suffix != "":
+        basename += "_" + suffix.replace(" ", "_")
+    return basename + ".dtb"
 
 
-def strings(filename, min=4):
-    # with open(filename, "rb") as f:           # Python 2.x
-    with open(filename, errors="ignore") as f:  # Python 3.x
+def get_dtb_model(filename, min_length=4):
+    """ Finds the first printable string in a file with length greater
+        than min_length. Replaces spaces with underscores.
+    """
+    with open(filename, errors="ignore") as f:
         result = ""
         for c in f.read():
             if c in string.printable:
                 result += c
                 continue
-            if len(result) >= min:
-                yield result
+            if len(result) >= min_length:
+                return result.replace(" ", "_")
             result = ""
-        if len(result) >= min:  # catch result at EOF
-            yield result
+        if len(result) >= min_length:  # catch result at EOF
+            return result.replace(" ", "_")
+    return None
 
 
 if __name__ == "__main__":