#ifndef _CRT_SECURE_NO_WARNINGS

#define _CRT_SECURE_NO_WARNINGS

#endif

#include <Windows.h>

#include <hexrays.hpp>

#include <struct.hpp>

#include <bytes.hpp>

#include <kernwin.hpp>

#include <algorithm>

#include <pro.h>

#include <srarea.hpp>

#include <auto.hpp>

#include <funcs.hpp>

#include <expr.hpp>//for VT_LONG

#include <frame.hpp>//for recalc_spd

#include "choosers.h"

#include "structures.h"

#include "negative_cast.h"

#include "helpers.h"

#include "struct_graph.h"

#include "classes_view.h"

#include "classes.h"

#include "new_struct.h"

#include "new_struc_view.h"

#include "ripped.h"

#include "zeroizer.h"

// Hex-Rays API pointer

hexdsp_t *hexdsp = NULL;

extern plugin_t PLUGIN;

static bool inited = false;

bool name_for_struct_crawler(vdui_t &vu, qstring & name, lvar_t ** out_var, ea_t * ea);

bool set_lvar_type(vdui_t * vu, lvar_t * lv, tinfo_t * ts);

bool structs_with_this_size(asize_t size);

/*

idea:

cexpr_t * call;

*/

struct hx_tree_match_node_t

{

hx_tree_match_node_t * x;

hx_tree_match_node_t * y;

hx_tree_match_node_t * z;

cexpr_t ** out;

virtual bool accept(cexpr_t * expr)

{

if (expr)

{

if(x && expr->x)

{

if(!x->accept(expr->x))

return false;

}

if(y && expr->y)

{

if(!y->accept(expr->y))

return false;

}

return true;

}

return false;

};

hx_tree_match_node_t():x(0),y(0),z(0),out(0){};

hx_tree_match_node_t(cexpr_t ** out_){out = out_;}

hx_tree_match_node_t(hx_tree_match_node_t * x_, hx_tree_match_node_t * y_):x(x_),y(y_),z(0),out(0){};

};

template <int goal> struct hx_tree_match_op_t: public hx_tree_match_node_t

{

//hx_tree_match_node_t * z;

//int goal; // cot_xxx

hx_tree_match_op_t(cexpr_t ** out_){out = out_;}

virtual bool accept(cexpr_t * expr)

{

if (expr && expr->op == goal)

{

if(x && expr->x)

{

if(!x->accept(expr->x))

return false;

}

if(y && expr->y)

{

if(!y->accept(expr->y))

return false;

}

if(out)

*out = expr;

return true;

}

return false;

}

};

//--------------------------------------------------------------------------

// Check if the item under the cursor is 'if' or 'else' keyword

// If yes, return pointer to the corresponding ctree item

static cinsn_t *find_if_statement(vdui_t &vu)

{

// 'if' keyword: straightforward check

if ( vu.item.is_citem() )

{

cinsn_t *i = vu.item.i;

if ( i->op == cit_block )

return i;

}

return NULL;

}

static cexpr_t *find_memptr_var(vdui_t &vu, bool * is_call = 0, ea_t * ea = 0 )

{

if ( vu.item.is_citem() )

{

cexpr_t *f;

cexpr_t *e;

e = f = vu.item.e;

if (is_call)

{

const citem_t *p = vu.cfunc->body.find_parent_of(e);

// call => cast => memptr

if (p && p->op == cot_cast)

{

p = vu.cfunc->body.find_parent_of(p);

}

if (p)

{

*is_call = p->op == cot_call;

//msg("tools, is call: %d\n", *is_call);

if (*is_call && ea)

{

*ea = p->ea;

//msg("tools, is call ea: %p\n", *ea);

}

}

}

if ( e->op == cot_memptr || e->op == cot_memref )

{

if ( e->x->op == cot_idx )

e = e->x;

typestring t = e->x->type;

while (t.is_ptr_or_array())

t.remove_ptr_or_array();

if (t.is_struct())

{

return f;

}

if ( (e->x->op == cot_obj) || (e->x->op == cot_var) || (e->x->op == cot_call) || (e->x->op == cot_memptr) || (e->x->op == cot_memref) )

{

return f;

}

}

}

return NULL;

}

