/* Custom viewer sample plugin.

* Copyright (c) 2007 by Ilfak Guilfanov, ig@hexblog.com

* Feel free to do whatever you want with this code.

*

* This sample plugin demonstates how to create and manipulate a simple

* custom viewer in IDA Pro v5.1

*

* Custom viewers allow you to create a view which displays colored lines.

* These colored lines are dynamically created by callback functions.

*

* Custom viewers are used in IDA Pro itself to display

* the disassembly listng, structure, and enumeration windows.

*

* This sample plugin just displays several sample lines on the screen.

* It displays a hint with the current line number.

* The right-click menu contains one sample command.

* It reacts to one hotkey.

*

* This plugin uses the simpleline_place_t class for the locations.

* Custom viewers can use any decendant of the place_t class.

* The place_t is responsible for supplying data to the viewer.

*/

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

#ifdef __NT__

#define WIN32_LEAN_AND_MEAN

#include <windows.h>

#else

#define VK_ESCAPE 27

#endif

#include <hexrays.hpp>

#include <ida.hpp>

#include <idp.hpp>

#include <loader.hpp>

#include <kernwin.hpp>

#include <struct.hpp>

#include "classplace.h"

#include "classes.h"

#include "helpers.h"

// Structure to keep all information about the our sample view

struct sample_info_t

{

TForm *form;

TCustomControl *cv;

strvec_t sv;

sample_info_t(TForm *f) : form(f), cv(NULL) {}

};

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

// get the word under the (keyboard or mouse) cursor

static bool get_current_word(TCustomControl *v, bool mouse, qstring &word)

{

// query the cursor position

int x, y;

if ( get_custom_viewer_place(v, mouse, &x, &y) == NULL )

return false;

// query the line at the cursor

char buf[MAXSTR];

const char *line = get_custom_viewer_curline(v, mouse);

tag_remove(line, buf, sizeof(buf));

if ( x >= (int)strlen(buf) )

return false;

// find the beginning of the word

char *ptr = buf + x;

while ( ptr > buf && !isspace(ptr[-1]) )

ptr--;

// find the end of the word

char *begin = ptr;

ptr = buf + x;

while ( !isspace(*ptr) && *ptr != '\0' )

ptr++;

word = qstring(begin, ptr-begin);

return true;

}

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

// Sample callback function for the right-click menu

static bool idaapi change_fnc_name_cb(void *ud)

{

sample_info_t *si = (sample_info_t *)ud;

/* qstring word;

if ( !get_current_word(si->cv, false, word) )

return false;

info("The current word is: %s", word.c_str());

*/

class_place_t * place;

place = (class_place_t *)get_custom_viewer_place(si->cv, false, NULL, NULL);

if(!place)

return false;

tid_t tid = get_class_by_idx(place->idx);

if(tid==BADNODE)

return false;

class_t * clas = get_class(tid);

if(!clas)

return false;

if (place->subsection >= clas->functions_ea.size())

return false;

ea_t ea = clas->functions_ea[place->subsection];

char name[MAXNAMESIZE];

get_true_name(BADADDR, ea, name, MAXNAMESIZE);

char * answer;

while(answer = askstr(HIST_TYPE, name, "[Hexrays-Tools] Enter new function name"))

{

qstrncpy(name, answer, MAXNAMESIZE);

if ( set_name(ea, name) )

break;

}

return true;

}

static bool set_first_parameter_for_one_place(class_place_t *place)

{

if (place->section != 2)

return false;

tid_t tid = get_class_by_idx(place->idx);

if(tid==BADNODE)

return false;

class_t * clas = get_class(tid);

if(!clas)

return false;

if (place->subsection >= clas->functions_ea.size())

return false;

ea_t ea = clas->functions_ea[place->subsection];

func_t * func = get_func(ea);

qtype type;

qtype fields;

if(!get_tinfo(ea, &type, &fields))

{

if(guess_func_tinfo(func,&type, &fields)!=GUESS_FUNC_OK)

return false;

}

func_type_data_t fti;

if(!build_funcarg_info(idati, type.c_str(), fields.c_str(), &fti, 0))

return false;

char class_name[MAXNAMELEN];

get_class_name(tid, class_name, MAXNAMELEN);

tinfo_t ptr = make_pointer(create_numbered_type_from_name(class_name));//make_pointer(create_typedef(class_name));

if (fti.size())

{

fti[0].type = ptr;

}

fields.clear();

type.clear();

build_func_type(&type, &fields, fti);

if ( !apply_tinfo(idati, ea, type.c_str(), fields.c_str(), 1) )

return false;

// vu.refresh_view(true);

return true;

}

