CFG-Tree-Pin/cfgtree.cpp at main · solidbird/CFG-Tree-Pin · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
#include <iostream>
#include <sstream>
#include <string>
#include <fstream>
#include <vector>

#include "pin.H"
using namespace std;


ofstream OutFile;
map<string, string> bblAddress;
vector<string> nextInstructionAddress;


VOID traceMain(TRACE trace, VOID* v){
	stringstream insBuilder;
	stringstream tmpBuilder;

	ADDRINT bblhead = TRACE_Address(trace);
	ADDRINT nextInstructionAddress = TRACE_Address(trace);
	string cfIns = "";

	IMG img = IMG_FindByAddress(TRACE_Address(trace));
	RTN rtn = TRACE_Rtn(trace);

	if(!IMG_Valid(img) || !(IMG_IsMainExecutable(img))) return;

	RTN_Open(rtn);
	for( BBL bbl = TRACE_BblHead(trace); BBL_Valid(bbl); bbl = BBL_Next(bbl) ){


		cout << "\t|" << endl;
		cout << "\t|" << endl;

		if(bblhead == nextInstructionAddress){
			cout << "\ttrue: " << cfIns << endl;

		}else{
			cout << "\tfalse: " << cfIns << endl;
		}

		cout << "\t|" << endl;
		cout << "\tv" << endl;

		bblhead = BBL_Address(bbl);
		cout << "--- BBL Head Address: " << hex << bblhead << " ---" << endl;

		cout << RTN_Name(rtn) << endl;
		for(INS ins = BBL_InsHead(bbl); INS_Valid(ins); ins = INS_Next(ins) ){
			cout << "\t" << " IP: " << hex << INS_Address(ins) << " IP+1: "  << INS_NextAddress(ins) << "\n\t\t << " << INS_Disassemble(ins) << endl;
			if(INS_IsControlFlow(ins)){
				if(!INS_IsIndirectControlFlow(ins)){

				}
				nextInstructionAddress = INS_NextAddress(ins);
				cfIns = INS_Disassemble(ins);
			}
		}


	}
	RTN_Close(rtn);
}

VOID mainIMG(IMG img, VOID* v){
	if(!IMG_IsMainExecutable(img)) return;

	for(SEC sec = IMG_SecHead(img); SEC_Valid(sec); sec = SEC_Next(sec)){
		if(SEC_Name(sec).find(".text") != string::npos){
			TRACE_AddInstrumentFunction(traceMain, NULL);
			//cout << SEC_Name(sec) << endl;
		}
	}
}

VOID EntryPoint(VOID *){
	cout << "-- Program Start --" << endl;
}

INT32 Usage(){
	cerr << "wrong" << endl;
	return -1;
}

VOID Fini(INT32 code, VOID *v){
	//for(auto x : bblAddress){
	//	cout << x.first << endl;
	//}
	OutFile.setf(ios::showbase);

	for(auto x : bblAddress /*nextInstructionAddress*/){
		//OutFile << x.second << endl;
		//OutFile << "--" << endl;
		//if(x.find(x.first) != 0){
		//}
	}

	OutFile.close();
}


int main(int argc, char* argv[]){

	PIN_InitSymbols();
	InitRegTables();
	if(PIN_Init(argc,argv)) return Usage();

	OutFile.open("data.out");


	PIN_AddApplicationStartFunction(EntryPoint, NULL);
	IMG_AddInstrumentFunction(mainIMG, NULL);
	PIN_AddFiniFunction(Fini, NULL);

	PIN_StartProgram();

	return 0;
}