startop: iorap script to create custom compiled TraceFile

#!/usr/bin/env python3

#

# Copyright 2019, The Android Open Source Project

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

#     http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

import argparse

import asyncio

import csv

import itertools

import os

import re

import struct

import sys

import tempfile

import time

import zipfile

from typing import Any, Callable, Dict, Generic, Iterable, List, NamedTuple, TextIO, Tuple, TypeVar, Optional, Union

# Include generated protos.

dir_name = os.path.dirname(os.path.realpath(__file__))

sys.path.append(dir_name + "/generated")

from TraceFile_pb2 import *

def parse_options(argv: List[str] = None):

  """Parse command line arguments and return an argparse Namespace object."""

  parser = argparse.ArgumentParser(description="Compile a TraceFile.proto from a manual text file.")

  # argparse considers args starting with - and -- optional in --help, even though required=True.

  # by using a named argument group --help will clearly say that it's required instead of optional.

  required_named = parser.add_argument_group('required named arguments')

  # optional arguments

  # use a group here to get the required arguments to appear 'above' the optional arguments in help.

  optional_named = parser.add_argument_group('optional named arguments')

  optional_named.add_argument('-opb', '--output-proto-binary', dest='output_proto_binary', action='store', help='Write binary proto output to file.')

  optional_named.add_argument('-pm', '--pinlist-meta', dest='pinlist_meta', action='store', help='Path to pinlist.meta (default=none) binary file.')

  optional_named.add_argument('-pmp', '--pinlist-meta-parent', dest='pinlist_meta_parent', action='store', help='Device path that the pinlist.meta applies to (e.g. /data/.../somefile.apk)')

  optional_named.add_argument('-i', '--input', dest='input', action='store', help='Input text file (default stdin).')

  optional_named.add_argument('-zp', '--zip_path', dest='zip_path', action='append', help='Directory containing zip files.')

  optional_named.add_argument('-d', '--debug', dest='debug', action='store_true', help='Add extra debugging output')

  optional_named.add_argument('-ot', '--output-text', dest='output_text', action='store', help='Output text file (default stdout).')

  return parser.parse_args(argv)

# TODO: refactor this with a common library file with analyze_metrics.py

def _debug_print(*args, **kwargs):

  """Print the args to sys.stderr if the --debug/-d flag was passed in."""

  if _debug:

    print(*args, **kwargs, file=sys.stderr)

class BadInputError(Exception):

  pass

InputRecord = NamedTuple('InputRecord', [('filepath', str), ('offset', int), ('length', int), ('remark', str)])

def find_zip_in_paths(original_name, zip_paths):

  # /foo/bar/bax.zip -> bax.zip

  file_basename = os.path.split(original_name)[1]

  # the file must be located in one of the --zip-path arguments

  matched = None

  for zip_path in zip_paths:

    for dir_entry in os.listdir(zip_path):

      if dir_entry == file_basename:

        matched = os.path.join(zip_path, dir_entry)

        break

    if matched:

      break

  if not matched:

    raise ValueError("%s could not be found in any of the --zip_path specified." %(file_basename))

  _debug_print("found zip file ", file_basename, " in ", matched)

  if not zipfile.is_zipfile(matched):

    raise ValueError("%s is not a zip file" %(matched))

  return matched

def handle_zip_entry(input_record, zip_paths):

  res = re.match("([^!]+)[!](.*)", input_record.filepath)

  if not res:

    return input_record

                         # 'foo!bar'

  in_filepath = res[1]   # -> 'foo'

  in_zip_entry = res[2]  # -> 'bar'

  matched = find_zip_in_paths(in_filepath, zip_paths)

  zip = zipfile.ZipFile(matched)

  try:

    zip_info = zip.getinfo(in_zip_entry)

  except KeyError:

    raise ValueError("%s is not an item in the zip file %s" %(in_zip_entry, matched))

  # TODO: do we also need to add header size to this?

  in_offset = zip_info.header_offset

  # TODO: if a range is specified, use that instead.

  in_length = zip_info.compress_size

  return InputRecord(in_filepath, in_offset, in_length, 'zip entry (%s)' %(in_zip_entry))

def parse_input_file(input: Iterable[str], zip_paths: List[str]) -> Iterable[InputRecord]:

  for line in input:

    line = line.strip()

    _debug_print("Line = ", line)

    if not line:

      _debug_print("  skip empty line", line)

      continue

    elif line[0] == "#":

      _debug_print("  skip commented line", line)

      continue

    res = re.match("([^\s]+)\s+(\d+)\s+(\d+)", line)

    if not res:

      raise BadInputError("Expected input of form: <str:filepath> <int:offset> <int:length>")

    in_filepath = res[1]

    in_offset = int(res[2])

    in_length = int(res[3])

    yield handle_zip_entry(InputRecord(in_filepath, in_offset, in_length, 'regular file'), zip_paths)

# format:

#   (<big_endian(i32):file_offset> <big_endian(i32):range_length>)+

PIN_META_FORMAT = ">ii"

PIN_META_READ_SIZE = struct.calcsize(PIN_META_FORMAT)

