unpacker.cpp

#include "pin.H"
#include <stdio.h>
#include <stdbool.h>
#include <vector>
#include <map>
#include <unistd.h>
#include <atomic>
#include "src/dta.h"

const size_t WRITE_THRESHOLD = 50;
TaintEngine g_taint;

using namespace std;

struct Page {
	int prot;
	bool writed;
	size_t c_writed;
	ADDRINT last_write;
};

struct ThreadS {
	ADDRINT syscallc;
	int prot;
	size_t length;
	ADDRINT addr;
	bool valid;
};

struct DumpReg {
	ADDRINT addr;
	size_t size;
	int prot;
};

TLS_KEY tls_key;
PIN_MUTEX e_regions_mutex;
map<ADDRINT, Page> vpages;
atomic<bool> dump_done{false};

vector<DumpReg> CollectRegions() {
	PIN_MutexLock(&e_regions_mutex);
	vector<pair<ADDRINT, int>> sorted_pages;
	for (auto i = vpages.begin(); i != vpages.end(); i++) {
		if (i->second.writed == true) {
			sorted_pages.push_back(make_pair(i->first, i->second.prot));
		}
	}
	PIN_MutexUnlock(&e_regions_mutex);
	vector<DumpReg> result;
	if (!sorted_pages.empty()) {
		sort(sorted_pages.begin(), sorted_pages.end());
		DumpReg current;
		current.addr = sorted_pages[0].first;
		current.size = page_size;
		current.prot = sorted_pages[0].second;
		for (size_t i = 1; i < sorted_pages.size(); i++) {
			if ((sorted_pages[i].first  == (current.addr + current.size)) && (sorted_pages[i].second == current.prot)) {
				current.size += page_size;
			} else {
				result.push_back(current);
        			current.addr = sorted_pages[i].first;
        			current.size = page_size;
        			current.prot = sorted_pages[i].second;
			}
		}
		result.push_back(current);
	}
	return result;
}

void DoDump(ADDRINT oep_addr){
	vector<DumpReg> regions = CollectRegions();
	if (regions.empty()) {
		printf("%s", "Nothing to dump");
		return;
	}
	char raw_name[256];
	char json_name[256];
	snprintf(raw_name, sizeof(raw_name), "dump_%d_%p.raw", PIN_GetPid(), (void*)oep_addr);
	snprintf(json_name, sizeof(json_name), "dump_%d_%p.json", PIN_GetPid(), (void*)oep_addr);
	FILE* raw = fopen(raw_name, "wb");
	FILE* json = fopen(json_name, "w");
	if (!(raw && json)) {
		printf("%s", "Failed to open\n");
		if (raw) fclose(raw);
		if (json) fclose(json);
		return;
	}
	size_t current_offset = 0;
	fprintf(json, "{\n  \"pid\": %d,\n  \"oep\": \"%p\",\n  \"regions\": [\n", PIN_GetPid(), (void*)oep_addr);
	for (size_t i = 0; i < regions.size(); i++) {
		void* buf = malloc(regions[i].size);
		if (!buf) {
		       	continue;
		}
		size_t copy = PIN_SafeCopy(buf, (void*)regions[i].addr, regions[i].size);
		if (copy == 0) continue;
		fwrite(buf, 1, copy, raw);
		fprintf(json, "    {\"addr\": \"%p\", \"size\": %zu, \"prot\": %d, \"offset\": %zu}",(void*)regions[i].addr, regions[i].size, regions[i].prot, copy);
		if (i + 1 < regions.size()) fprintf(json, ",");
		fprintf(json, "\n");
		current_offset += copy;
		free(buf);

	}
	fprintf(json, "  ]\n}\n");
	fclose(raw);
	fclose(json);
	printf("%s", "Dump complete\n");
	fflush(stdout);
	PIN_ExitApplication(0);
}

VOID ThreadStart(THREADID tid, CONTEXT* ctx, INT32 flags, VOID* v) {
	ThreadS* state = new ThreadS();
	state->syscallc = 0;
	state->prot = 0;
	state->length = 0;
	state->addr = 0;
	state->valid = false;
	PIN_SetThreadData(tls_key, state, tid);	
}

VOID ThreadFini(THREADID tid, const CONTEXT* ctx, INT32 code, VOID* v) {
	ThreadS* state = (ThreadS*)PIN_GetThreadData(tls_key, tid);
	delete state;	
}

