Data Structure Recovery using PIN and PyGraphviz

This is a simple POC PIN tool to recover data structures in dynamically linked stripped executables, mainly for analyzing small programs. The PIN tool keeps track of heap allocations done by executable and traces all the write operations to the allocated heap memory. The trace file having allocations and write operations will be used to generate graph using pygraphviz.

Tracing Size and Address of Allocations

Size of Allocation

Right now, we track libc functions malloc, realloc, calloc, sbrk, mmap and free. All these routines are instrumented using Rtn_InsertCall to fetch the size of requested allocation.

For example, for tracing malloc

RTN_InsertCall( rtn,
IPOINT_BEFORE,
(AFUNPTR)AllocBefore,
IARG_ADDRINT,
funcname,
IARG_G_ARG0_CALLEE,
IARG_RETURN_IP,
IARG_END);

We fetch the size of requested allocation using IARG_G_ARG0_CALLEE and IPOINT_BEFORE. Also, we need to identify the malloc calls that are only called from our main executable. To find this we use IARG_RETURN_IP to check if the return address of the call is part of main executable, if not we don't trace the allocation.

Address of Allocation

IARG_RETURN_IP is valid only at function entry point, so we cannot use IPOINT_AFTER along with IARG_RETURN_IP. As a work around, we save the return address during IPOINT_BEFORE. Then in instruction trace, if instruction pointer equals return address of an allocation call, we fetch the EAX value. This gives the address of allocation.

if (insaddr == retaddress){
INS_InsertCall( ins,
IPOINT_BEFORE,
(AFUNPTR)AllocAfter,
#ifdef __i386__
IARG_REG_VALUE, LEVEL_BASE::REG_EAX,
#else
IARG_REG_VALUE, LEVEL_BASE::REG_RAX,
#endif
IARG_END);
}
Now we have both address and size of allocation. These details are stored as dictionary as pairs of address : size. Also we don't remove an address when free is called upon that, instead if an already existing address is returned during an allocation call ie. reallocation, we just update the size of existing allocation for the new allocation request.

if(allocations.count(address) == 0) {
allocations.insert(std::make_pair(address, allocsize));
}
else {
std::map::iterator it = allocations.find(retval);
it->second = allocsize;
}

.data and .bss sections

data and bss sections are also added to dictionary. The size and address of these segments are fetched from main executable and added as part of allocations

if(!strcmp(sec_name.c_str(), ".bss") || !strcmp(sec_name.c_str(), ".data")) {

ADDRINT addr = SEC_Address(sec);
USIZE size = SEC_Size(sec);

if(allocations.count(addr) == 0) {
allocations.insert(std::make_pair(addr, size));
}
}
Tracing Memory Writes

We trace instructions that writes into the allocated memory. As of now only XED_ICLASS_MOV class of instructions are traced. For all XED_ICLASS_MOV instruction, we check if its a memory write instruction using INS_IsMemoryWrite and is part of main executable.

In this case, we fetch the destination address of write operation using IARG_MEMORYWRITE_EA. Then we check if the destination address is part of any allocation, on success this instruction is traced.

for (it = allocations.begin(); it != allocations.end(); it++) {
if ((des_addr >= it->first) && (des_addr < it->first+it->second)) return true;
}
Sample Trace


.data[0x804b02c,0x8]
.bss[0x804b040,0xfc4]
[email protected][0x420]
ret[0x98de000]
0x8048565 @mov dword ptr [0x804c000], eax : WRREG MEM[0x804c000] VAL[0x98de000]
0x8048575 @mov dword ptr [eax+0x8], 0x0 : WRIMM MEM[0x98de008] VAL[0]
0x804857f @mov dword ptr [edx+0x4], eax : WRREG MEM[0x98de004] VAL[0]
0x8048587 @mov dword ptr [eax], 0x10 : WRIMM MEM[0x98de000] VAL[0x10]
0x80485a0 @mov dword ptr [eax+0x4], edx : WRREG MEM[0x98de004] VAL[0x98de010]
0x80485ac @mov dword ptr [eax+0x8], edx : WRREG MEM[0x98de018] VAL[0x98de000]
Graphing

Node Create

