#!/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 by a bug report (see fx bugreport).
## 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()