VOID BeforeSyscall(THREADID tid, CONTEXT* ctx, SYSCALL_STANDARD std, VOID* v) {
	ThreadS* state = (ThreadS*)PIN_GetThreadData(tls_key, tid);
	state->syscallc = PIN_GetSyscallNumber(ctx, std);	
	state->valid = 0;
	if (state->syscallc == 9) { //mmap
		state->prot = PIN_GetSyscallArgument(ctx, std, 2);
		if ((state->prot & 0x4) || (state->prot & 0x2)) {
			state->valid = 1;
			state->length = PIN_GetSyscallArgument(ctx, std, 1);
			printf("%s", "EXEC or WRITE bit mmap\n");
			fflush(stdout);
		}
	} else if (state->syscallc == 10) { //mprotect
		state->prot = PIN_GetSyscallArgument(ctx, std, 2);
		if ((state->prot & 0x4) || (state->prot & 0x2)) {
			state->valid = 1;
			state->addr = PIN_GetSyscallArgument(ctx, std, 0);
			state->length = PIN_GetSyscallArgument(ctx, std, 1);
			printf("%s", "EXEC or WRITE bit mproterct\n");
			fflush(stdout);
		}
	}  else if (state->syscallc == 11) { //munmap
		state->valid = 1;
		state->addr = PIN_GetSyscallArgument(ctx, std, 0);
		state->length = PIN_GetSyscallArgument(ctx, std, 1);
	} else {
		state->valid = 0;
	}
}

VOID AfterSyscall(THREADID tid, CONTEXT* ctx, SYSCALL_STANDARD std, VOID* v) {
	ThreadS* state = (ThreadS*)PIN_GetThreadData(tls_key, tid);
	if (!state->valid) {
		return;
	}
	if (state->syscallc == 9) {
		state->addr =  PIN_GetSyscallReturn(ctx, std);
		if (state->addr != (ADDRINT) -1){ // 0xFFFFF...
			PIN_MutexLock(&e_regions_mutex);
			ADDRINT page_start = state->addr & ~(page_size - 1);
			ADDRINT page_end = (state->addr + state->length + page_size - 1) & ~(page_size - 1);
			for (ADDRINT i = page_start; i < page_end; i += page_size) {
    				auto lookup = vpages.find(i);
    				if (lookup == vpages.end()) {
        				vpages[i] = {state->prot, 0, 0, 0};
    				} else {
        				lookup->second.prot = state->prot;
    				}
			}
			PIN_MutexUnlock(&e_regions_mutex);
			printf("%p\n", (void *)state->addr);
			fflush(stdout);
		}
	} else if (state->syscallc == 10 && ((state->prot & 0x4) || (state->prot & 0x2)) && (PIN_GetSyscallReturn(ctx, std) == 0)) {
		PIN_MutexLock(&e_regions_mutex);
		ADDRINT page_start = state->addr & ~(page_size - 1);
		ADDRINT page_end = (state->addr + state->length + page_size - 1) & ~(page_size - 1);
		for (ADDRINT i = page_start; i < page_end; i += page_size) {
			auto lookup = vpages.find(i);
    			if (lookup == vpages.end()) {
        			vpages[i] = {state->prot, 0, 0, 0};
    			} else {
        			lookup->second.prot = state->prot;
    			}
		}
		PIN_MutexUnlock(&e_regions_mutex);
	} else if (state->syscallc == 11 && PIN_GetSyscallReturn(ctx, std) == 0) {
		PIN_MutexLock(&e_regions_mutex);
		ADDRINT page_start = state->addr & ~(page_size - 1);
		ADDRINT page_end = (state->addr + state->length + page_size - 1) & ~(page_size - 1);
		for (ADDRINT i = page_start; i < page_end; i += page_size) {
			auto lookup = vpages.find(i);
			if (lookup != vpages.end() && !lookup->second.writed) {
				vpages.erase(lookup);
    			}
		}
		PIN_MutexUnlock(&e_regions_mutex);
	}  
	state->valid = 0;
}

VOID CheckMemWrite(ADDRINT mem_addr, ADDRINT instr_addr, UINT32 size) {
	if (mem_addr == 0 || size == 0) return;

	ADDRINT write_end = mem_addr + size;
	PIN_MutexLock(&e_regions_mutex);
	ADDRINT current = mem_addr;
	while (current < write_end) {
		ADDRINT page = current & ~(page_size - 1);
    		ADDRINT page_next = page + page_size;
		ADDRINT reg_end;
		if (page_next < write_end) {
			reg_end = page_next;
		} else {
			reg_end = write_end;
		}
        	auto lookup = vpages.find(page);
        	if (lookup != vpages.end()) {
            		lookup->second.writed = true;
            		lookup->second.c_writed++;
            		lookup->second.last_write = instr_addr;
        		*g_taint.shadow_byte(current) = TAINT_UNPACKED;
		}
		current = reg_end;
	}
	PIN_MutexUnlock(&e_regions_mutex);
}