static bool idaapi set_first_parameter_cb(void *ud)

{

sample_info_t *si = (sample_info_t *)ud;

twinpos_t b,e;

if (!readsel2(si->cv, &b, &e))

{

//just one

class_place_t * place;

place = (class_place_t *)get_custom_viewer_place(si->cv, false, NULL, NULL);

if(!place)

return false;

return set_first_parameter_for_one_place(place);

};

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

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

do

{

if (sb->section == 2)

{

set_first_parameter_for_one_place(sb);

}

if (!sb->next(0))

return false;

} while(sb->compare(se)!=1);

return true;

}

static bool idaapi changetype_cb(void *ud)

{

sample_info_t *si = (sample_info_t *)ud;

/* qstring word;

if ( !get_current_word(si->cv, false, word) )

return false;

info("The current word is: %s", word.c_str());

*/

class_place_t * place;

place = (class_place_t *)get_custom_viewer_place(si->cv, false, NULL, NULL);

if(!place)

return false;

tid_t tid = get_class_by_idx(place->idx);

if(tid==BADNODE)

return false;

class_t * clas = get_class(tid);

if(!clas)

return false;

if (place->subsection >= clas->functions_ea.size())

return false;

ea_t ea = clas->functions_ea[place->subsection];

char name[MAXSTR];

get_true_name(BADADDR, ea, name, MAXSTR);

qtype type;

qtype fields;

if (!get_tinfo(ea, &type, &fields))

{

if (!guess_func_tinfo(get_func(ea), &type, &fields))

goto pokracuj;

}

print_type_to_one_line(name, MAXSTR, idati, type.c_str(), name, 0, fields.c_str(), 0);

pokracuj:

char declaration[MAXSTR];

char * answer;

while(answer = askstr(HIST_TYPE, name, "[Hexrays-Tools] Enter new function type"))

{

qstrncpy(declaration, answer, MAXSTR);

//notice:

//there is always some 0 in memory

//so no need to check non-nullness of strchr result

if( *(strchr(declaration, 0)-1)!= ';')

qstrncat(declaration, ";", MAXSTR);

if (parse_decl(idati, declaration, NULL, &type, &fields, PT_VAR))

{

if (apply_tinfo(idati, ea, type.c_str(), fields.c_str(), 1))

{

return true;

}

}

}

return false;

}

extern void show_function_with_this_struct_tid_as_parameter(tid_t goal);

static bool idaapi scan_cb(void *ud)

{

sample_info_t *si = (sample_info_t *)ud;

class_place_t * place;

place = (class_place_t *)get_custom_viewer_place(si->cv, false, NULL, NULL);

if(!place)

return false;

tid_t tid = get_class_by_idx(place->idx);

if(tid==BADNODE)

return false;

show_function_with_this_struct_tid_as_parameter(tid);

return true;

}

bool is_at_function(Controls::TCustomControl * cv, class_place_t ** place_out = 0)

{

class_place_t * place;

place = (class_place_t *)get_custom_viewer_place(cv, false, NULL, NULL);

if(!place)

return false;

if (place_out)

*place_out = place;

return place->section == 2;

}

bool is_at_vftable(Controls::TCustomControl * cv)

{

class_place_t * place;

place = (class_place_t *)get_custom_viewer_place(cv, false, NULL, NULL);

if(!place)

return false;

return place->section == 1;

}

bool is_at_classname(Controls::TCustomControl * cv)

{

class_place_t * place;

place = (class_place_t *)get_custom_viewer_place(cv, false, NULL, NULL);

if(!place)

return false;

return place->section == 0;

}

static bool class_add_constructor(class_t * clas)

{

if(!clas)

return false;

if(clas->virt_table_ea == BADADDR)

return false;

int cnt = 0;

for(ea_t ea = get_first_dref_to(clas->virt_table_ea); ea!=BADADDR; ea = get_next_dref_to(clas->virt_table_ea, ea))

{

func_t * fnc = get_func(ea);

if(fnc)

{

clas->functions_ea.add_unique(fnc->startEA);

cnt++;

}

}

return cnt>0;

}

