1#!/usr/bin/env python
2#
3# Copyright (C) 2018 The Android Open Source Project
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#      http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16
17"""
18A tool to extract kernel information from a kernel image.
19"""
20
21import argparse
22import subprocess
23import sys
24import re
25
26CONFIG_PREFIX = b'IKCFG_ST'
27GZIP_HEADER = b'\037\213\010'
28COMPRESSION_ALGO = (
29    (["gzip", "-d"], GZIP_HEADER),
30    (["xz", "-d"], b'\3757zXZ\000'),
31    (["bzip2", "-d"], b'BZh'),
32    (["lz4", "-d", "-l"], b'\002\041\114\030'),
33
34    # These are not supported in the build system yet.
35    # (["unlzma"], b'\135\0\0\0'),
36    # (["lzop", "-d"], b'\211\114\132'),
37)
38
39# "Linux version " UTS_RELEASE " (" LINUX_COMPILE_BY "@"
40# LINUX_COMPILE_HOST ") (" LINUX_COMPILER ") " UTS_VERSION "\n";
41LINUX_BANNER_PREFIX = b'Linux version '
42LINUX_BANNER_REGEX = LINUX_BANNER_PREFIX.decode() + \
43    r'(?P<release>(?P<version>[0-9]+[.][0-9]+[.][0-9]+).*) \(.*@.*\) \((?P<compiler>.*)\) .*\n'
44
45
46def get_from_release(input_bytes, start_idx, key):
47  null_idx = input_bytes.find(b'\x00', start_idx)
48  if null_idx < 0:
49    return None
50  try:
51    linux_banner = input_bytes[start_idx:null_idx].decode()
52  except UnicodeDecodeError:
53    return None
54  mo = re.match(LINUX_BANNER_REGEX, linux_banner)
55  if mo:
56    return mo.group(key)
57  return None
58
59
60def dump_from_release(input_bytes, key):
61  """
62  Helper of dump_version and dump_release
63  """
64  idx = 0
65  while True:
66    idx = input_bytes.find(LINUX_BANNER_PREFIX, idx)
67    if idx < 0:
68      return None
69
70    value = get_from_release(input_bytes, idx, key)
71    if value:
72      return value.encode()
73
74    idx += len(LINUX_BANNER_PREFIX)
75
76
77def dump_version(input_bytes):
78  """
79  Dump kernel version, w.x.y, from input_bytes. Search for the string
80  "Linux version " and do pattern matching after it. See LINUX_BANNER_REGEX.
81  """
82  return dump_from_release(input_bytes, "version")
83
84
85def dump_compiler(input_bytes):
86  """
87  Dump kernel version, w.x.y, from input_bytes. Search for the string
88  "Linux version " and do pattern matching after it. See LINUX_BANNER_REGEX.
89  """
90  return dump_from_release(input_bytes, "compiler")
91
92
93def dump_release(input_bytes):
94  """
95  Dump kernel release, w.x.y-..., from input_bytes. Search for the string
96  "Linux version " and do pattern matching after it. See LINUX_BANNER_REGEX.
97  """
98  return dump_from_release(input_bytes, "release")
99
100
101def dump_configs(input_bytes):
102  """
103  Dump kernel configuration from input_bytes. This can be done when
104  CONFIG_IKCONFIG is enabled, which is a requirement on Treble devices.
105
106  The kernel configuration is archived in GZip format right after the magic
107  string 'IKCFG_ST' in the built kernel.
108  """
109
110  # Search for magic string + GZip header
111  idx = input_bytes.find(CONFIG_PREFIX + GZIP_HEADER)
112  if idx < 0:
113    return None
114
115  # Seek to the start of the archive
116  idx += len(CONFIG_PREFIX)
117
118  sp = subprocess.Popen(["gzip", "-d", "-c"], stdin=subprocess.PIPE,
119                        stdout=subprocess.PIPE, stderr=subprocess.PIPE)
120  o, _ = sp.communicate(input=input_bytes[idx:])
121  if sp.returncode == 1: # error
122    return None
123
124  # success or trailing garbage warning
125  assert sp.returncode in (0, 2), sp.returncode
126
127  return o
128
129
130def try_decompress_bytes(cmd, input_bytes):
131  sp = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE,
132                        stderr=subprocess.PIPE)
133  o, _ = sp.communicate(input=input_bytes)
134  # ignore errors
135  return o
136
137
138def try_decompress(cmd, search_bytes, input_bytes):
139  idx = 0
140  while True:
141    idx = input_bytes.find(search_bytes, idx)
142    if idx < 0:
143      return
144
145    yield try_decompress_bytes(cmd, input_bytes[idx:])
146    idx += 1
147
148
149def decompress_dump(func, input_bytes):
150  """
151  Run func(input_bytes) first; and if that fails (returns value evaluates to
152  False), then try different decompression algorithm before running func.
153  """
154  o = func(input_bytes)
155  if o:
156    return o
157  for cmd, search_bytes in COMPRESSION_ALGO:
158    for decompressed in try_decompress(cmd, search_bytes, input_bytes):
159      if decompressed:
160        o = decompress_dump(func, decompressed)
161        if o:
162          return o
163    # Force decompress the whole file even if header doesn't match
164    decompressed = try_decompress_bytes(cmd, input_bytes)
165    if decompressed:
166      o = decompress_dump(func, decompressed)
167      if o:
168        return o
169
170
171def dump_to_file(f, dump_fn, input_bytes, desc):
172  """
173  Call decompress_dump(dump_fn, input_bytes) and write to f. If it fails, return
174  False; otherwise return True.
175  """
176  if f is not None:
177    o = decompress_dump(dump_fn, input_bytes)
178    if o:
179      f.write(o)
180    else:
181      sys.stderr.write(
182          "Cannot extract kernel {}".format(desc))
183      return False
184  return True
185
186def to_bytes_io(b):
187  """
188  Make b, which is either sys.stdout or sys.stdin, receive bytes as arguments.
189  """
190  return b.buffer if sys.version_info.major == 3 else b
191
192def main():
193  parser = argparse.ArgumentParser(
194      formatter_class=argparse.RawTextHelpFormatter,
195      description=__doc__ +
196      "\nThese algorithms are tried when decompressing the image:\n    " +
197      " ".join(tup[0][0] for tup in COMPRESSION_ALGO))
198  parser.add_argument('--input',
199                      help='Input kernel image. If not specified, use stdin',
200                      metavar='FILE',
201                      type=argparse.FileType('rb'),
202                      default=to_bytes_io(sys.stdin))
203  parser.add_argument('--output-configs',
204                      help='If specified, write configs. Use stdout if no file '
205                           'is specified.',
206                      metavar='FILE',
207                      nargs='?',
208                      type=argparse.FileType('wb'),
209                      const=to_bytes_io(sys.stdout))
210  parser.add_argument('--output-version',
211                      help='If specified, write version. Use stdout if no file '
212                           'is specified.',
213                      metavar='FILE',
214                      nargs='?',
215                      type=argparse.FileType('wb'),
216                      const=to_bytes_io(sys.stdout))
217  parser.add_argument('--output-release',
218                      help='If specified, write kernel release. Use stdout if '
219                           'no file is specified.',
220                      metavar='FILE',
221                      nargs='?',
222                      type=argparse.FileType('wb'),
223                      const=to_bytes_io(sys.stdout))
224  parser.add_argument('--output-compiler',
225                      help='If specified, write the compiler information. Use stdout if no file '
226                           'is specified.',
227                      metavar='FILE',
228                      nargs='?',
229                      type=argparse.FileType('wb'),
230                      const=to_bytes_io(sys.stdout))
231  parser.add_argument('--tools',
232                      help='Decompression tools to use. If not specified, PATH '
233                           'is searched.',
234                      metavar='ALGORITHM:EXECUTABLE',
235                      nargs='*')
236  args = parser.parse_args()
237
238  tools = {pair[0]: pair[1]
239           for pair in (token.split(':') for token in args.tools or [])}
240  for cmd, _ in COMPRESSION_ALGO:
241    if cmd[0] in tools:
242      cmd[0] = tools[cmd[0]]
243
244  input_bytes = args.input.read()
245
246  ret = 0
247  if not dump_to_file(args.output_configs, dump_configs, input_bytes,
248                      "configs in {}".format(args.input.name)):
249    ret = 1
250  if not dump_to_file(args.output_version, dump_version, input_bytes,
251                      "version in {}".format(args.input.name)):
252    ret = 1
253  if not dump_to_file(args.output_release, dump_release, input_bytes,
254                      "kernel release in {}".format(args.input.name)):
255    ret = 1
256
257  if not dump_to_file(args.output_compiler, dump_compiler, input_bytes,
258                      "kernel compiler in {}".format(args.input.name)):
259    ret = 1
260
261  return ret
262
263
264if __name__ == '__main__':
265  sys.exit(main())
266