cpu_features.py

'''
Crude script to analyze which CPU features are used by a given executable/library.

Created on Jun 2, 2019
@author: SleepProgger
'''

import json
import argparse
import sys
import os.path as path
import re
import urllib2
from collections import defaultdict
import subprocess

def parse_gas_specs(filename):
    if filename.startswith(("http://", "https://")):
        print("Downloading file: '%s'" % filename)
        resp = urllib2.urlopen(filename)
        gas = resp.read()
    else:
        with open(filename, "rb") as fd:
            gas = fd.read() 
    print("Parsing gas.vim")
    tmp = re.findall(r"syn keyword (gasOpcode_[^\t ]+)[ \t]+(.*)", gas)
    opcodes = defaultdict(set)
    op_count = defaultdict()
    for k, v in tmp:
        codes = v.split(" ")
        for code in codes:
            op_count[code] = op_count.get(code, 0) + 1
        opcodes[k].update(codes)
    if False:
        # cleanup non unique ops
        for k, v in opcodes.items():
            for op in list(v):
                if op_count[op] > 1:
                    v.remove(op)
    result = defaultdict(dict)
    re_op_mapping = r"call <SID>MapOpcode\('([^']+)'[ \t]*,[ \t]*'([^']+)'[ \t]*,[ \t]*'([^']+)'\)"
    for k, arch, feature in re.findall(re_op_mapping, gas):
        result[arch][feature] = list(opcodes[k])
    return result

def find_command(cmd, specs):
    cmd = re.compile("^%s$" % cmd)
    for arch, features in specs.items():
        for feature, ops in features.items():
            for op in ops:
                if cmd.match(op):
                    print("%s %s => %s" % (arch, feature, op)) 
    

if __name__ == '__main__':
    own_path = path.dirname(sys.argv[0])
    parser = argparse.ArgumentParser(description='Tries to detect which CPU features where used in a given binary.')
    parser.add_argument("-j", "--json-specs", required=False, default=str(path.join(own_path, "specs.json")), help="json file containing a command to feature mapping.")
    parser.add_argument("-o", "--json-output", required=False, default=str(path.join(own_path, "specs.json")), help="json file to save the command to feature mapping parsed from an gas.vim file. Defaults to same folder as this scipt/specs.json")
    parser.add_argument("-g", "--gas", required=False, default="https://github.com/Shirk/vim-gas/raw/master/syntax/gas.vim", help="gas.vim file to convert to feature mapping.")
    parser.add_argument("-nw", "--no-json-save", action="store_true", required=False, help="Do not save converted mapping from gas.vim file.")
    parser.add_argument("-b", "--include-base", action="store_true", required=False, help="Include base instructions in the search.")
    parser.add_argument("-l", "--lookup-op", action="store_true", required=False, help="Lookup arch and feature for given command. Can be regex.")
    parser.add_argument("executable", help="The executable to analyze or the command to lookup if -l is set.")
    args = parser.parse_args()
    
    specs = None
    if path.isfile(args.json_specs):
        try:
            with open(args.json_specs, "rb") as fp:
                specs = json.load(fp)
        except Exception as e: # TODO: dirty
            print("Failed loading specs.json file from '%s'" % args.json_specs)
            print(e)
            exit(1) 
    if specs is None:
        try:
            specs = parse_gas_specs(args.gas)
            if not args.no_json_save:
                with open(args.json_output, "wb") as fp:
                    json.dump(specs, fp, sort_keys=True, indent=2)
        except Exception as e: # TODO: dirty
            print("Failed loading gas.vim file from '%s'" % args.gas)
            print(e)
            exit(1)
    
    if args.lookup_op:
        find_command(args.executable, specs)
        exit(0)
    
    print("Running objdump")
    found_ops = set()
    out = subprocess.check_output(["objdump", "-M", "intel", "--no-show-raw-insn", "-D", args.executable])
    for line in out.split("\n"):
        line = line.split()
        if len(line) < 3:
            continue
        op = line[1]
        found_ops.add(op)
    
    print("Found %i ops" % len(found_ops))
    found_features = set()
    for op in found_ops:
        for arch, features in specs.items():
            for feature, fops in features.items():
                if not args.include_base and feature == "base":
                    continue 
                if op in fops:
                    found_features.add((arch, feature, op))
                    
    foo = defaultdict(lambda: defaultdict(list))
    for arch, feature, op in found_features:
        foo[arch][feature].append(op)
    for arch in foo:
        print("\n%s:" % arch)
        for feature in foo[arch]:
            print("\t%s:" % feature)
            print("\t\t%s" % "\n\t\t".join(foo[arch][feature]))