#!/usr/bin/env python2
'''
show the features extracted by capa.
'''
import sys
import logging

import argparse

import capa.main
import capa.rules
import capa.engine
import capa.features
import capa.features.freeze
import capa.features.extractors.viv


def main(argv=None):
    if argv is None:
        argv = sys.argv[1:]

    formats = [
        ('auto', '(default) detect file type automatically'),
        ('pe', 'Windows PE file'),
        ('sc32', '32-bit shellcode'),
        ('sc64', '64-bit shellcode'),
        ('freeze', 'features previously frozen by capa'),
    ]
    format_help = ', '.join(['%s: %s' % (f[0], f[1]) for f in formats])

    parser = argparse.ArgumentParser(description='detect capabilities in programs.')
    parser.add_argument('sample', type=str,
                        help='Path to sample to analyze')
    parser.add_argument('-f', '--format', choices=[f[0] for f in formats], default='auto',
                        help='Select sample format, %s' % format_help)
    parser.add_argument('-F', '--function', type=lambda x: int(x, 0),
                        help='Show features for specific function')
    args = parser.parse_args(args=argv)

    logging.basicConfig(level=logging.INFO)
    logging.getLogger().setLevel(logging.INFO)

    if args.format == 'freeze':
        with open(args.sample, 'rb') as f:
            extractor = capa.features.freeze.load(f.read())
    else:
        vw = capa.main.get_workspace(args.sample, args.format)
        extractor = capa.features.extractors.viv.VivisectFeatureExtractor(vw, args.sample)

    if not args.function:
        for feature, va in extractor.extract_file_features():
            if va:
                print('file: 0x%08x: %s' % (va, feature))
            else:
                print('file: 0x00000000: %s' % (feature))

    functions = extractor.get_functions()

    if args.function:
        if args.format == 'freeze':
            functions = filter(lambda f: f == args.function, functions)
        else:
            functions = filter(lambda f: f.va == args.function, functions)

    for f in functions:
        for feature, va in extractor.extract_function_features(f):
            print('func: 0x%08x: %s' % (va, feature))

        for bb in extractor.get_basic_blocks(f):
            for feature, va in extractor.extract_basic_block_features(f, bb):
                print('bb  : 0x%08x: %s' % (va, feature))

            for insn in extractor.get_instructions(f, bb):
                for feature, va in extractor.extract_insn_features(f, bb, insn):
                    print('insn: 0x%08x: %s' % (va, feature))

    return 0


if __name__ == '__main__':
    sys.exit(main())