#!/usr/local/bin/python2.7

import os
import imp
import sys
import time
import inspect
import cProfile
import argparse
import threading
import traceback


import vivisect
import vivisect.cli as viv_cli
import envi.config as e_config
import envi.threads as e_threads
import vivisect.parsers as viv_parsers

def main():
    parser = argparse.ArgumentParser(prog='vivbin', usage='%(prog)s [options] <workspace|binaries...>')
    parser.add_argument('-M', '--module', dest='modname', default=None, action='store',
                        help='run the file listed as an analysis module in non-gui mode and exit')
    parser.add_argument('-A', '--skip-analysis', dest='doanalyze', default=True, action='store_false',
                        help='Do *not* do an initial auto-analysis pass')
    parser.add_argument('-B', '--bulk', dest='bulk', default=False, action='store_true',
                        help='Do *not* start the gui, just load, analyze and save')
    parser.add_argument('-C', '--cprofile', dest='cprof', default=False, action='store_true',
                        help='Output vivisect performace profiling (cProfile) info')
    parser.add_argument('-O', '--option', dest='option', default=None, action='store',
                        help='<secname>.<optname>=<optval> (optval must be json syntax)')
    parser.add_argument('-p', '--parser', dest='parsemod', default=None, action='store',
                        help='Manually specify the parser module (pe/elf/blob/...)')
    parser.add_argument('-s', '--storage', dest='storage_name', default=None, action='store',
                        help='Specify a storage module by name')
    parser.add_argument('-v', '--verbose', dest='verbose', default=False, action='store_true',
                        help='Enable verbose mode')
    parser.add_argument('-V', '--version', dest='version', default=None, action='store',
                        help='Add file version (if available) to save file name')
    parser.add_argument('file', nargs='*')
    args = parser.parse_args()

    vw = viv_cli.VivCli()
    vw.verbose = args.verbose

    if args.option is not None:
        if args.option in ('-h','?'):
            print(vw.config.reprConfigPaths())
            sys.exit(-1)

        try:
            vw.config.parseConfigOption(args.option)
        except e_config.ConfigNoAssignment, e:
            print(vw.config.reprConfigPaths() + "\n")
            print(e)
            print("syntax: \t-O <secname>.<optname>=<optval> (optval must be json syntax)")
            sys.exit(-1)

        except Exception, e:
            print(vw.config.reprConfigPaths())
            print("With entry: %s" % args.option)
            print(e)
            sys.exit(-1)



    if args.storage_name is not None:
        vw.setMeta("StorageModule", args.storage_name)

    # If we're not gonna load files, no analyze
    if args.file is None:
        args.doanalyze = False

    # Load in any additional files...
    needanalyze = False
    if args.file is not None:
        for fname in args.file:
            if args.parsemod == None:
                args.parsemod = viv_parsers.guessFormatFilename(fname)

            start = time.time()
            if args.parsemod == 'viv':
                vw.loadWorkspace(fname)
            else:
                needanalyze = True
                vw.loadFromFile(fname, fmtname=args.parsemod)

            end = time.time()
            print('Loaded (%.4f sec) %s' % (end - start, fname))

    if args.bulk:
        if args.doanalyze:
            if args.cprof:
                cProfile.run("vw.analyze()")
            else:
                start = time.time()
                vw.analyze()
                end = time.time()
                print "ANALYSIS TIME: %s" % (end-start)

        if args.modname is not None:
            module = imp.load_module("custom_analysis", file(modname, "rb"), modname, ('.py', 'U', 1))
            module.analyze(vw)

        print('stats: %r' % (vw.getStats(),))
        print("Saving workspace: %s" % (vw.getMeta('StorageName')))

        vw.saveWorkspace()

    else:

        import vivisect.qt.main as viv_qt_main

        # If we are interactive, lets turn on extended output...
        vw.verbose = True
        if args.doanalyze and needanalyze:
            e_threads.firethread(vw.analyze)()

        viv_qt_main.main(vw)

if __name__ == '__main__':
    try:
        # psyco makes disasm much faster (2-3X)
        import psyco
        #psyco.log()
        psyco.full()
    except ImportError:
        pass

    main()