For each allocation in trace file generated by PIN tool, a new node is created in the graph. Each node is uniquely identified using a node id which is assigned sequentially. An ordered dictionary is maintained, key being node id and value is dictionary of address and size of allocation. New allocations are added to the start of ordered dictionary.

An edge count is associated with each of created node. This will be used for pruning away nodes without any edges.

Separate nodes are created for bss and data sections. But this is optional.

Example

Say a structure is allocated in heap using malloc, this is how a node will look like

[email protected][0x20]
ret[0x8fcf030]


-------------------
| [0] 0x8fcf030 |
-------------------
[0] is the node id, this could signify the order of allocation. Every new allocator call gets a new id, irrespective of the return address

0x8fcf030 is the address returned by allocator call

Node Update

For each instruction, fetch the target address of write operation. If the target address is part of any allocation, update the node to which the target address belongs to. Basically we create a new port in the record node.

A new port signifies an element of an allocation, say element of a structure.

Then check if the source value is part of any allocation. If yes, we consider the source value as an address. Then update the node to which the source address belongs to. This operation could be interpreted as a pointer assignment [or link creation]

[email protected][0x20]
ret[0x8fcf030]
0x8048957 @mov byte ptr [eax+edx*1], 0x0 : WRIMM MEM[0x8fcf031] VAL[0]
0x80489bb @mov dword ptr [eax+0x14], edx : WRREG MEM[0x8fcf044] VAL[0x8fcf058]
0x8048a40 @mov dword ptr [eax+0x18], edx : WRREG MEM[0x8fcf048] VAL[0x8fcf008]
0x8048a4e @mov dword ptr [eax+0x1c], edx : WRREG MEM[0x8fcf04c] VAL[0x8fcf008]

-------------------------------------------------------------------------
| [0] 0x8fcf030 | 0x8fcf031 | 0x8fcf044 | 0x8fcf048 | 0x8fcf04c |
-------------------------------------------------------------------------
Now first field [0] 0x8fcf030 is the meta data for the chunk ie node id and return address of allocator call. The rest of 4 fields signifies 4 individual write operations into this allocation [example, 4 elements of a structure]

Create Link

If both source and destination values are valid address and belongs to a traced allocation, we link both ports of the nodes. Whenever a link is created, edge count of source and destination are incremented.

Similarly, during memory overwrite an edge is removed and edge count is decremented.

Example,

[email protected][0x20]
ret[0x8fcf008]
…...

[email protected][0x20]
ret[0x8fcf030]
…...
0x80489bb @mov dword ptr [eax+0x14], edx : WRREG MEM[0x8fcf044] VAL[0x8fcf058]
0x8048a40 @mov dword ptr [eax+0x18], edx : WRREG MEM[0x8fcf048] VAL[0x8fcf008]
0x8048a4e @mov dword ptr [eax+0x1c], edx : WRREG MEM[0x8fcf04c] VAL[0x8fcf008]
Above is a series of pointer writes into memory allocated at 0x8fcf030. The address points to another allocation at 0x8fcf008. Hence we link both

Prune Node

Finally after parsing all instructions, remove nodes that doesn't have any edges. For this, check if the edge count for a node is 0. If yes, remove the node.

Other Options

By default, we consider only the first non-NULL write operation for node update and link creation. This might be good enough to reveal some of data structures. Any memory writes to an address after first write non-NULL are skipped. But one can use relink option to consider more than single write operation for graphing. This could be useful when relink operations are done, say circular linked list.

NULL writes can also be enabled as option. This might be useful along with relink.

The tool itself doesn't itself have the intelligence to say what data structure is used, but can graph the allocation and links to help someone understand a data structure from the revealed shape.

Example - Singly Linked List

Example - Binary Tree



Example - HackIM Mixme Circular Doubly Linked List



The POC code which I use for CTF is available here. To repeat again, this works on small binaries, as things get complex the graph might make less sense. There is lot of scope for improvement though.

References

AES Whitebox Unboxing: No Such Problem, this served as excellent reference for the usage of PIN tool and pygraphviz to visualize memory access

Thanks to Danny K, for help with Intel PIN Framework.
Uncategorized