1143aa5cbSKonrad Sztyber#!/usr/bin/env python3 2*17538bdcSpaul luse# SPDX-License-Identifier: BSD-3-Clause 3*17538bdcSpaul luse# Copyright (C) 2021 Intel Corporation 4*17538bdcSpaul luse# All rights reserved. 5*17538bdcSpaul luse# 6143aa5cbSKonrad Sztyber 7143aa5cbSKonrad Sztyberfrom argparse import ArgumentParser 8143aa5cbSKonrad Sztyberimport os 9143aa5cbSKonrad Sztyberimport re 10143aa5cbSKonrad Sztyberimport subprocess 11143aa5cbSKonrad Sztyberimport sys 12143aa5cbSKonrad Sztyber 13143aa5cbSKonrad Sztyber 14143aa5cbSKonrad Sztyberclass TraceProcess: 15143aa5cbSKonrad Sztyber def __init__(self, pid): 16143aa5cbSKonrad Sztyber self._path = os.readlink(f'/proc/{pid}/exe') 17143aa5cbSKonrad Sztyber self._pid = pid 18143aa5cbSKonrad Sztyber self._probes = self._init_probes() 19143aa5cbSKonrad Sztyber 20143aa5cbSKonrad Sztyber def _init_probes(self): 21143aa5cbSKonrad Sztyber lines = subprocess.check_output(['bpftrace', '-l', '-p', str(self._pid)], text=True) 22143aa5cbSKonrad Sztyber probes = {} 23143aa5cbSKonrad Sztyber for line in lines.split('\n'): 24143aa5cbSKonrad Sztyber parts = line.split(':') 25143aa5cbSKonrad Sztyber if len(parts) < 3: 26143aa5cbSKonrad Sztyber continue 27143aa5cbSKonrad Sztyber ptype, path, function = parts[0], parts[1], parts[-1] 28143aa5cbSKonrad Sztyber probes[(ptype, function)] = path 29143aa5cbSKonrad Sztyber return probes 30143aa5cbSKonrad Sztyber 31143aa5cbSKonrad Sztyber def fixup(self, script): 32143aa5cbSKonrad Sztyber pregs = [re.compile(r'({}):__EXE__:(\w+)'.format(ptype)) for ptype in ['usdt', 'uprobe']] 33143aa5cbSKonrad Sztyber with open(script, 'r') as file: 34143aa5cbSKonrad Sztyber lines = file.readlines() 35143aa5cbSKonrad Sztyber result = '' 36143aa5cbSKonrad Sztyber for line in lines: 37143aa5cbSKonrad Sztyber for regex in pregs: 38143aa5cbSKonrad Sztyber match = regex.match(line) 39143aa5cbSKonrad Sztyber if match is not None: 40143aa5cbSKonrad Sztyber ptype, function = match.groups() 41143aa5cbSKonrad Sztyber path = self._probes.get((ptype, function), self._path) 42143aa5cbSKonrad Sztyber line = line.replace('__EXE__', path) 43143aa5cbSKonrad Sztyber break 44143aa5cbSKonrad Sztyber result += line.replace('__EXE__', self._path).replace('__PID__', str(self._pid)) 45143aa5cbSKonrad Sztyber return result 46143aa5cbSKonrad Sztyber 47143aa5cbSKonrad Sztyber 48143aa5cbSKonrad Sztyberif __name__ == '__main__': 49143aa5cbSKonrad Sztyber parser = ArgumentParser(description='bpftrace script generator replacing special ' + 50143aa5cbSKonrad Sztyber 'variables in the scripts with appropriate values') 51143aa5cbSKonrad Sztyber parser.add_argument('-p', '--pid', type=int, required=True, help='PID of a traced process') 52143aa5cbSKonrad Sztyber parser.add_argument('scripts', metavar='SCRIPTS', type=str, nargs='+', 53143aa5cbSKonrad Sztyber help='bpftrace scripts to process') 54143aa5cbSKonrad Sztyber args = parser.parse_args(sys.argv[1:]) 55143aa5cbSKonrad Sztyber proc = TraceProcess(args.pid) 56143aa5cbSKonrad Sztyber for script in args.scripts: 57143aa5cbSKonrad Sztyber print(proc.fixup(script)) 58