funcdb_util.py

   1 #!/usr/bin/env python3
   2 import sys
   3 import os
   4 import argparse
   5
   6 import yaml
   7
   8 import yamlutils
   9
  10
  11 argp = argparse.ArgumentParser(description="Perform various transformations on a function database")
  12 argp.add_argument("file", help="function database file (YAML)")
  13 argp.add_argument("command", help="transformation to perform")
  14 argp.add_argument("args", nargs="*", help="transformation arguments (optional)")
  15 args = argp.parse_args()
  16
  17 with open(args.file) as f:
  18     FUNC_DB = yaml.load(f)
  19
  20 label2addr = {}
  21
  22 for addr, props in FUNC_DB.items():
  23     label2addr[props["label"]] = addr
  24
  25
  26 if args.command == "label2addr":
  27
  28     for addr, props in FUNC_DB.items():
  29         props["calls_addr"] = [label2addr.get(x, x) for x in props["calls"]]
  30
  31 elif args.command == "addr2label":
  32     for addr, props in FUNC_DB.items():
  33         print(addr)
  34         props["calls"] = [FUNC_DB[x]["label"] if x in FUNC_DB else x for x in props["calls_addr"]]
  35
  36 elif args.command == "called_by":
  37     called = {}
  38
  39     for addr, props in FUNC_DB.items():
  40         calls_unk = []
  41         name = props["label"]
  42         for callee in props.get("calls", []):
  43             called.setdefault(callee, set()).add(name)
  44             if callee not in label2addr:
  45                 calls_unk.append(callee)
  46         if calls_unk:
  47             props["calls_unk"] = calls_unk
  48
  49     for addr, props in FUNC_DB.items():
  50         name = props["label"]
  51         if name in called:
  52             props["called_by"] = called[name]
  53
  54 elif args.command == "returns":
  55     for addr, props in FUNC_DB.items():
  56         if "modifieds" in props and "callsites_live_out" in props:
  57             props["returns"] = set(props["modifieds"]) & set(props["callsites_live_out"])
  58
  59 elif args.command == "select-subgraph":
  60     if len(args.args) > 1:
  61         dirname = args.args[1]
  62     else:
  63         dirname = args.args[0] + ".subgraph"
  64     if not os.path.isdir(dirname):
  65         os.makedirs(dirname)
  66     else:
  67         print("Warning: already exists:", dirname)
  68
  69     queue = [args.args[0]]
  70     seen = set()
  71     while queue:
  72         func = queue.pop()
  73         if func in seen:
  74             continue
  75         seen.add(func)
  76         addr = label2addr[func]
  77         funcinfo = FUNC_DB[addr]
  78         print(func)
  79         #print(funcinfo)
  80         queue.extend(funcinfo["calls"])
  81         fname = "%s-%s.lst" % (addr, func)
  82         try:
  83             os.symlink("../funcs/" + fname, dirname + "/" + fname)
  84         except FileExistsError as e:
  85             print("Warning:", e)
  86
  87 else:
  88     argp.error("Unknown command: " + args.command)
  89
  90 os.rename(args.file, args.file + ".bak")
  91
  92 with open(args.file, "w") as f:
  93     yaml.dump(FUNC_DB, f)