extract-kernel: Fix indexing

The original script always tries to extract from index
0 even if header is found in positive indices. Fix that.

Also, continue to try other positions if previous run failed.

Test: builds
Bug: 139348603

Change-Id: Ia54fc709de9ae587fc64b59d94a1fa4ae669c14f

1 files changed, 20 insertions, 14 deletions

diff --git a/tools/extract_kernel.py b/tools/extract_kernel.py
index 42561cf0c6..8ca11d1cc6 100755
--- a/tools/extract_kernel.py
+++ b/tools/extract_kernel.py
@@ -100,19 +100,25 @@ def dump_configs(input_bytes):
   return o
 
 
-def try_decompress(cmd, search_bytes, input_bytes):
-  idx = input_bytes.find(search_bytes)
-  if idx < 0:
-    return None
-
-  idx = 0
+def try_decompress_bytes(cmd, input_bytes):
   sp = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE,
                         stderr=subprocess.PIPE)
-  o, _ = sp.communicate(input=input_bytes[idx:])
+  o, _ = sp.communicate(input=input_bytes)
   # ignore errors
   return o
 
 
+def try_decompress(cmd, search_bytes, input_bytes):
+  idx = 0
+  while True:
+    idx = input_bytes.find(search_bytes, idx)
+    if idx < 0:
+      raise StopIteration()
+
+    yield try_decompress_bytes(cmd, input_bytes[idx:])
+    idx += 1
+
+
 def decompress_dump(func, input_bytes):
   """
   Run func(input_bytes) first; and if that fails (returns value evaluates to
@@ -122,15 +128,15 @@ def decompress_dump(func, input_bytes):
   if o:
     return o
   for cmd, search_bytes in COMPRESSION_ALGO:
-    decompressed = try_decompress(cmd, search_bytes, input_bytes)
-    if decompressed:
-      o = func(decompressed)
-      if o:
-        return o
+    for decompressed in try_decompress(cmd, search_bytes, input_bytes):
+      if decompressed:
+        o = decompress_dump(func, decompressed)
+        if o:
+          return o
     # Force decompress the whole file even if header doesn't match
-    decompressed = try_decompress(cmd, b"", input_bytes)
+    decompressed = try_decompress_bytes(cmd, input_bytes)
     if decompressed:
-      o = func(decompressed)
+      o = decompress_dump(func, decompressed)
       if o:
         return o