#!/usr/bin/env python
#
# USAGE:    place-probe [-h] [--dry-run] [--debug] PID METHOD
#
# This tool helps place dynamic probes on .NET methods that were
# CrossGen-generated (compiled ahead of time). To use the tool,
# the CrossGen-generated assemblies need to have perfmaps generated
# by CrossGen /CreatePerfMap, expected in the /tmp directory.
#
# Copyright (C) 2018, Sasha Goldshtein
# Licensed under the MIT License

import argparse
import os
import re
import subprocess

class Section(object):
    def __init__(self, start, perms, offset, path):
        self.start = int(start, 16)
        self.perms = perms
        self.offset = int(offset, 16)
        self.path = path

def assembly_from_map_file(map_file):
    return re.match("/tmp/(.*)\.ni\.{.*}.map", map_file).group(1) + ".dll"

def all_sections(pid):
    sections = {}
    with open("/proc/%d/maps" % pid, "r") as maps:
        for line in maps:
            match = re.match(r"(\S+)-\S+\s+(\S+)\s+(\S+)\s+\S+\s+\S+\s+(\S+)", line.strip())
            if match is None:
                continue
            start, perms, offset, path = match.group(1, 2, 3, 4)
            if '/' not in path:
                continue
            filename = os.path.basename(path)
            section = Section(start, perms, offset, path)
            if filename in sections:
                sections[filename].append(section)
            else:
                sections[filename] = [section]
    return sections

def place_probe(path, offset):
    command = "sudo perf probe -x %s --add 0x%x" % (path, offset)
    if args.dry_run:
        print(command)
    else:
        subprocess.check_call(command, shell=True)

parser = argparse.ArgumentParser(description="Place dynamic tracing probes on a managed method " +
    "that resides in a crossgen-compiled assembly. For .NET Core on Linux.",
    epilog="EXAMPLE: ./place-probe.py 1234 'Thread::Sleep'")
parser.add_argument("pid", type=int, help="the dotnet process id")
parser.add_argument("symbol", type=str, help="the symbol on which to place the probe")
parser.add_argument("--dry-run", action="store_true",
                    help="print the symbol and the command but don't place the probe")
parser.add_argument("--debug", action="store_true", help="print diagnostic information")
args = parser.parse_args()

sections = all_sections(args.pid)

output = subprocess.check_output("grep '%s' /tmp/*.ni.*.map" % args.symbol,
                                 shell=True)
for line in output.strip().split('\n'):
    parts = line.split()
    map_file, address = parts[0].split(':')
    address = int(address, 16)
    assembly = assembly_from_map_file(map_file)

    symbol = str.join(' ', parts[2:])
    if args.dry_run:
        print("\n" + symbol)

    first_section = sections[assembly][0]
    exec_section = [section for section in sections[assembly] if 'r-xp' == section.perms][0]
    offset_from_first = exec_section.start - first_section.start
    offset_in_file = exec_section.offset
    final_address = address - offset_from_first + offset_in_file

    if args.debug:
        print("address: %x, offset_from_first: %x, offset_in_file: %x, final_address: %x" %
              (address, offset_from_first, offset_in_file, final_address))

    place_probe(exec_section.path, final_address)