VOID CheckOEP(ADDRINT target_addr, ADDRINT instr_addr) { //check on OEP
	if (target_addr == 0) return;
	if (dump_done.load()) return;

	ADDRINT page = target_addr & ~(page_size - 1);
	bool tainted = g_taint.is_tainted(target_addr);
	PIN_MutexLock(&e_regions_mutex);
	auto lookup = vpages.find(page);
	if (lookup == vpages.end()) {
		PIN_MutexUnlock(&e_regions_mutex);
		return;
	}
	bool is_exec = lookup->second.prot & 0x4;
	if (tainted && is_exec) {
		fprintf(stderr, "[DTA-SINK] OEP confirmed: tainted jump to %p from %p\n",
                (void*)target_addr, (void*)instr_addr);
		dump_done = true;
		PIN_MutexUnlock(&e_regions_mutex);
		DoDump(target_addr);
		return;
	}
        if (!tainted && lookup->second.writed && is_exec && (lookup->second.c_writed >= WRITE_THRESHOLD)) {
            		printf("[NOT-DTA] OEP candidate: jump to %p from %p, page %p, writes: %zu\n", (void*)target_addr, (void*)instr_addr, (void*)page, lookup->second.c_writed);
			fflush(stdout);
			dump_done = true;
			PIN_MutexUnlock(&e_regions_mutex);
			DoDump(target_addr);
			return;
        	}
	PIN_MutexUnlock(&e_regions_mutex);
}
VOID InstructionInstrumentation(INS ins, VOID* v) {
	UINT32 memOperands = INS_MemoryOperandCount(ins);
	for (UINT32 memOp = 0; memOp < memOperands; memOp++) {
		if (INS_MemoryOperandIsWritten(ins, memOp) && INS_IsStandardMemop(ins)) {
			INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR)CheckMemWrite, IARG_MEMORYOP_EA, memOp, IARG_INST_PTR, IARG_MEMORYOP_SIZE, memOp, IARG_END);
		}
	}
	if (INS_IsControlFlow(ins)) {
		if (INS_IsValidForIpointTakenBranch(ins) || !INS_HasFallThrough(ins)) {
			INS_InsertCall(ins, IPOINT_TAKEN_BRANCH, (AFUNPTR)CheckOEP, IARG_BRANCH_TARGET_ADDR, IARG_INST_PTR, IARG_END);
		} else {
			INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR)CheckOEP, IARG_BRANCH_TARGET_ADDR, IARG_INST_PTR, IARG_END);
        	}
	}
}
/* change
VOID Fini(INT32 code, VOID* v) {
    printf("\n=== REGIONS DUMP ===\n");
    PIN_MutexLock(&e_regions_mutex);
    for (size_t i = 0; i < e_regions.size(); i++) {
        printf("region[%zu]: %p-%p prot=0x%x writed=%d writes=%zu\n", i,
            (void*)e_regions[i].address, (void*)(e_regions[i].address + e_regions[i].size),
            e_regions[i].prot, e_regions[i].writed, e_regions[i].c_writed);
    }
    PIN_MutexUnlock(&e_regions_mutex);
}
*/

int main(int argc, char** argv) {
	PIN_Init(argc, argv);
	PIN_SetSmcSupport(SMC_ENABLE);
	PIN_MutexInit(&e_regions_mutex);
	if (!g_taint.init_shadow_memory()) {
		fprintf(stderr, "Failed to init shadow memory\n");
		return 1;
	}
	tls_key = PIN_CreateThreadDataKey(NULL);
	PIN_AddThreadStartFunction(ThreadStart, NULL);
	PIN_AddThreadFiniFunction(ThreadFini, NULL);
	PIN_AddSyscallEntryFunction(BeforeSyscall, NULL);
	PIN_AddSyscallExitFunction(AfterSyscall, NULL);
	INS_AddInstrumentFunction(InstructionInstrumentation, NULL);
	//PIN_AddFiniFunction(Fini, NULL);
	PIN_StartProgram();
	return 0;
}