'''
Freeze capa features.

Example usage:
  freeze_features.py <test files dir>
  freeze_features.py samples\benign
'''

import os
import sys
import time
import logging

import argparse

from scripts.testbed import FREEZE_EXTENSION
from capa.features.freeze import main as freeze_features


# only process files with these extensions
TARGET_EXTENSIONS = [
    '.mal_',
    '.exe_',
    '.dll_',
    '.sys_'
]


logger = logging.getLogger('check_rule')


def freeze(input_path, reprocess):
    if not os.path.exists(input_path):
        raise IOError('%s does not exist or cannot be accessed' % input_path)

    if os.path.isfile(input_path):
        outfile = '%s%s' % (input_path, FREEZE_EXTENSION)
        freeze_file(input_path, outfile, reprocess)

    elif os.path.isdir(input_path):
        logger.info('freezing features of %s files in %s', '|'.join(TARGET_EXTENSIONS), input_path)
        for root, dirs, files in os.walk(input_path):
            for file in files:
                if not os.path.splitext(file)[1] in TARGET_EXTENSIONS:
                    logger.debug('skipping non-target file: %s', file)
                    continue
                path = os.path.join(root, file)
                outfile = '%s%s' % (path, FREEZE_EXTENSION)
                freeze_file(path, outfile, reprocess)


def freeze_file(path, output, reprocess=False):
    logger.info('freezing features of %s', path)

    if os.path.exists(output) and not reprocess:
        logger.info('%s already exists, provide -r argument to reprocess', output)
        return

    try:
        freeze_features([path, output])  # args: sample, output
    except Exception as e:
        logger.error('could not freeze features for %s: %s', path, str(e))


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

    parser = argparse.ArgumentParser(description='Freeze capa features of a file or of files in a directory')
    parser.add_argument('file_path', type=str,
                        help='Path to file or directory to analyze')
    parser.add_argument('-r', '--reprocess', action='store_true', default=False,
                        help='Overwrite existing analysis')
    parser.add_argument('-v', '--verbose', action='store_true',
                        help='Enable verbose output')
    parser.add_argument('-q', '--quiet', action='store_true',
                        help='Disable all output but errors')
    args = parser.parse_args(args=argv)

    if args.quiet:
        logging.basicConfig(level=logging.ERROR)
        logging.getLogger().setLevel(logging.ERROR)
    elif args.verbose:
        logging.basicConfig(level=logging.DEBUG)
        logging.getLogger().setLevel(logging.DEBUG)
    else:
        logging.basicConfig(level=logging.INFO)
        logging.getLogger().setLevel(logging.INFO)

    time0 = time.time()
    try:
        freeze(args.file_path, args.reprocess)
    except IOError as e:
        logger.error('%s', str(e))
        return -1

    logger.info('freezing features took %d seconds', time.time() - time0)
    return 0


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