static bool idaapi add_constructor_destuctor_cb(void *ud)

{

sample_info_t * si = (sample_info_t *)ud;

class_place_t * place;

place = (class_place_t *)get_custom_viewer_place(si->cv, false, NULL, NULL);

if(!place)

return false;

tid_t tid = get_class_by_idx(place->idx);

if(tid==BADNODE)

return false;

class_t * clas = get_class(tid);

class_add_constructor(clas);

save_class(clas);

return true;

}

static bool idaapi add_from_structs_cb(void *ud)

{

for( uval_t idx = get_first_struc_idx(); idx!=BADNODE; idx=get_next_struc_idx(idx) )

{

tid_t id = get_struc_by_idx(idx);

struc_t * struc = get_struc(id);

if (!struc->memqty)

continue;

char name[MAXNAMESIZE];

if(!get_member_name(struc->members[0].id, name, MAXNAMESIZE))

continue;

if((strnicmp(name, "VT", 2)==0) || (stristr(name, "vt_")!=0))

{

//msg("adding class %p\n", id);

add_class(id);

tinfo_t t;

if( get_member_type(&struc->members[0], &t) )

{

tid_t child = get_struc_from_typestring(t);

struc_t * vt = get_struc(child);

if(vt)

{

char comment[MAXSTR];

get_struc_cmt(vt->id, true, comment, sizeof(comment));

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

if ( strpos )

{

ea_t ea;

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

{

class_t * clas = get_class(id);

if(!clas)

continue;

//TODO: class_set_vft();

clas->virt_table_ea = ea;

//msg("adding vt %p to class %p\n", ea, id);

save_class(clas);

{

while (1)

{

ea_t fncea = get_long(ea);

flags_t fnc_flags = getFlags(fncea);

if ( isFunc(fnc_flags) )

{

clas->functions_ea.add_unique(fncea);

func_t * func = get_func(fncea);

if (func && (func->flags & FUNC_THUNK))

{

ea_t target;

ea_t target2;

if ((target2 = calc_thunk_func_target(func, &target))!=BADADDR)

{

clas->functions_ea.add_unique(target2);

if (target!=BADADDR)

clas->functions_ea.add_unique(target);

}

}

}

ea += 4;

//this doesn't work if vtable is already declared as structure

//ea = next_head(ea, BADADDR);

flags_t ea_flags = getFlags(ea);

if (has_any_name(ea_flags))

break;

}

class_add_constructor(clas);

save_class(clas);

}

}

}

}

}

}

}

return true;

}

static bool jumping(sample_info_t *si, int shift)

{

class_place_t * place;

place = (class_place_t *)get_custom_viewer_place(si->cv, false, NULL, NULL);

switch (place->section)

{

case 0:

{

tid_t id = get_class_by_idx( place->idx );

open_structs_window(id);

msg("[Hexrays-Tools] Enter has been pressed");

}

break;

case 1:

{

//TODO

/*tid_t id = get_class_by_idx( place->idx );

open_structs_window(id);

msg("Enter has been pressed");

*/

}

break;

case 2:

{

tid_t id = get_class_by_idx( place->idx );

if (id==BADADDR)

return false;

class_t * clas = get_class(id);

if (!clas)

return false;

if (place->subsection >= clas->functions_ea.size())

return false;

if(shift)

{

jumpto(clas->functions_ea[place->subsection]);

}

else

{

open_pseudocode(clas->functions_ea[place->subsection], -1);

}

}

break;

}

return true;

}

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

// The user double clicked

static bool idaapi cb_dblclick(TCustomControl *cv, int shift, void *ud)

{

sample_info_t *si = (sample_info_t *)ud;

return jumping(si, shift);

}

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

// Keyboard callback

static bool idaapi ct_keyboard(TCustomControl * /*v*/, int key, int shift, void *ud)

