#!/usr/bin/env python3

# Copyright 2020 The Fuchsia Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

#### CATEGORY=Run, inspect and debug
### Extract pprof data from inspect.json
##
## inspect.json holds Inspect data generated as part of a snapshot (see fx snapshot).
## pprof is a tool used to analyze profiling data (for more details, see
## https://github.com/google/pprof). Some components are configured to provide
## the Inspect tool with pprof data that can later be used to analyze the running
## program's state (memory, threads, etc.). This is especially useful after a
## crash where developers can observe the last known state before the crash.
##
## This tool extracts the pprof data for some component from inspect.json and
## decodes it from its stored format (base64 encoded) into the format the pprof
## tool will understand (raw binary). The pprof files will be extracted to a
## temporary directory.
##
## usage: fx extract_pprof --inspect FILE --component CMX
##
##   --inspect   Path to the inspect.json file
##   --component Component to extract pprof data for (i.e. netstack.cmx)

import argparse
import base64
import json
import os
import sys
import tempfile


def main():
  parser = argparse.ArgumentParser(
      description="Extract pprof data from inspect.json")
  parser.add_argument(
      "--inspect", required=True, help="Path to inspect.json file")
  parser.add_argument(
      "--component",
      required=True,
      help="Component to extract pprof data for (i.e. netstack.cmx)")
  args = parser.parse_args()

  tempdir = tempfile.mkdtemp()

  print("inspect.json @ %s" % (args.inspect))
  print("component = %s" % (args.component))
  print("output @ %s" % (tempdir))

  with open(args.inspect, "r") as f:
    inspect = json.loads(f.read())

  for entry in inspect:
    if entry["data_source"] != "Inspect" or entry["moniker"] != args.component:
      continue

    filepath = entry["metadata"]["filename"]
    # We only care about pprof data.
    if not filepath.startswith("pprof"):
      continue

    filepath = os.path.join(tempdir, filepath)
    os.makedirs(filepath, exist_ok=True)

    data = entry["payload"]["root"]["pprof"]
    for k, v in data.items():
      # The first 4 characters are used to indicate that the file is a
      # base64 encoded file.
      prefix, v = v[:4], v[4:]
      if prefix != "b64:":
        print("%s at %s is not base64 encoded" % (k, filepath))

      with open(os.path.join(filepath, k), "wb") as f:
        f.write(base64.b64decode(v))


if __name__ == "__main__":
  main()