def parse_pin_meta(pin_meta_file, pinlist_meta_parent, zip_paths):

  if not pin_meta_file:

    return ()

  global PIN_META_FORMAT

  global PIN_META_READ_SIZE

  # '/data/app/com.google.android.GoogleCamera-aNQhzSznf4h_bvJ_MRbweQ==/base.apk'

  #  -> 'com.google.android.GoogleCamera'

  package_name_match = re.match('/.*/(.*)-.*=/base.apk', pinlist_meta_parent)

  if not package_name_match:

    raise ValueError("%s did not contain the <packagename>.apk" %(pinlist_meta_parent))

  package_name = package_name_match[1]

  # "com.google.android.GoogleCamera" -> "GoogleCamera.apk"

  apk_name = package_name.split(".")[-1] + ".apk"

  path_to_zip_on_host = find_zip_in_paths(apk_name, zip_paths)

  apk_file_size = os.path.getsize(path_to_zip_on_host)

  _debug_print("APK path '%s' file size '%d'" %(path_to_zip_on_host, apk_file_size))

  while True:

    data = pin_meta_file.read(PIN_META_READ_SIZE)

    if not data:

      break

    (pin_offset, pin_length) = struct.unpack(PIN_META_FORMAT, data)  # (offset, length)

    remark = 'regular file (pinlist.meta)'

    remaining_size = apk_file_size - pin_offset

    if remaining_size < 0:

      print("WARNING: Clamp entry (%d, %d), offset too large (max file size = %d)" %(pin_offset, pin_length, apk_file_size))

      pin_length = pin_length + remaining_size

      pin_offset = pin_offset + remaining_size

      if pin_offset < 0:

        pin_offset = 0

      remark += '[clamped.offset]'

    pin_last_offset = pin_offset + pin_length

    remaining_size = apk_file_size - pin_last_offset

    if remaining_size < 0:

      print("WARNING: Clamp entry (%d, %d), length too large (max file size = %d)" %(pin_offset, pin_length, apk_file_size))

      pin_length = pin_length + remaining_size

      remark += '[clamped.length]'

    yield InputRecord(pinlist_meta_parent, pin_offset, pin_length, remark)

def write_text_file_output(input_records: Iterable[InputRecord], output_text_file):

  for rec in input_records:

    output_text_file.write("%s %d %d #%s\n" %(rec.filepath, rec.offset, rec.length, rec.remark))

def build_trace_file(input_records: Iterable[InputRecord]) -> TraceFile:

  trace_file = TraceFile()

  trace_file_index = trace_file.index

  file_id_counter = 0

  file_id_map = {} # filename -> id

  stats_length_total = 0

  filename_stats = {} # filename -> total size

  for rec in input_records:

    filename = rec.filepath

    file_id = file_id_map.get(filename)

    if not file_id:

      file_id = file_id_counter

      file_id_map[filename] = file_id_counter

      file_id_counter = file_id_counter + 1

      file_index_entry = trace_file_index.entries.add()

      file_index_entry.id = file_id

      file_index_entry.file_name = filename

    # already in the file index, add the file entry.

    file_entry = trace_file.list.entries.add()

    file_entry.index_id = file_id

    file_entry.file_length = rec.length

    stats_length_total += file_entry.file_length

    file_entry.file_offset = rec.offset

    filename_stats[filename] = filename_stats.get(filename, 0) + file_entry.file_length

  return trace_file

def main():

  global _debug

  options= parse_options()

  _debug = options.debug

  _debug_print("parsed options: ", options)

  if not options.input:

    input_file = sys.stdin

    _debug_print("input = stdin")

  else:

    input_file = open(options.input)

    _debug_print("input = (file)", options.input)

  if not options.output_proto_binary:

    output_proto_file = None

  else:

    output_proto_file = open(options.output_proto_binary, 'wb')

  _debug_print("output_proto_binary = ", output_proto_file)

  pinlist_meta_parent = options.pinlist_meta_parent

  if options.pinlist_meta:

    pin_meta_file = open(options.pinlist_meta, 'rb')

  else:

    pin_meta_file = None

  if (pinlist_meta_parent == None) != (pin_meta_file == None):

    print("Options must be used together: --pinlist-meta and --pinlist-meta-path")

    return 1

  if not options.output_text:

    output_text_file = sys.stdout

    _debug_print("output = stdout")

  else:

    output_text_file = open(options.output_text, 'w')

    _debug_print("output = (file)", options.output_text)

  zip_paths = options.zip_path or []

  input_records = list(parse_pin_meta(pin_meta_file, pinlist_meta_parent, zip_paths))

  input_records = input_records + list(parse_input_file(input_file, zip_paths))

  for p in input_records:

    _debug_print(p)

  write_text_file_output(input_records, output_text_file)

  output_text_file.close()

  out_proto = build_trace_file(input_records)

  if output_proto_file:

    output_proto_file.write(out_proto.SerializeToString())

    output_proto_file.close()

  return 0

if __name__ == '__main__':

  sys.exit(main())