#ifndef __TRACE__
#define __TRACE__

#include        "bitmap.h"
       
typedef struct  _x86regs{
        ULONG   reg_eip;
        ULONG   reg_eflags;
        ULONG   reg_edi;
        ULONG   reg_esi;
        ULONG   reg_ebp;
        ULONG   reg_dummy;
        ULONG   reg_ebx;
        ULONG   reg_edx;
        ULONG   reg_ecx;
        ULONG   reg_eax;
        ULONG   reg_esp;
}x86regs, *px86regs;

typedef struct _stack_syscall_ret{
        unsigned long stack_base;               //base read from __readfsdword(4)
        unsigned long stack_index;              //stack index (__readfsdword - esp) might be even wrong
                                                //if stack is changed to be outside of __readfsdword range
                                                //but it will give me correct stack address to work with
        unsigned long stack_ret;                //ret address which was taken from there...
        unsigned long stack_syscall;            //number of syscall to which this matches...
}stack_syscall_ret, *pstack_syscall_ret;

typedef struct _context{
        unsigned long   reg_eax;        
        unsigned long   reg_ecx;        
        unsigned long   reg_edx;        
        unsigned long   reg_ebx;        
        unsigned long   reg_esp;        
        unsigned long   reg_ebp;        
        unsigned long   reg_esi;        
        unsigned long   reg_edi;        
        unsigned long   reg_eflags;     
        unsigned long   reg_eip;       
        unsigned long   current_bbl;     
        unsigned long   syscall_number;         //syscall number
        unsigned long   syscall_arguments;      //on x32 point to stack which hold arguments!!!
                          
        unsigned long   exception_eip;
        unsigned long   exception_code;
        pstack_syscall_ret psyscall_ret;
        unsigned long   syscall_ret_depth;
        
        ULONG_PTR       thread_stack;
        ULONG_PTR       thread_stack_base;
        ULONG_PTR       teb;
        ULONG_PTR       thread_real_esp;
        
        ULONG_PTR       real_thread_stack_base;
        ULONG_PTR       real_thread_stack_limit;
        
        unsigned long   tf_set;
        unsigned long   tf_inject;
             
}context, *pcontext;

typedef struct _BBL_STRUCT{
        ULONG_PTR       bbl_id;                 //not really needed...
        ULONG_PTR       bbl_start;              //start of basic block (original EIP)
        ULONG_PTR       bbl_end;                //end of basic block   (original EIP)
        ULONG_PTR       bbl_buffer;             //buffer where rebuilt instructions are stored...
        ULONG_PTR       bbl_inst_count;
        unsigned long   bbl_refcount;           //if it's in use, we can't free it in case there is write
                                                //operation...
        unsigned long   bbl_crc32;              //check if block has changed if it's write block!!!
        unsigned long   bbl_protection;         //always extracted from vmmaps...
        unsigned long   bbl_syscall_int2e;
}BBL_STRUCT, *PBBL_STRUCT;

typedef struct _bbl_struct{
        LIST            Next;
        ULONG_PTR       bbl_eip_redirection;
        ULONG_PTR       bbl_id;
        ULONG_PTR       bbl_start;
        ULONG_PTR       bbl_end;
        ULONG_PTR       bbl_buffer;
        ULONG_PTR       bbl_alternate_buffer;   //buffer which has exact copy of bbl as is in memory
        ULONG_PTR       bbl_inst_count;
        unsigned long   bbl_refcount;
        unsigned long   bbl_syscall_int2e;
        unsigned long   bbl_syscall_int2b;
        bitmap_struct   bitmap;
}bbl_struct, *pbbl_struct;

void    traceProcessNextBuffer(__in context *pctx);

VOID traceSysCallInstrument(__in px86regs pregs);
VOID traceSysCallPreInstrument(__in px86regs pregs);
VOID    *traceBuildBbl(__in unsigned char *inst_ptr, __in unsigned long b_one_instruction);
VOID    traceInitThread(__in px86regs pregs);
VOID    tracePostSysenter(__in context *pctx);
VOID    tracePostInt2e(__in context *pctx);
extern  ULONG_PTR g_traceProcessNextBuffer;
VOID    traceFreeBbl(__in BBL_STRUCT *pbbl);

unsigned long tracePreSyscall(__in context *pctx);
VOID    tracePostSyscall(__in context *pctx);

VOID    traceInjectException(__in context *pctx, __in DWORD ExceptionCode, __in DWORD ExceptionAddress, __in DWORD b_write);
VOID    traceFreeBblRange(__in PVOID lpBaseAddress, __in ULONG dwSize);
typedef VOID (WINAPI *EXECBBL)(VOID);

#endif