static cexpr_t *find_cast(vdui_t &vu)

{

if ( vu.item.is_citem() )

{

cexpr_t *e = vu.item.e;

if ( e->op == cot_cast )

{

return e;

}

}

return NULL;

}

static cexpr_t *find_var(vdui_t &vu)

{

if ( vu.item.is_citem() )

{

cexpr_t *e = vu.item.e;

if ( e->op == cot_var )

{

return e;

}

}

return NULL;

}

bool jumptype(tinfo_t t, int offset, bool virtual_calls = false)

{

while(t.is_ptr_or_array())

t.remove_ptr_or_array();

if(t.is_struct())

{

tid_t id = get_struc_from_typestring( t );

if (id == BADADDR)

{

msg("!get_struc_id\n");

return false;

}

if (virtual_calls)

{

char comment[MAXSTR];

get_struc_cmt(id, true, comment, sizeof(comment));

char * strpos = strstr(comment, "@0x");

if (strpos)

{

ea_t ea;

if ( my_atoea( strpos + 1, &ea) )

{

ea_t fnc = get_long(ea + offset);

flags_t flags = getFlags(fnc);

if (isFunc(flags))

{

open_pseudocode(fnc, -1);

}

else

{

if (isCode(flags))

{

jumpto(fnc);

}

else

{

jumpto(ea + offset);

}

}

return true;

}

}

}

open_structs_window(id, offset);

return true;

}

else

return false;

}

static bool has_VT(void * ud)

{

vdui_t &vu = *(vdui_t *)ud;

bool is_call = false;

cexpr_t *e = find_memptr_var(vu, &is_call);

if (!e)

return false;

int offset = e->m;

cexpr_t * var = e->x;

if(var->type.empty())

return false;

tid_t id = get_struc_from_typestring(var->type);

if (id == BADNODE)

return false;

char comment[MAXSTR];

get_struc_cmt(id, true, comment, sizeof(comment));

char * strpos = strstr(comment, "@0x");

if (!strpos)

return false;

return true;

}

static bool is_virtual_call(void * ud)

{

vdui_t &vu = *(vdui_t *)ud;

bool is_call = false;

cexpr_t *e = find_memptr_var(vu, &is_call);

if (!e)

return false;

return is_call;

}

static bool idaapi jump_to_VT(void * ud)

{

vdui_t &vu = *(vdui_t *)ud;

bool is_call = false;

cexpr_t *e = find_memptr_var(vu, &is_call);

if (e)

{

int offset = e->m;

if (e->op == cot_idx)

e = e->x;

cexpr_t * var = e->x;

if(var->type.empty())

{

return false;

}

return jumptype(var->type, offset, true);

}

e = find_cast(vu);

if (e)

{

if(e->type.empty())

{

return false;

}

return jumptype(e->type, 0);

}

return false;

}

//resets all calls to ea

static void reset_alt_calls(ea_t ea)

{

static const char * nname;

if ( ph.id == PLFM_MIPS )

nname = "$ mips";

else if ( ph.id == PLFM_ARM )

nname = " $arm";

else

nname = "$ vmm functions";

netnode n(nname);

for(nodeidx_t idx= n.alt1st(); idx != BADNODE; idx = n.altnxt(idx))

{

if(n.altval(idx)-1 == ea)

{

n.altdel(idx);

noUsed(idx); // reanalyze the current instruction

recalc_spd(idx);

}

//n.altdel_all(atag);

}

}

static const char reset_alt_idc_args[] = { VT_LONG, 0 };

static error_t idaapi reset_alt_idc(idc_value_t *argv, idc_value_t *res)

{

msg("[Hexrays-Tools] reset_alt_idc is called with arg0=%x\n", argv[0].num);

reset_alt_calls(argv[0].num);

return eOk;

}

static const char structs_of_size_args[] = { VT_LONG, 0 };

static error_t idaapi structs_of_size(idc_value_t *argv, idc_value_t *res)

{

msg("[Hexrays-Tools] structs_of_size is called with arg0=%x\n", argv[0].num);

asize_t sz = argv[0].num;

structs_with_this_size(sz);

return eOk;

}

void register_idc_functions()