{

//if ( shift == 0 )

{

sample_info_t *si = (sample_info_t *)ud;

switch ( key )

{

case 'N':

if (is_at_function(si->cv))

{

change_fnc_name_cb(si);

return true;

}

msg("[Hexrays-Tools] The hotkey 'N' has been pressed\n");

return false;

case 'Y':

if (is_at_function(si->cv))

{

changetype_cb(si);

return true;

}

return false;

case VK_ESCAPE:

close_tform(si->form, FORM_SAVE);

return true;

case VK_DELETE:

{

class_place_t * class_pl;

if (is_at_function(si->cv, &class_pl))

{

tid_t tid = get_class_by_idx(class_pl->idx);

class_t * clas = get_class(tid);

eavec_t::iterator i = clas->functions_ea.begin();

i += class_pl->subsection;

clas->functions_ea.erase(i);

save_class(clas);

refresh_custom_viewer(si->cv);

return true;

}

}

break;

case VK_RETURN:

return jumping(si, shift);

}

}

return false;

}

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

#define add_popup add_custom_viewer_popup_item

bool is_something_selected(TCustomControl *v)

{

twinpos_t b,e;

return readsel2(v, &b, &e);

}

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

static void idaapi ct_popup(TCustomControl *v, void *ud)

{

set_custom_viewer_popup_menu(v, NULL);

// Create right-click menu on the fly

add_popup(v, "[Hexrays-Tools] Import from Structures", "", add_from_structs_cb, ud);

if ( is_at_function(v) )

{

if (!is_something_selected(v))

{

add_popup(v, "[Hexrays-Tools] Rename function", "N", change_fnc_name_cb, ud);

add_popup(v, "[Hexrays-Tools] Retype function", "Y", changetype_cb, ud);

}

add_popup(v, "[Hexrays-Tools] Set first argument type", "", set_first_parameter_cb, ud);

}

add_popup(v, "[Hexrays-Tools] Scan for class functions", "S", scan_cb, ud);

add_popup(v, "[Hexrays-Tools] add constructor / destructor", "", add_constructor_destuctor_cb, ud);

}

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

// This callback will be called each time the keyboard cursor position

// is changed

static void idaapi ct_curpos(TCustomControl *v, void *)

{

qstring word;

/*

if ( get_current_word(v, false, word) )

msg("Current word is: %s\n", word.c_str());

*/

}

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

static int idaapi ui_callback(void *ud, int code, va_list va)

{

sample_info_t *si = (sample_info_t *)ud;

switch ( code )

{

// how to implement a simple hint callback

case ui_get_custom_viewer_hint:

{

TCustomControl *viewer = va_arg(va, TCustomControl *);

place_t *place = va_arg(va, place_t *);

int *important_lines = va_arg(va, int *);

qstring &hint = *va_arg(va, qstring *);

if ( si->cv == viewer ) // our viewer

{

if ( place == NULL )

return 0;

class_place_t *spl = (class_place_t *)place;

hint.sprnt("[Hexrays-Tools] Hint for line %d", spl->idx);

*important_lines = 1;

return 1;

}

break;

}

case ui_tform_invisible:

{

TForm *f = va_arg(va, TForm *);

if ( f == si->form )

{

delete si;

unhook_from_notification_point(HT_UI, ui_callback, NULL);

}

}

break;

}

return 0;

}

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

// Create a custom view window

bool idaapi show_classes_view(void *ud)

{

HWND hwnd = NULL;

TForm *form = find_tform("[Hexrays-Tools] Classes");

if ( form != NULL )

{

switchto_tform(form, true);

return true;

}

form = create_tform("[Hexrays-Tools] Classes", &hwnd);

// allocate block to hold info about our sample view

sample_info_t *si = new sample_info_t(form);

// create two place_t objects: for the minimal and maximal locations

class_place_t s1;

class_t * clas = get_class(get_class_by_idx(get_last_class_idx()));

class_place_t s2(get_last_class_idx(), 2, 100);

// create a custom viewer

si->cv = create_custom_viewer("", (TWinControl *)form, &s1, &s2, &s1, 0, &si->sv);

// set the handlers so we can communicate with it

set_custom_viewer_handlers(si->cv, ct_keyboard, ct_popup, NULL, cb_dblclick, ct_curpos, NULL, si);

// also set the ui event callback

hook_to_notification_point(HT_UI, ui_callback, si);

// finally display the form on the screen

open_tform(form, FORM_TAB|FORM_MENU|FORM_RESTORE|FORM_QWIDGET);

return true;

}