1#!/usr/bin/env python3 2 3from argparse import ArgumentParser 4import os 5import re 6import subprocess 7import sys 8 9 10class TraceProcess: 11 def __init__(self, pid): 12 self._path = os.readlink(f'/proc/{pid}/exe') 13 self._pid = pid 14 self._probes = self._init_probes() 15 16 def _init_probes(self): 17 lines = subprocess.check_output(['bpftrace', '-l', '-p', str(self._pid)], text=True) 18 probes = {} 19 for line in lines.split('\n'): 20 parts = line.split(':') 21 if len(parts) < 3: 22 continue 23 ptype, path, function = parts[0], parts[1], parts[-1] 24 probes[(ptype, function)] = path 25 return probes 26 27 def fixup(self, script): 28 pregs = [re.compile(r'({}):__EXE__:(\w+)'.format(ptype)) for ptype in ['usdt', 'uprobe']] 29 with open(script, 'r') as file: 30 lines = file.readlines() 31 result = '' 32 for line in lines: 33 for regex in pregs: 34 match = regex.match(line) 35 if match is not None: 36 ptype, function = match.groups() 37 path = self._probes.get((ptype, function), self._path) 38 line = line.replace('__EXE__', path) 39 break 40 result += line.replace('__EXE__', self._path).replace('__PID__', str(self._pid)) 41 return result 42 43 44if __name__ == '__main__': 45 parser = ArgumentParser(description='bpftrace script generator replacing special ' + 46 'variables in the scripts with appropriate values') 47 parser.add_argument('-p', '--pid', type=int, required=True, help='PID of a traced process') 48 parser.add_argument('scripts', metavar='SCRIPTS', type=str, nargs='+', 49 help='bpftrace scripts to process') 50 args = parser.parse_args(sys.argv[1:]) 51 proc = TraceProcess(args.pid) 52 for script in args.scripts: 53 print(proc.fixup(script)) 54