{

set_idc_func_ex("reset_alt_calls", reset_alt_idc, reset_alt_idc_args, 0);

set_idc_func_ex("structs_of_size", structs_of_size, structs_of_size_args, 0);

//set_idc_func_ex("get_member_offset", idc_get_member_by_fullname, idc_get_member_by_fullname_args, 0);

}

void unregister_idc_functions()

{

set_idc_func_ex("reset_alt_calls", NULL, NULL, 0);

set_idc_func_ex("structs_of_size", NULL, NULL, 0);

//set_idc_func_ex("get_member_offset", NULL, NULL, 0);

}

static void set_call(ea_t from, ea_t to)

{

static const char * nname;

if ( ph.id == PLFM_MIPS )

nname = "$ mips";

else if ( ph.id == PLFM_ARM )

nname = " $arm";

else

nname = "$ vmm functions";

netnode n(nname);

ea_t ea = from; // get current address

if ( !isCode(get_flags_novalue(ea)) ) return; // not an instruction

ea_t callee = n.altval(ea)-1; // get the callee address from the database

// remove thumb bit for arm

if ( ph.id == PLFM_ARM )

callee &= ~1;

//char buf[MAXSTR];

//qsnprintf(buf, sizeof(buf), form, help);

//if ( AskUsingForm_c(buf, &callee) )

callee = to;

{

if ( callee == BADADDR )

{

n.altdel(ea);

}

else

{

#define T 20

if ( ph.id == PLFM_ARM && (callee & 1) == 0 )

{

// if we're calling a thumb function, set bit 0

sel_t tbit = getSR(callee, T);

if ( tbit != 0 && tbit != BADSEL )

callee |= 1;

}

#undef T

n.altset(ea, callee+1); // save the new address

}

noUsed(ea); // reanalyze the current instruction

}

}

static bool idaapi add_VT_cref(void * ud)

{

vdui_t &vu = *(vdui_t *)ud;

bool is_call = false;

vu.get_current_item(USE_KEYBOARD);

ea_t caller_ea;

cexpr_t *e = find_memptr_var(vu, &is_call, &caller_ea);

if (e)

{

int offset = e->m;

cexpr_t * var = e->x;

if(var->type.empty())

{

return false;

}

tid_t id = get_struc_from_typestring( var->type );

if (id==BADNODE)

return false;

char comment[MAXSTR];

get_struc_cmt(id, true, comment, sizeof(comment));

char * strpos = strstr(comment, "@0x");

if ( strpos )

{

ea_t ea;

if ( my_atoea( strpos + 1, &ea) )

{

ea_t fnc = get_long(ea + offset);

flags_t flags = getFlags(fnc);

if (isFunc(flags))

{

msg("[Hexrays-Tools] adding cref from %p to %p\n", caller_ea, fnc);

//add_cref(caller_ea, fnc, fl_CF);

set_call(caller_ea, fnc);

//n.altset(ea, callee+1); // save the new address

//set_cmt();

//char comment[MAXSTR];

comment[0]=0;

char new_comment[MAXSTR];

new_comment[0]=0;

get_cmt(caller_ea, false, comment, sizeof(comment));

qsnprintf( new_comment, MAXSTR, "%s\ncalls %p", comment, fnc );

set_cmt(caller_ea, new_comment, false);

return true;

}

else

{

msg("!isFunc\n");

}

}

else

{

msg("!atoea %s \n", strpos + 1);

}

}

else

{

msg("!strpos\n");

}

}

return false;

}

static bool idaapi jump_to_struct(void * ud)

{

vdui_t &vu = *(vdui_t *)ud;

bool is_call = false;

vu.get_current_item(USE_KEYBOARD);

cexpr_t *e = find_memptr_var(vu, &is_call);

if (e)

{

int offset = e->m;

if (e->op == cot_idx)

e = e->x;

cexpr_t * var = e->x;

if( var->type.empty() )

{

return false;

}

return jumptype(var->type, offset, is_call);

}

e = find_cast(vu);

if (e)

{

if(e->type.empty())

{

return false;

}

return jumptype(e->type, 0);

}

e = find_var(vu);

if (e)

{

if(e->type.empty())

{

return false;

}

return jumptype(e->type, 0);

}

return false;

}

bool show_VT_from_tid(tid_t id)

