#!/home/other/Compil/python-haystack/venv/bin/python
# -*- coding: utf-8 -*-
#
# Copyright (C) 2011 Loic Jaquemet loic.jaquemet+python@gmail.com
#

__author__ = "Loic Jaquemet loic.jaquemet+python@gmail.com"

__doc__ = '''
Search for struct in memory.
MemoryHandler can be a file dump or a live process.

Output offset and values of possible matches.
Output type can be txt, json, pickled.
'''

import argparse
import logging
import sys

import os

from haystack import argparse_utils
from haystack import cli

log = logging.getLogger('haystack')


def argparser():
    """
        Builds the argparse tree.
        See the command line --help .
    """
    rootparser = argparse.ArgumentParser(prog='haystack', description='Live memory tools on heap allocators.')
    rootparser.add_argument('--debug', dest='debug', action='store_const', const=True, help='setLevel to DEBUG')
    rootparser.add_argument('--quiet', dest='quiet', action='store_const', const=True, help='setLevel to ERROR only')
    rootparser.add_argument('--interactive', dest='interactive', action='store_const', const=True, help='drop to python command line after action')
    rootparser.add_argument('--nommap', dest='mmap', action='store_const', const=False, default=True, help='disable mmap()-ing')

    target = rootparser.add_mutually_exclusive_group(required=True)
    target.add_argument('--pid', type=int, help='Target PID')
    target.add_argument('--memfile', type=argparse.FileType('r'), help='Use a file memory dump instead of a live process ID')
    target.add_argument('--dumpname', type=argparse_utils.readable, help='Use a haystack memory dump instead of a live process ID')
    target.add_argument('--volname', type=argparse.FileType('r'), help='Use a OS memory dump')

    #FIXME
    rootparser.add_argument('--volpid', type=int, help='Target PID for the OS --volname memory parser')

    output = rootparser.add_mutually_exclusive_group(required=True)
    output.add_argument('--string', dest='human', action='store_const', const=True, help='Print results as human readable string')
    output.add_argument('--json', dest='json', action='store_const', const=True, help='Print results as json readable string')
    output.add_argument('--pickled', dest='pickled', action='store_const', const=True, help='Print results as pickled string')

    subparsers = rootparser.add_subparsers(help='sub-command help')
    search_parser = subparsers.add_parser('search', help='search help')
    search_parser.add_argument('struct_name', type=str, help='Structure type name')
    search_parser.add_argument('--constraints_file', type=argparse.FileType('r'), help='Constraints filename for the module')
    search_parser.add_argument('--extended', dest='extended_search', action='store_const', const=True, help='do no restrict search to allocated chunks')
    search_parser.set_defaults(func=cli.search_cmdline)
    #
    refresh_parser = subparsers.add_parser('refresh', help='refresh help')
    refresh_parser.add_argument('struct_name', type=str, help='Structure type name')
    refresh_parser.add_argument('addr', type=argparse_utils.int16, help='Structure memory address')
    refresh_parser.add_argument('--validate', dest='validate', action='store_const', const=True, help='validate the record after loading')
    refresh_parser.add_argument('--constraints_file', type=argparse.FileType('r'), help='Constraints filename for the module')
    refresh_parser.set_defaults(func=cli.refresh)
    #
    watch_parser = subparsers.add_parser('watch', help='watch a structure in a live process')
    watch_parser.add_argument('struct_name', type=str, help='Structure type name')
    watch_parser.add_argument('addr', type=argparse_utils.int16, help='Structure memory address')
    watch_parser.add_argument('refresh_rate', type=int, action='store', default=0, help='Seconds between refresh')
    watch_parser.add_argument('varname', type=str, action='store', default=None, help='structure member name (eg. pointername.valuename)')
    watch_parser.set_defaults(func=cli.watch)
    return rootparser


def main(argv):

    parser = argparser()
    opts = parser.parse_args(argv)

    level = logging.INFO
    if opts.debug:
        level=logging.DEBUG
    elif opts.quiet:
        level=logging.ERROR

    if opts.debug:
        flog = os.path.normpath('log')
        # FORMAT = '%(relativeCreated)d %(message)s'
        # logging.basicConfig(format=FORMAT, level=level, filename=flog, filemode='w')
        logging.basicConfig(level=level, filename=flog, filemode='w')
        print ('[+] **** COMPLETE debug log to %s'%(flog))
    else:
        logging.basicConfig(level=level)

    if opts.volname:
        opts.pid = opts.volpid

    sh = logging.StreamHandler(sys.stdout) # 2.6, 2.7 compat
    # logging.getLogger('abouchet').addHandler( sh )
    logging.getLogger('haystack').addHandler(sh)

    if opts.json:
        log.warning('the JSON feature is experimental and probably wont work.')
    try:
        opts.func(opts)
    except ImportError,e:
        log.error('Structure type does not exists.')
        log.error('sys.path is %s'%sys.path)
        print e

    if opts.pid:
        log.debug("done for pid %d"%opts.pid)
    elif opts.memfile:
        log.debug("done for file %s"%opts.memfile.name)
    elif opts.dumpname:
        log.debug("done for file %s"%opts.dumpname)
    return 0


if __name__ == "__main__":
    sys.path.append(os.getcwd())
    main(sys.argv[1:])


