From 791378cda108a145f8e390e36d617c1c5c56b6a4 Mon Sep 17 00:00:00 2001
From: cfig <yuyezhong@gmail.com>
Date: Sun, 7 Feb 2021 17:01:26 +0800
Subject: [PATCH] extract_kernel: make it py2/py3 compatible

use py3 for AOSP build, while the script itself can be
executed under py2/py3 environment directly.

TEST: extract info from Pixel 4 kernel

Change-Id: I00db6dabb1ff93f0758017a666c476816caefae7
---
 tools/Android.bp             |  4 ++--
 tools/extract_kernel.py      | 23 ++++++++++++++---------
 tools/test_extract_kernel.py | 12 ++++++------
 3 files changed, 22 insertions(+), 17 deletions(-)

diff --git a/tools/Android.bp b/tools/Android.bp
index e0f373928..357ee2ad1 100644
--- a/tools/Android.bp
+++ b/tools/Android.bp
@@ -62,10 +62,10 @@ python_binary_host {
   srcs: ["extract_kernel.py"],
   version: {
     py2: {
-      enabled: true,
+      enabled: false,
     },
     py3: {
-      enabled: false,
+      enabled: true,
     },
   },
 }
diff --git a/tools/extract_kernel.py b/tools/extract_kernel.py
index 0046b38fa..44fbcdfd3 100755
--- a/tools/extract_kernel.py
+++ b/tools/extract_kernel.py
@@ -39,12 +39,12 @@ COMPRESSION_ALGO = (
 # "Linux version " UTS_RELEASE " (" LINUX_COMPILE_BY "@"
 # LINUX_COMPILE_HOST ") (" LINUX_COMPILER ") " UTS_VERSION "\n";
 LINUX_BANNER_PREFIX = b'Linux version '
-LINUX_BANNER_REGEX = LINUX_BANNER_PREFIX + \
+LINUX_BANNER_REGEX = LINUX_BANNER_PREFIX.decode() + \
     r'(?P<release>(?P<version>[0-9]+[.][0-9]+[.][0-9]+).*) \(.*@.*\) \((?P<compiler>.*)\) .*\n'
 
 
 def get_from_release(input_bytes, start_idx, key):
-  null_idx = input_bytes.find('\x00', start_idx)
+  null_idx = input_bytes.find(b'\x00', start_idx)
   if null_idx < 0:
     return None
   try:
@@ -69,7 +69,7 @@ def dump_from_release(input_bytes, key):
 
     value = get_from_release(input_bytes, idx, key)
     if value:
-      return value
+      return value.encode()
 
     idx += len(LINUX_BANNER_PREFIX)
 
@@ -140,7 +140,7 @@ def try_decompress(cmd, search_bytes, input_bytes):
   while True:
     idx = input_bytes.find(search_bytes, idx)
     if idx < 0:
-      raise StopIteration()
+      return
 
     yield try_decompress_bytes(cmd, input_bytes[idx:])
     idx += 1
@@ -183,6 +183,11 @@ def dump_to_file(f, dump_fn, input_bytes, desc):
       return False
   return True
 
+def to_bytes_io(b):
+  """
+  Make b, which is either sys.stdout or sys.stdin, receive bytes as arguments.
+  """
+  return b.buffer if sys.version_info.major == 3 else b
 
 def main():
   parser = argparse.ArgumentParser(
@@ -194,35 +199,35 @@ def main():
                       help='Input kernel image. If not specified, use stdin',
                       metavar='FILE',
                       type=argparse.FileType('rb'),
-                      default=sys.stdin)
+                      default=to_bytes_io(sys.stdin))
   parser.add_argument('--output-configs',
                       help='If specified, write configs. Use stdout if no file '
                            'is specified.',
                       metavar='FILE',
                       nargs='?',
                       type=argparse.FileType('wb'),
-                      const=sys.stdout)
+                      const=to_bytes_io(sys.stdout))
   parser.add_argument('--output-version',
                       help='If specified, write version. Use stdout if no file '
                            'is specified.',
                       metavar='FILE',
                       nargs='?',
                       type=argparse.FileType('wb'),
-                      const=sys.stdout)
+                      const=to_bytes_io(sys.stdout))
   parser.add_argument('--output-release',
                       help='If specified, write kernel release. Use stdout if '
                            'no file is specified.',
                       metavar='FILE',
                       nargs='?',
                       type=argparse.FileType('wb'),
-                      const=sys.stdout)
+                      const=to_bytes_io(sys.stdout))
   parser.add_argument('--output-compiler',
                       help='If specified, write the compiler information. Use stdout if no file '
                            'is specified.',
                       metavar='FILE',
                       nargs='?',
                       type=argparse.FileType('wb'),
-                      const=sys.stdout)
+                      const=to_bytes_io(sys.stdout))
   parser.add_argument('--tools',
                       help='Decompression tools to use. If not specified, PATH '
                            'is searched.',
diff --git a/tools/test_extract_kernel.py b/tools/test_extract_kernel.py
index 1a1cfcbda..002e38734 100644
--- a/tools/test_extract_kernel.py
+++ b/tools/test_extract_kernel.py
@@ -15,16 +15,16 @@
 # limitations under the License.
 
 import unittest
-from extract_kernel import get_version, dump_version
+from extract_kernel import dump_version
 
 class ExtractKernelTest(unittest.TestCase):
   def test_extract_version(self):
-    self.assertEqual("4.9.100", get_version(
-        b'Linux version 4.9.100-a123 (a@a) (a) a\n\x00', 0))
-    self.assertEqual("4.9.123", get_version(
-        b'Linux version 4.9.123 (@) () \n\x00', 0))
+    self.assertEqual("4.9.100", dump_version(
+        b'Linux version 4.9.100-a123 (a@a) (a) a\n\x00'))
+    self.assertEqual("4.9.123", dump_version(
+        b'Linux version 4.9.123 (@) () \n\x00'))
 
   def test_dump_self(self):
     self.assertEqual("4.9.1", dump_version(
         b"trash\x00Linux version 4.8.8\x00trash\x00"
-        "other trash Linux version 4.9.1-g3 (2@s) (2) a\n\x00"))
+        b"other trash Linux version 4.9.1-g3 (2@s) (2) a\n\x00"))