{

if (id == BADADDR)

{

//msg("!get_struc_id\n");

return false;

}

{

char comment[MAXSTR];

get_struc_cmt(id, true, comment, sizeof(comment));

char * strpos = strstr(comment, "@0x");

if (strpos)

{

ea_t ea;

if ( my_atoea( strpos + 1, &ea) )

{

function_list fl;

{

while (1)

{

ea_t fncea = get_long(ea);

flags_t fnc_flags = getFlags(fncea);

if ( isFunc(fnc_flags) )

{

fl.functions.push_back(fncea);

}

ea = next_head(ea, BADADDR);

flags_t ea_flags = getFlags(ea);

if (has_any_name(ea_flags))

break;

}

}

fl.choose("[Hexrays-Tools] virtual table");

fl.open_pseudocode();

return true;

}

}

struc_t * struc = get_struc(id);

member_t * member = get_member(struc, 0);

if(member)

{

tinfo_t t;

if( get_member_type(member, &t) )

{

tid_t child = get_struc_from_typestring(t);

if(child!=BADNODE)

{

return show_VT_from_tid(child);

}

}

}

}

return false;

}

static bool idaapi show_VT(void * ud)

{

vdui_t &vu = *(vdui_t *)ud;

cexpr_t *e = find_memptr_var(vu, 0);

if (e)

{

int offset = e->m;

if (e->op == cot_idx)

e = e->x;

cexpr_t * var = e->x;

tinfo_t t = var->type;

if(t.empty())

{

return false;

}

while(t.is_ptr_or_array())

t.remove_ptr_or_array();

if(t.is_struct())

{

tid_t id = get_struc_from_typestring( t );

show_VT_from_tid(id);

//open_structs_window(id, offset);

}

//return jumptype(var->type, offset, true);

}

return false;

}

struct offset_locator_t : public ctree_parentee_t

{

lvars_t * lvars;

lvar_t * var_hled;

std::set<int> idxs;

std::set<int> offsets;

std::map<int, typestring> types;

bool is_our(int idx)

{

return idxs.find(idx)!=idxs.end();

}

offset_locator_t(lvars_t * lvars_, lvar_t * lvar_) : lvars(lvars_)

{

int idx = get_idx_of(lvars_, lvar_);

//deprecated index of

/*int c=0;

for(lvars_t::iterator i = lvars_->begin(); i!=lvars_->end(); i++)

{

if (&*i == lvar_)

{

idx = c;

break;

}

c++;

}*/

//QASSERT(0, idx!=-1);

idxs.insert(idx);

}

int idaapi visit_expr(cexpr_t * e)

{

if ( e->op == cot_asg )

{

if (e->x->op == cot_var && e->y->op == cot_var)

{

if(idxs.find(e->y->v.idx)!=idxs.end() && idxs.find(e->x->v.idx)==idxs.end())

{

idxs.insert(e->x->v.idx);

}

}

return 0;

}

if ( e->op == cot_memptr )

{

if ( e->x->op == cot_var )

{

cexpr_t * var = e->x;

if(idxs.find(var->v.idx)!=idxs.end())

{

msg( "var: %s.%d\n", (*lvars)[var->v.idx].name.c_str(), e->m );

offsets.insert(e->m);

}

}

}

{

bool is_lvar = (e->op == cot_var && is_our(e->v.idx));

bool is_glob_obj_ptr = (e->op == cot_obj && is_our(e->obj_ea));

int index = 0;

if (is_lvar)

{

index = e->v.idx;

}

/*if (is_glob_obj_ptr)

{

index = e->obj_ea;

}*/

if ( is_lvar /*|| is_glob_obj_ptr*/ )

{ // found our var. are we inside a pointer expression?

uint64 delta1 = 0;

uint64 delta2 = 0;

bool delta_defined1 = true;

bool delta_defined2 = false;

bool is_array = false;

int i = parents.size() - 1;

if ( parents[i]->op == cot_add ) // possible delta

{

cexpr_t *d = ((cexpr_t*)parents[i])->theother(e);

delta_defined1 = d->get_const_value(&delta1);

i--;

//possible array indexing

if ( parents[i]->op == cot_add /*|| parents[i]->op == cot_sub*/ )

{

cexpr_t *d = ((cexpr_t*)parents[i])->theother((cexpr_t*)parents[i+1]);

delta_defined2 = d->get_const_value(&delta2);

is_array = true;

i--;

}

}

uint64 delta;

if (delta_defined1)

delta = delta1;

if (delta_defined2)

delta = delta2;

typestring t;

if ( delta_defined1 || delta_defined2 )

{

// skip casts

while ( parents[i]->op == cot_cast )

{

t = ((cexpr_t*)parents[i])->type;

i--;

}

//is_ptr = type.type.is_ptr();

t = remove_pointer(t);

offsets.insert((int)delta);

if (t.size()>0)

{

types[(int)delta] = t;

char buff[MAXSTR];

t.print(buff, sizeof(buff));

msg("[Hexrays-Tools] new var: %s.%d [%s][array: %d]\n", (*lvars)[index].name.c_str(), (int)delta, buff, is_array);

}

else

msg("[Hexrays-Tools] new var: %s.%d[array:%d]\n", (*lvars)[index].name.c_str(), (int)delta, is_array);

}

}

}

if ( e->op == cot_ptr )

{

if(e->x->op == cot_cast)

{

cexpr_t * cast = e->x;

if( cast->x->op == cot_add )

{

cexpr_t * add = cast->x;

if ( add->x->op == cot_var )

{

cexpr_t * var = add->x;

if (add->y->op == cot_num)

{

cexpr_t * num = add->y;

if(idxs.find(var->v.idx)!=idxs.end() )

{

msg("[Hexrays-Tools] var: %s.%d\n", (*lvars)[var->v.idx].name.c_str(), (int)num->numval());

offsets.insert((int)num->numval());

types[(int)num->numval()] = cast->type;

}

}

else

{

if(&(*lvars)[var->v.idx] == var_hled)

{

msg("[Hexrays-Tools] xvar: %s\n", (*lvars)[var->v.idx].name.c_str());

}

}

}

}

else

if( cast->x->op == cot_var )

{

cexpr_t * var = cast->x;

if(idxs.find(var->v.idx)!=idxs.end() )

{

msg("cast var: %s\n", (*lvars)[var->v.idx].name.c_str());

offsets.insert(0);

types[0] = cast->type;

}

}

}

}

return 0;

}

int idaapi visit_insn(cinsn_t *i)

{

return 0; // continue enumeration

}

};

