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