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

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

__doc__ = '''
  Reverse heap analysis.
'''

import argparse
import logging
import sys

import os

from haystack import argparse_utils
from haystack.reverse import config
from haystack.reverse import context
from haystack.reverse.heuristics import reversers
from haystack.reverse.heuristics import signature
from haystack import dump_loader

log = logging.getLogger('haystack-reverse')



def reverse_instances(opt):
    # go through the 4 step process. Double linked list, basic type, Pointers and graph.
    memory_handler = dump_loader.load(opt.dumpname)
    ctx = api.reverse_instances(memory_handler)
    return

def writeReversedTypes(opt):
    return NotImplementedError("obsolete")
    """
    reverse types from a memorydump, and write structure definition to file
    :param opt:
    :return:
    """
    ctx, sizeCache = signature.makeSizeCaches(opt.dumpname)
    ctx = signature.makeReversedTypes(ctx, sizeCache)
    outfile = file(config.get_cache_filename(config.REVERSED_TYPES_FILENAME, ctx.dumpname),'w')
    for revStructType in ctx.list_reversed_types():
        outfile.write(revStructType.to_string())
        #
        #print revStructType._instances
        #for a,lst in revStructType._instances.items():
        #    print lst.__dict__
        #    for i in lst:
        #        print i
        #        outfile.write(i.toString())
    outfile.close()
    log.info('[+] Wrote to %s', outfile.name)
    return

def groupStructures(opt):
    """
    show sorted structure instances groups to stdout
    :param opt:
    :return:
    """
    ctx, sizeCache = signature.makeSizeCaches(opt.dumpname)
    for chains in signature.buildStructureGroup(ctx, sizeCache, opt.size ):
        signature.printStructureGroups(ctx, chains, opt.address )
    return

def saveSignatures(opt):
    """
    translate a memdump into a signature based file NULL,POINTERS,OTHERS
    :param opt:
    :return:
    """
    ctx, sig = signature.makeSignatures(opt.dumpname)
    outfile = config.get_cache_filename(config.SIGNATURES_FILENAME, ctx.dumpname)
    file(outfile,'w').write(sig)
    log.info('[+] Signature written to %s'%(outfile))
    return


def show(opt):
    """
    Show a structure
    :param opt:
    :return:
    """
    log.info('[+] Load ctx')
    memory_handler = dump_loader.load(opt.dumpname)
    ctx = context.get_context_for_address(memory_handler, opt.address)
    log.info('[+] Find Structure at: @%x', opt.address)
    try:
        st = ctx.get_record_at_address(opt.address)
        st.decodeFields()
        print st.to_string()
    except ValueError,e:
        log.info('[+] Found no structure.')
        return
    return

def printParents(opt):
    """
    print the parental allocators
    :param opt:
    :return:
    """
    log.info('[+] Load ctx')
    memory_handler = dump_loader.load(opt.dumpname)
    ctx = context.get_context_for_address(memory_handler, opt.address)
    log.info('[+] find offsets of struct_addr:%x', opt.address)
    i = 0
    try:
        child_address = ctx.get_record_address_at_address(opt.address)
        for st in ctx.listStructuresForPointerValue(child_address):
            st.decodeFields()
            print st.to_string()
            i+=1
    except ValueError,e:
        log.info('[+] Found no allocators.')
        return
    log.info('[+] Found %d allocators.'%( i ))
    return


def clean(opt):
    log.info('[+] Cleaning cache')
    ctx = config.remove_cache_folder(opt.dumpname)


def graph(opt):
    ''' show sorted structure instances groups to gefx '''
    #log.info('[+] Graphing')
    #ctx, sizeCache = signature.makeSizeCaches(opt.dumpname)
    #for chains in signature.buildStructureGroup(ctx, sizeCache, opt.size ):
    #  signature.graphStructureGroups(ctx, chains, opt.address )
    # TODO change to generic fn, and output graph
    return


def argparser():
    rootparser = argparse.ArgumentParser(prog='haystack-reverse',
                                         description='Several tools to reverse engineer allocators on the heap.')

    rootparser.add_argument('--debug', action='store_true', help='Debug mode on.')
    rootparser.add_argument('dumpname', type=argparse_utils.readable, action='store', help='Source memory dump by haystack.')

    subparsers = rootparser.add_subparsers(help='sub-command help')

    instances = subparsers.add_parser('instances',
                                      help='List all allocators instances with virtual address, member types guess and info.')
    instances.set_defaults(func=reverse_instances)

    # not refactored yet
    typemap = subparsers.add_parser('typemap',
                                    help='Try to reverse generic types from instances\' similarities.')
    typemap.set_defaults(func=writeReversedTypes)

    # not refactored yet
    #groupparser = subparsers.add_parser('group', help='Show structure instances groups by size and signature.')
    #groupparser.add_argument('--size', type=int, action='store', default=None,
    #                         help='Limit to a specific structure size')
    #groupparser.add_argument('--address', type=argparse_utils.int16, action='store', default=None,
    #                         help='Limit to structure similar to the structure pointed at <address>')
    #groupparser.set_defaults(func=groupStructures)

    # not refactored yet
    #parent = subparsers.add_parser('parent', help='Print the parent allocators pointing to the structure located at this address.')
    #parent.add_argument('address', type=argparse_utils.int16, action='store', default=None,
    #                    help='Hex address of the child structure.')
    #parent.set_defaults(func=printParents)

    # not refactored yet
    #graphparser = subparsers.add_parser('graph', help='DISABLED - Show sorted structure instances groups by size and signature in a graph.')
    #graphparser.add_argument('--size', type=int, action='store', default=None,
    #                         help='Limit to a specific structure size')
    #graphparser.add_argument('--address', type=argparse_utils.int16, action='store', default=None,
    #                         help='Limit to structure similar to the structure pointed at <address>')
    #graphparser.set_defaults(func=graph)

    # not refactored yet
    #showparser = subparsers.add_parser('show', help='Show one structure instance.')
    #showparser.add_argument('address', type=argparse_utils.int16, action='store', default=None,
    #                        help='Specify the address of the structure, or of a structure member.')
    #showparser.set_defaults(func=show)

    # FIXME delete ?
    makesig = subparsers.add_parser('makesig', help='Create a simple signature file of the heap - NULL, POINTERS, OTHER VALUES.')
    makesig.set_defaults(func=saveSignatures)

    cleanp = subparsers.add_parser('clean', help='Clean the memory dump from cached info.')
    cleanp.set_defaults(func=clean)

    return rootparser


def main(argv):

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

    level=logging.WARNING
    if opts.debug :
        level=logging.DEBUG
        flog = os.path.normpath('log')
        logging.basicConfig(level=level, filename=flog, filemode='w')
        logging.getLogger('haystack-reverse').setLevel(logging.DEBUG)
        logging.getLogger('signature').setLevel(logging.DEBUG)
        logging.getLogger('reversers').setLevel(logging.DEBUG)
        print ('[+] **** COMPLETE debug log to %s'%(flog))
    else:
        logging.getLogger('haystack-reverse').setLevel(logging.INFO)
        logging.getLogger('signature').setLevel(logging.INFO)
        logging.getLogger('reversers').setLevel(logging.INFO)
        logging.getLogger('ctx').setLevel(logging.INFO)
    sh=logging.StreamHandler(sys.stdout) # 2.6, 2.7 compat
    logging.getLogger('signature').addHandler( sh )
    logging.getLogger('reversers').addHandler( sh )
    logging.getLogger('haystack-reverse').addHandler( sh )

    opts.func(opts)
    return

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