uint32 idaapi sizer(void *obj)

{

intvec_t * vec = (intvec_t*)obj;

return vec->size() + 1;

}

void idaapi get_type_line(void *obj,uint32 n,char * const *arrptr)

{

intvec_t * vec = (intvec_t*)obj;

if (n==0)

{

qstpncpy( arrptr[0], "type", MAXSTR );

qstpncpy( arrptr[1], "sz", MAXSTR );

qstpncpy( arrptr[2], "sz hx", MAXSTR );

}

else if (n==1)

{

qstpncpy( arrptr[0], "<create new>", MAXSTR );

arrptr[1][0] = 0;

arrptr[2][0] = 0;

}

else

{

tid_t id = (*vec)[n-2];

asize_t size = get_struc_size(id);

qstring tmpname;

get_struc_name(&tmpname, id);

qsnprintf(arrptr[0], MAXSTR, "%s", tmpname.c_str());

qsnprintf(arrptr[1], MAXSTR, "%d", size);

qsnprintf(arrptr[2], MAXSTR, "0x%08x", size);

}

}

/*

void idaapi enter(void *obj, uint32 n)

{

intvec_t * idcka = (intvec_t *)obj;

uint32 choosed = n;

if(n>1)

{

open_structs_window((*idcka)[n], 0);

}

}

*/

bool idaapi extract_substruct_callback(void *ud)

{

//Controls::TCustomControl * cc = (Controls::TCustomControl *)ud;

Controls::TCustomControl * cc = get_current_viewer();

if(!cc)

return false;

twinpos_t b,e;

if (!readsel2(cc, &b, &e))

return false;

structplace_t * sb = (structplace_t *)b.at;

structplace_t * se = (structplace_t *)e.at;

msg("%d, %d\n", sb->offset, se->offset);

if (sb->idx != se->idx)

return false;

if ( extract_substruct(sb->idx, sb->offset, se->offset) )

{

//this should be called automatically

//refresh_custom_viewer(cc);

return true;

}

return false;

}

