#ifndef __PY_GRAPH__
#define __PY_GRAPH__
//
class py_graph_t : public py_customidamemo_t
{
typedef py_customidamemo_t inherited;
protected:
virtual void node_info_modified(int n, const node_info_t *ni, uint32 flags)
{
if ( ni == NULL )
{
node_cache.erase(n);
}
else
{
nodetext_cache_t *c = node_cache.get(n);
if ( c != NULL )
{
if ( (flags & NIF_TEXT) == NIF_TEXT )
c->text = ni->text;
if ( (flags & NIF_BG_COLOR) == NIF_BG_COLOR )
c->bgcolor = ni->bg_color;
}
}
}
void collect_class_callbacks_ids(callbacks_ids_t *out);
private:
enum
{
GRCODE_HAVE_USER_HINT = 0x00010000,
GRCODE_HAVE_CLICKED = 0x00020000,
GRCODE_HAVE_DBL_CLICKED = 0x00040000,
GRCODE_HAVE_GOTFOCUS = 0x00080000,
GRCODE_HAVE_LOSTFOCUS = 0x00100000,
GRCODE_HAVE_CHANGED_CURRENT = 0x00200000,
GRCODE_HAVE_COMMAND = 0x00400000,
GRCODE_HAVE_CREATING_GROUP = 0x00800000,
GRCODE_HAVE_DELETING_GROUP = 0x01000000,
GRCODE_HAVE_GROUP_VISIBILITY = 0x02000000,
};
struct nodetext_cache_t
{
qstring text;
bgcolor_t bgcolor;
nodetext_cache_t(const nodetext_cache_t &rhs): text(rhs.text), bgcolor(rhs.bgcolor) { }
nodetext_cache_t(const char *t, bgcolor_t c): text(t), bgcolor(c) { }
nodetext_cache_t() { }
};
class nodetext_cache_map_t: public std::map
{
public:
nodetext_cache_t *get(int node_id)
{
iterator it = find(node_id);
if ( it == end() )
return NULL;
return &it->second;
}
nodetext_cache_t *add(const int node_id, const char *text, bgcolor_t bgcolor = DEFCOLOR)
{
return &(insert(std::make_pair(node_id, nodetext_cache_t(text, bgcolor))).first->second);
}
};
class cmdid_map_t: public std::map
{
private:
Py_ssize_t uid;
public:
cmdid_map_t()
{
// We start by one and keep zero for error id
uid = 1;
}
void add(py_graph_t *pyg)
{
(*this)[uid] = pyg;
++uid;
}
const Py_ssize_t id() const
{
return uid;
}
void clear(py_graph_t *pyg)
{
iterator e = end();
for (iterator it=begin();it!=end();)
{
if ( it->second == pyg )
{
iterator temp = it++;
erase(temp);
}
else
++it;
}
}
py_graph_t *get(Py_ssize_t id)
{
iterator it = find(id);
return it == end() ? NULL : it->second;
}
};
static cmdid_map_t cmdid_pyg;
bool refresh_needed;
nodetext_cache_map_t node_cache;
// instance callback
int gr_callback(int code, va_list va);
// static callback
static int idaapi s_callback(void *obj, int code, va_list va)
{
QASSERT(30453, py_customidamemo_t::lookup_info.find_by_py_view(NULL, NULL, (py_graph_t *) obj));
PYW_GIL_GET;
return ((py_graph_t *)obj)->gr_callback(code, va);
}
static bool idaapi s_menucb(void *ud)
{
PYW_GIL_GET;
Py_ssize_t id = (Py_ssize_t)ud;
py_graph_t *_this = cmdid_pyg.get(id);
if ( _this != NULL )
_this->on_command(id);
return true;
}
void on_command(Py_ssize_t id)
{
// Check return value to OnRefresh() call
PYW_GIL_CHECK_LOCKED_SCOPE();
newref_t ret(PyObject_CallMethod(self.o, (char *)S_ON_COMMAND, "n", id));
PyW_ShowCbErr(S_ON_COMMAND);
}
// Refresh user-defined graph node number and edges
// It calls Python method and expects that the user already filled
// the nodes and edges. The nodes and edges are retrieved and passed to IDA
void on_user_refresh(mutable_graph_t *g);
// Retrieves the text for user-defined graph node
// It expects either a string or a tuple (string, bgcolor)
bool on_user_text(mutable_graph_t * /*g*/, int node, const char **str, bgcolor_t *bg_color);
// Retrieves the hint for the user-defined graph
// Calls Python and expects a string or None
int on_user_hint(mutable_graph_t *, int mousenode, int /*mouseedge_src*/, int /*mouseedge_dst*/, char **hint);
// graph is being destroyed
void on_graph_destroyed(mutable_graph_t * /*g*/ = NULL)
{
refresh_needed = true;
node_cache.clear();
}
// graph is being clicked
int on_clicked(
graph_viewer_t * /*view*/,
selection_item_t * /*item1*/,
graph_item_t *item2)
{
// in: graph_viewer_t *view
// selection_item_t *current_item1
// graph_item_t *current_item2
// out: 0-ok, 1-ignore click
// this callback allows you to ignore some clicks.
// it occurs too early, internal graph variables are not updated yet
// current_item1, current_item2 point to the same thing
// item2 has more information.
// see also: kernwin.hpp, custom_viewer_click_t
if ( item2->n == -1 )
return 1;
PYW_GIL_CHECK_LOCKED_SCOPE();
newref_t result(
PyObject_CallMethod(
self.o,
(char *)S_ON_CLICK,
"i",
item2->n));
PyW_ShowCbErr(S_ON_CLICK);
return result == NULL || !PyObject_IsTrue(result.o);
}
// a graph node has been double clicked
int on_dblclicked(graph_viewer_t * /*view*/, selection_item_t *item)
{
// in: graph_viewer_t *view
// selection_item_t *current_item
// out: 0-ok, 1-ignore click
//graph_viewer_t *v = va_arg(va, graph_viewer_t *);
//selection_item_t *s = va_arg(va, selection_item_t *);
if ( item == NULL || !item->is_node )
return 1;
PYW_GIL_CHECK_LOCKED_SCOPE();
newref_t result(
PyObject_CallMethod(
self.o,
(char *)S_ON_DBL_CLICK,
"i",
item->node));
PyW_ShowCbErr(S_ON_DBL_CLICK);
return result == NULL || !PyObject_IsTrue(result.o);
}
// a graph viewer got focus
void on_gotfocus(graph_viewer_t * /*view*/)
{
if ( self.o == NULL )
return;
PYW_GIL_CHECK_LOCKED_SCOPE();
newref_t result(
PyObject_CallMethod(
self.o,
(char *)S_ON_ACTIVATE,
NULL));
PyW_ShowCbErr(S_ON_ACTIVATE);
}
// a graph viewer lost focus
void on_lostfocus(graph_viewer_t * /*view*/)
{
if ( self.o == NULL )
return;
PYW_GIL_CHECK_LOCKED_SCOPE();
newref_t result(
PyObject_CallMethod(
self.o,
(char *)S_ON_DEACTIVATE,
NULL));
PyW_ShowCbErr(S_ON_DEACTIVATE);
}
// a new graph node became the current node
int on_changed_current(graph_viewer_t * /*view*/, int curnode)
{
// in: graph_viewer_t *view
// int curnode
// out: 0-ok, 1-forbid to change the current node
if ( curnode < 0 )
return 0;
PYW_GIL_CHECK_LOCKED_SCOPE();
newref_t result(
PyObject_CallMethod(
self.o,
(char *)S_ON_SELECT,
"i",
curnode));
PyW_ShowCbErr(S_ON_SELECT);
return !(result != NULL && PyObject_IsTrue(result.o));
}
// a group is being created
int on_creating_group(mutable_graph_t *my_g, intvec_t *my_nodes)
{
PYW_GIL_CHECK_LOCKED_SCOPE();
newref_t py_nodes(PyList_New(my_nodes->size()));
int i;
intvec_t::const_iterator p;
for ( i = 0, p=my_nodes->begin(); p != my_nodes->end(); ++p, ++i )
PyList_SetItem(py_nodes.o, i, PyInt_FromLong(*p));
newref_t py_result(
PyObject_CallMethod(
self.o,
(char *)S_ON_CREATING_GROUP,
"O",
py_nodes.o));
PyW_ShowCbErr(S_ON_CREATING_GROUP);
return (py_result == NULL || !PyInt_Check(py_result.o)) ? 1 : PyInt_AsLong(py_result.o);
}
// a group is being deleted
int on_deleting_group(mutable_graph_t * /*g*/, int old_group)
{
PYW_GIL_CHECK_LOCKED_SCOPE();
// TODO
return 0;
}
// a group is being collapsed/uncollapsed
int on_group_visibility(mutable_graph_t * /*g*/, int group, bool expand)
{
PYW_GIL_CHECK_LOCKED_SCOPE();
// TODO
return 0;
}
void show()
{
TForm *form;
if ( lookup_info.find_by_py_view(&form, NULL, this) )
open_tform(form, FORM_TAB|FORM_MENU|FORM_QWIDGET);
}
void jump_to_node(int nid)
{
ref_t nodes(PyW_TryGetAttrString(self.o, S_M_NODES));
if ( nid >= PyList_Size(nodes.o) )
return;
viewer_center_on(view, nid);
int x, y;
// will return a place only when a node was previously selected
place_t *old_pl = get_custom_viewer_place(view, false, &x, &y);
if ( old_pl != NULL )
{
user_graph_place_t *new_pl = (user_graph_place_t *) old_pl->clone();
new_pl->node = nid;
jumpto(view, new_pl, x, y);
delete new_pl;
}
}
virtual void refresh()
{
refresh_needed = true;
inherited::refresh();
}
int initialize(PyObject *self, const char *title)
{
PYW_GIL_CHECK_LOCKED_SCOPE();
if ( !collect_pyobject_callbacks(self) )
return -1;
// Create form
HWND hwnd = NULL;
TForm *form = create_tform(title, &hwnd);
if ( hwnd != NULL ) // Created new tform
{
lookup_info_t::entry_t &e = lookup_info.new_entry(this);
// get a unique graph id
netnode id;
char grnode[MAXSTR];
qsnprintf(grnode, sizeof(grnode), "$ pygraph %s", title);
id.create(grnode);
graph_viewer_t *pview = create_graph_viewer(form, id, s_callback, this, 0);
open_tform(form, FORM_TAB | FORM_MENU | FORM_QWIDGET);
if ( pview != NULL )
viewer_fit_window(pview);
bind(self, pview);
install_custom_viewer_handlers();
refresh();
lookup_info.commit(e, form, view);
}
else
{
show();
}
viewer_fit_window(view);
return 0;
}
Py_ssize_t add_command(const char *title, const char *hotkey)
{
if ( !has_callback(GRCODE_HAVE_COMMAND) || view == NULL)
return 0;
Py_ssize_t cmd_id = cmdid_pyg.id();
bool ok = viewer_add_menu_item(view, title, s_menucb, (void *)cmd_id, hotkey, 0);
if ( !ok )
return 0;
cmdid_pyg.add(this);
return cmd_id;
}
public:
py_graph_t()
{
// form = NULL;
refresh_needed = true;
}
virtual ~py_graph_t()
{
// Remove all associated commands from the list
cmdid_pyg.clear(this);
}
static void SelectNode(PyObject *self, int nid)
{
if ( nid < 0 )
return;
py_graph_t *_this = view_extract_this(self);
if ( _this == NULL || !lookup_info.find_by_py_view(NULL, NULL, _this) )
return;
_this->jump_to_node(nid);
}
static Py_ssize_t AddCommand(PyObject *self, const char *title, const char *hotkey)
{
py_graph_t *_this = view_extract_this(self);
if ( _this == NULL || !lookup_info.find_by_py_view(NULL, NULL, _this) )
return 0;
return _this->add_command(title, hotkey);
}
static void Close(PyObject *self)
{
TForm *form;
py_graph_t *_this = view_extract_this(self);
if ( _this == NULL || !lookup_info.find_by_py_view(&form, NULL, _this) )
return;
close_tform(form, FORM_CLOSE_LATER);
}
static py_graph_t *Show(PyObject *self)
{
PYW_GIL_CHECK_LOCKED_SCOPE();
py_graph_t *py_graph = view_extract_this(self);
// New instance?
if ( py_graph == NULL )
{
qstring title;
if ( !PyW_GetStringAttr(self, S_M_TITLE, &title) )
return NULL;
// Form already created? try to get associated py_graph instance
// so that we reuse it
graph_viewer_t *found_view;
TForm *form = find_tform(title.c_str());
if ( form != NULL )
lookup_info.find_by_form(&found_view, (py_customidamemo_t**) &py_graph, form);
if ( py_graph == NULL )
{
py_graph = new py_graph_t();
}
else
{
// unbind so we are rebound
py_graph->unbind(false);
py_graph->refresh_needed = true;
}
if ( py_graph->initialize(self, title.c_str()) < 0 )
{
delete py_graph;
py_graph = NULL;
}
}
else
{
py_graph->show();
}
return py_graph;
}
};
//-------------------------------------------------------------------------
void py_graph_t::collect_class_callbacks_ids(callbacks_ids_t *out)
{
inherited::collect_class_callbacks_ids(out);
out->add(S_ON_REFRESH, 0);
out->add(S_ON_GETTEXT, 0);
out->add(S_M_EDGES, -1);
out->add(S_M_NODES, -1);
out->add(S_ON_HINT, GRCODE_HAVE_USER_HINT);
out->add(S_ON_CLICK, GRCODE_HAVE_CLICKED);
out->add(S_ON_DBL_CLICK, GRCODE_HAVE_DBL_CLICKED);
out->add(S_ON_COMMAND, GRCODE_HAVE_COMMAND);
out->add(S_ON_SELECT, GRCODE_HAVE_CHANGED_CURRENT);
out->add(S_ON_ACTIVATE, GRCODE_HAVE_GOTFOCUS);
out->add(S_ON_DEACTIVATE, GRCODE_HAVE_LOSTFOCUS);
out->add(S_ON_CREATING_GROUP, GRCODE_HAVE_CREATING_GROUP);
out->add(S_ON_DELETING_GROUP, GRCODE_HAVE_DELETING_GROUP);
out->add(S_ON_GROUP_VISIBILITY, GRCODE_HAVE_GROUP_VISIBILITY);
}
//-------------------------------------------------------------------------
void py_graph_t::on_user_refresh(mutable_graph_t *g)
{
if ( !refresh_needed || self == NULL /* Happens at creation-time */ )
return;
// Check return value to OnRefresh() call
PYW_GIL_CHECK_LOCKED_SCOPE();
newref_t ret(PyObject_CallMethod(self.o, (char *)S_ON_REFRESH, NULL));
PyW_ShowCbErr(S_ON_REFRESH);
if ( ret != NULL && PyObject_IsTrue(ret.o) )
{
// Refer to the nodes
ref_t nodes(PyW_TryGetAttrString(self.o, S_M_NODES));
if ( ret != NULL && PyList_Check(nodes.o) )
{
// Refer to the edges
ref_t edges(PyW_TryGetAttrString(self.o, S_M_EDGES));
if ( ret != NULL && PyList_Check(edges.o) )
{
// Resize the nodes
int max_nodes = abs(int(PyList_Size(nodes.o)));
g->clear();
g->resize(max_nodes);
// Mark that we refreshed already
refresh_needed = false;
// Clear cached nodes
node_cache.clear();
// Get the edges
for ( int i=(int)PyList_Size(edges.o)-1; i>=0; i-- )
{
// Each list item is a sequence (id1, id2)
borref_t item(PyList_GetItem(edges.o, i));
if ( !PySequence_Check(item.o) )
continue;
// Get and validate each of the two elements in the sequence
int edge_ids[2];
int j;
for ( j=0; j max_nodes )
break;
edge_ids[j] = v;
}
// Incomplete?
if ( j != qnumber(edge_ids) )
break;
// Add the edge
g->add_edge(edge_ids[0], edge_ids[1], NULL);
}
}
}
}
}
//-------------------------------------------------------------------------
bool py_graph_t::on_user_text(mutable_graph_t * /*g*/, int node, const char **str, bgcolor_t *bg_color)
{
// If already cached then return the value
nodetext_cache_t *c = node_cache.get(node);
if ( c != NULL )
{
*str = c->text.c_str();
if ( bg_color != NULL )
*bg_color = c->bgcolor;
return true;
}
// Not cached, call Python
PYW_GIL_CHECK_LOCKED_SCOPE();
newref_t result(PyObject_CallMethod(self.o, (char *)S_ON_GETTEXT, "i", node));
PyW_ShowCbErr(S_ON_GETTEXT);
if ( result == NULL )
return false;
bgcolor_t cl = bg_color == NULL ? DEFCOLOR : *bg_color;
const char *s;
// User returned a string?
if ( PyString_Check(result.o) )
{
s = PyString_AsString(result.o);
if ( s == NULL )
s = "";
c = node_cache.add(node, s, cl);
}
// User returned a sequence of text and bgcolor
else if ( PySequence_Check(result.o) && PySequence_Size(result.o) == 2 )
{
newref_t py_str(PySequence_GetItem(result.o, 0));
newref_t py_color(PySequence_GetItem(result.o, 1));
if ( py_str == NULL || !PyString_Check(py_str.o) || (s = PyString_AsString(py_str.o)) == NULL )
s = "";
if ( py_color != NULL && PyNumber_Check(py_color.o) )
cl = bgcolor_t(PyLong_AsUnsignedLong(py_color.o));
c = node_cache.add(node, s, cl);
}
*str = c->text.c_str();
if ( bg_color != NULL )
*bg_color = c->bgcolor;
return true;
}
//-------------------------------------------------------------------------
int py_graph_t::on_user_hint(mutable_graph_t *, int mousenode, int /*mouseedge_src*/, int /*mouseedge_dst*/, char **hint)
{
// 'hint' must be allocated by qalloc() or qstrdup()
// out: 0-use default hint, 1-use proposed hint
// We dispatch hints over nodes only
if ( mousenode == -1 )
return 0;
PYW_GIL_CHECK_LOCKED_SCOPE();
newref_t result(PyObject_CallMethod(self.o, (char *)S_ON_HINT, "i", mousenode));
PyW_ShowCbErr(S_ON_HINT);
bool ok = result != NULL && PyString_Check(result.o);
if ( ok )
*hint = qstrdup(PyString_AsString(result.o));
return ok; // use our hint
}
//-------------------------------------------------------------------------
int py_graph_t::gr_callback(int code, va_list va)
{
int ret;
switch ( code )
{
//
case grcode_user_text:
{
mutable_graph_t *g = va_arg(va, mutable_graph_t *);
int node = va_arg(va, int);
const char **result = va_arg(va, const char **);
bgcolor_t *bgcolor = va_arg(va, bgcolor_t *);
ret = on_user_text(g, node, result, bgcolor);
break;
}
//
case grcode_destroyed:
on_graph_destroyed(va_arg(va, mutable_graph_t *));
ret = 0;
break;
//
case grcode_clicked:
if ( has_callback(GRCODE_HAVE_CLICKED) )
{
graph_viewer_t *view = va_arg(va, graph_viewer_t *);
selection_item_t *item = va_arg(va, selection_item_t *);
graph_item_t *gitem = va_arg(va, graph_item_t *);
ret = on_clicked(view, item, gitem);
}
else
{
// Ignore the click
ret = 1;
}
break;
//
case grcode_dblclicked:
if ( has_callback(GRCODE_HAVE_DBL_CLICKED) )
{
graph_viewer_t *view = va_arg(va, graph_viewer_t *);
selection_item_t *item = va_arg(va, selection_item_t *);
ret = on_dblclicked(view, item);
}
else
ret = 0; // We don't want to ignore the double click, but rather
// fallback to the default behavior (e.g., double-clicking
// on an edge will to jump to the node on the other side
// of that edge.)
break;
//
case grcode_gotfocus:
if ( has_callback(GRCODE_HAVE_GOTFOCUS) )
on_gotfocus(va_arg(va, graph_viewer_t *));
ret = 0;
break;
//
case grcode_lostfocus:
if ( has_callback(GRCODE_HAVE_LOSTFOCUS) )
on_lostfocus(va_arg(va, graph_viewer_t *));
ret = 0;
break;
//
case grcode_user_refresh:
on_user_refresh(va_arg(va, mutable_graph_t *));
ret = 1;
break;
//
case grcode_user_hint:
if ( has_callback(GRCODE_HAVE_USER_HINT) )
{
mutable_graph_t *g = va_arg(va, mutable_graph_t *);
int mousenode = va_arg(va, int);
int mouseedge_src = va_arg(va, int);
int mouseedge_dest = va_arg(va, int);
char **hint = va_arg(va, char **);
ret = on_user_hint(g, mousenode, mouseedge_src, mouseedge_dest, hint);
}
else
{
ret = 0;
}
break;
//
case grcode_changed_current:
if ( has_callback(GRCODE_HAVE_CHANGED_CURRENT) )
{
graph_viewer_t *view = va_arg(va, graph_viewer_t *);
int cur_node = va_arg(va, int);
ret = on_changed_current(view, cur_node);
}
else
ret = 0; // allow selection change
break;
//
case grcode_creating_group: // a group is being created
if ( has_callback(GRCODE_HAVE_CREATING_GROUP) )
{
mutable_graph_t *g = va_arg(va, mutable_graph_t*);
intvec_t *nodes = va_arg(va, intvec_t*);
ret = on_creating_group(g, nodes);
}
else
{
ret = 0; // Ok to create
}
break;
//
case grcode_deleting_group: // a group is being deleted
if ( has_callback(GRCODE_HAVE_DELETING_GROUP) )
{
mutable_graph_t *g = va_arg(va, mutable_graph_t*);
int old_group = va_arg(va, int);
ret = on_deleting_group(g, old_group);
}
else
{
ret = 0; // Ok to delete
}
break;
//
case grcode_group_visibility: // a group is being collapsed/uncollapsed
if ( has_callback(GRCODE_HAVE_GROUP_VISIBILITY) )
{
mutable_graph_t *g = va_arg(va, mutable_graph_t*);
int group = va_arg(va, int);
bool expand = bool(va_arg(va, int));
ret = on_group_visibility(g, group, expand);
}
else
{
ret = 0; // Ok.
}
break;
//
default:
ret = 0;
break;
}
//grcode_changed_graph, // new graph has been set
//grcode_user_size, // calculate node size for user-defined graph
//grcode_user_title, // render node title of a user-defined graph
//grcode_user_draw, // render node of a user-defined graph
return ret;
}
//-------------------------------------------------------------------------
py_graph_t::cmdid_map_t py_graph_t::cmdid_pyg;
bool pyg_show(PyObject *self)
{
return py_graph_t::Show(self) != NULL;
}
void pyg_close(PyObject *self)
{
py_graph_t::Close(self);
}
PyObject *pyg_add_command(PyObject *self, const char *title, const char *hotkey)
{
PYW_GIL_CHECK_LOCKED_SCOPE();
return Py_BuildValue("n", py_graph_t::AddCommand(self, title, hotkey));
}
void pyg_select_node(PyObject *self, int nid)
{
py_graph_t::SelectNode(self, nid);
}
//
//--------------------------------------------------------------------------
//
void pyg_close(PyObject *self);
PyObject *pyg_add_command(PyObject *self, const char *title, const char *hotkey);
void pyg_select_node(PyObject *self, int nid);
bool pyg_show(PyObject *self);
//
#endif