bool idaapi which_struct_matches_here_callback(void *ud)

{

//Controls::TCustomControl * cc = (Controls::TCustomControl *)ud;

Controls::TCustomControl * cc = get_current_viewer();

if(!cc)

return false;

twinpos_t b,e;

if (!readsel2(cc, &b, &e))

return false;

structplace_t * sb = (structplace_t *)b.at;

structplace_t * se = (structplace_t *)e.at;

if (sb->idx != se->idx)

return false;

if ( which_struct_matches_here(sb->idx, sb->offset, se->offset) )

{

//we dont need refresh after this (I guess)

return false;//true;

}

return false;

}

tid_t create_VT_struc(ea_t VT_ea, char * name, uval_t idx = BADADDR, unsigned int * vt_len = NULL )

{

char name_[MAXNAMELEN];

get_name(BADADDR, VT_ea, name_, sizeof(name_));

if(qstrstr(name_,"_VT_"))

{

char * ch = strchr(name_, 0);

*(--ch) = 0;

name = name_;

}

else

if( !name )

{

qsnprintf(name_, MAXNAMELEN, "VT_%p", VT_ea);

name = name_;

}

{//TODO: do this better

ea_t fncea = get_long(VT_ea);

flags_t fnc_flags = getFlags(fncea);

if (!isFunc(fnc_flags))

return BADNODE;

}

tid_t newid = get_struc_id(name);

if (newid == BADADDR)

newid = add_struc(idx, name);

if (newid == BADADDR)

return BADNODE;

struc_t * newstruc = get_struc(newid);

if (!newstruc)

return BADNODE;

char comment[MAXSTR];

qsnprintf(comment, sizeof(comment), "@0x%p", VT_ea);

set_struc_cmt(newid, comment, true);

ea_t ea = VT_ea;

int offset = 0;

int len = 0;

while (1)

{

offset = ea - VT_ea;

char funcname[MAXSTR];

memset(funcname, 0, sizeof(funcname));

ea_t fncea = get_long(ea);

//there are no holes in vftables (I hope)

if(fncea == 0)

break;

//there are also no false pointers in vftables

if (!isEnabled(fncea))

break;

flags_t fnc_flags = getFlags(fncea);

char * fncname = NULL;

if (isFunc(fnc_flags) /*&& has_user_name(fnc_flags)*/ )

{

//if ( get_short_name(BADADDR, fncea, funcname, sizeof(funcname)) )

if(get_func_name(fncea, funcname, MAXSTR))

fncname = funcname;

}

add_struc_member(newstruc, NULL, offset, dwrdflag(), NULL, 4);

len++;

if (fncname)

{

set_member_cmt(get_member(newstruc, offset), fncname, true);

if(!set_member_name( newstruc, offset, fncname ))

{

get_name(NULL, fncea, funcname, sizeof(funcname));

set_member_name( newstruc, offset, fncname );

}

}

//ea = next_head(ea, BADADDR);

//TODO: this is only for 32bit processors but so is hxrs

ea += 4;

flags_t ea_flags = getFlags(ea);

if (has_any_name(ea_flags))

break;

}

if(vt_len)

*vt_len = len;

return newid;

}

static void set_vt_type(struc_t *struc, char name_of_vt_struct[MAXSTR])

{

typestring type = create_numbered_type_from_name(name_of_vt_struct);//create_typedef(name);

type = make_pointer(type);

member_t * member = get_member(struc, 0);

if(!member)

{

if(add_struc_member(struc, NULL, 0, dwrdflag(), NULL, 4) != 0)

return;

member = get_member(struc, 0);

}

set_member_tinfo(idati, struc, member, 0, type.c_str(), /*fields.c_str()*/NULL, SET_MEMTI_MAY_DESTROY);

set_member_name(struc, 0, "VT");

}

static bool create_VT(uval_t idx)

{

tid_t id = get_struc_by_idx( idx );

if (is_union(id))

{

msg("[Hexrays-Tools] union!\n");

return false;

}

struc_t * struc = get_struc(id);

if(!struc)

{

msg("[Hexrays-Tools] !struc\n");

return false;

}