import r2pipe
import pygraphviz as pgv
import networkx as nx
import json
import matplotlib.pyplot as plt
from networkx.drawing.nx_pydot import write_dot
import random
import sys
import os

if len(sys.argv) < 2:
    print("Usage: python script.py <argument>")
    sys.exit(1)
binary_path = sys.argv[1]
#fs = sys.argv[2]

r = r2pipe.open(binary_path)
print(binary_path)
##############################################################

r.cmd('aaa')

ARM32_SYSCALLS = {
0 : "restart_syscall",
1 : "exit",
2 : "fork",
3 : "read",
4 : "write",
5 : "open",
6 : "close",
7 : "not implemented",
8 : "creat",
9 : "link",
10 : "unlink",
11 : "execve",
12 : "chdir",
13 : "not implemented",
14 : "mknod",
15 : "chmod",
16 : "lchown",
17 : "not implemented",
18 : "not implemented",
19 : "lseek",
20 : "getpid",
21 : "mount",
22 : "not implemented",
23 : "setuid",
24 : "getuid",
25 : "not implemented",
26 : "ptrace",
27 : "alarm",
28 : "not implemented",
29 : "pause",
30 : "utime",
31 : "stty",
32 : "gtty",
33 : "access",
34 : "nice",
35 : "ftime",
36 : "sync",
37 : "kill",
38 : "rename",
39 : "mkdir",
40 : "rmdir",
41 : "dup",
42 : "pipe",
43 : "times",
44 : "prof",
45 : "brk",
46 : "setgid",
47 : "getgid",
48 : "signal",
49 : "geteuid",
50 : "getegid",
51 : "acct",
52 : "umount",
53 : "lock",
54 : "ioctl",
55 : "fcntl",
56 : "mpx",
57 : "setpgid",
58 : "ulimit",
59 : "olduname",
60 : "umask",
61 : "chroot",
62 : "ustat",
63 : "dup2",
64 : "getppid",
65 : "getpgrp",
66 : "setsid",
67 : "sigaction",
68 : "sgetmask",
69 : "ssetmask",
70 : "setreuid",
71 : "setregid",
72 : "sigsuspend",
73 : "sigpending",
74 : "sethostname",
75 : "setrlimit",
76 : "getrlimit",
77 : "getrusage",
78 : "gettimeofday",
79 : "settimeofday",
80 : "getgroups",
81 : "setgroups",
82 : "select",
83 : "symlink",
84 : "oldlstat",
85 : "readlink",
86 : "uselib",
87 : "swapon",
88 : "reboot",
89 : "readdir",
90 : "mmap",
91 : "munmap",
92 : "truncate",
93 : "ftruncate",
94 : "fchmod",
95 : "fchown",
96 : "getpriority",
97 : "setpriority",
98 : "profil",
99 : "statfs",
100 : "fstatfs",
101 : "ioperm",
102 : "socketcall",
103 : "syslog",
104 : "setitimer",
105 : "getitimer",
106 : "stat",
107 : "lstat",
108 : "fstat",
109 : "olduname",
110 : "iopl",
111 : "vhangup",
112 : "idle",
113 : "vm86",
114 : "wait4",
115 : "swapoff",
116 : "sysinfo",
117 : "ipc",
118 : "fsync",
119 : "sigreturn",
120 : "clone",
121 : "setdomainname",
122 : "uname",
123 : "modify_ldt",
124 : "adjtimex",
125 : "mprotect",
126 : "sigprocmask",
127 : "create_module",
128 : "init_module",
129 : "delete_module",
130 : "get_kernel_syms",
131 : "quotactl",
132 : "getpgid",
133 : "fchdir",
134 : "bdflush",
135 : "sysfs",
136 : "personality",
137 : "afs_syscall",
138 : "setfsuid",
139 : "setfsgid",
140 : "_llseek",
141 : "getdents",
142 : "newselect",
143 : "flock",
144 : "msync",
145 : "readv",
146 : "writev",
147 : "getsid",
148 : "fdatasync",
149 : "_sysctl",
150 : "mlock",
151 : "munlock",
152 : "mlockall",
153 : "munlockall",
154 : "sched_setparam",
155 : "sched_getparam",
156 : "sched_setscheduler",
157 : "sched_getscheduler",
158 : "sched_yield",
159 : "sched_get_priority_max",
160 : "sched_get_priority_min",
161 : "sched_rr_get_interval",
162 : "nanosleep",
163 : "mremap",
164 : "setresuid",
165 : "getresuid",
166 : "vm86old",
167 : "query_module",
168 : "poll",
169 : "nfsservctl",
170 : "setresgid",
171 : "getresgid",
172 : "prctl",
173 : "rt_sigreturn",
174 : "rt_sigaction",
175 : "rt_sigprocmask",
176 : "rt_sigpending",
177 : "rt_sigtimedwait",
178 : "rt_sigqueueinfo",
179 : "rt_sigsuspend",
180 : "pread64",
181 : "pwrite64",
182 : "chown",
183 : "getcwd",
184 : "capget",
185 : "capset",
186 : "sigaltstack",
187 : "sendfile",
188 : "getpmsg",
189 : "putpmsg",
190 : "vfork",
191 : "ugetrlimit",
192 : "mmap2",
193 : "truncate64",
194 : "ftruncate64",
195 : "stat64",
196 : "lstat64",
197 : "fstat64",
198 : "lchown32",
199 : "getuid32",
200 : "getgid32",
201 : "geteuid32",
202 : "getegid32",
203 : "setreuid32",
204 : "setregid32",
205 : "getgroups32",
206 : "setgroups32",
207 : "fchown32",
208 : "setresuid32",
209 : "getresuid32",
210 : "setresgid32",
211 : "getresgid32",
212 : "chown32",
213 : "setuid32",
214 : "setgid32",
215 : "setfsuid32",
216 : "setfsgid32",
217 : "pivot_root",
218 : "mincore",
219 : "madvise",
220 : "madvise1",
221 : "getdents64",
222 : "fcntl64",
223 : "not implemented",
224 : "gettid",
225 : "readahead",
226 : "setxattr",
227 : "lsetxattr",
228 : "fsetxattr",
229 : "getxattr",
230 : "lgetxattr",
231 : "fgetxattr",
232 : "listxattr",
233 : "llistxattr",
234 : "flistxattr",
235 : "removexattr",
236 : "lremovexattr",
237 : "fremovexattr",
238 : "tkill",
239 : "sendfile64",
240 : "futex",
241 : "sched_setaffinity",
242 : "sched_getaffinity",
243 : "io_setup",
244 : "io_destroy",
245 : "io_getevents",
246 : "io_submit",
247 : "io_cancel",
248 : "exit_group",
249 : "lookup_dcookie",
250 : "epoll_create",
251 : "epoll_ctl",
252 : "epoll_wait",
253 : "remap_file_pages",
254 : "not implemented",
255 : "not implemented",
256 : "set_tid_address",
257 : "timer_create",
258 : "timer_settime",
259 : "timer_gettime",
260 : "timer_getoverrun",
261 : "timer_delete",
262 : "clock_settime",
263 : "clock_gettime",
264 : "clock_getres",
265 : "clock_nanosleep",
266 : "statfs64",
267 : "fstatfs64",
268 : "tgkill",
269 : "utimes",
270 : "arm_fadvise64_64",
271 : "pciconfig_iobase",
272 : "pciconfig_read",
273 : "pciconfig_write",
274 : "mq_open",
275 : "mq_unlink",
276 : "mq_timedsend",
277 : "mq_timedreceive",
278 : "mq_notify",
279 : "mq_getsetattr",
280 : "waitid",
281 : "socket",
282 : "bind",
283 : "connect",
284 : "listen",
285 : "accept",
286 : "getsockname",
287 : "getpeername",
288 : "socketpair",
289 : "send",
290 : "sendto",
291 : "recv",
292 : "recvfrom",
293 : "shutdown",
294 : "setsockopt",
295 : "getsockopt",
296 : "sendmsg",
297 : "recvmsg",
298 : "semop",
299 : "semget",
300 : "semctl",
301 : "msgsnd",
302 : "msgrcv",
303 : "msgget",
304 : "msgctl",
305 : "shmat",
306 : "shmdt",
307 : "shmget",
308 : "shmctl",
309 : "add_key",
310 : "request_key",
311 : "keyctl",
312 : "semtimedop",
313 : "vserver",
314 : "ioprio_set",
315 : "ioprio_get",
316 : "inotify_init",
317 : "inotify_add_watch",
318 : "inotify_rm_watch",
319 : "mbind",
320 : "get_mempolicy",
321 : "set_mempolicy",
322 : "openat",
323 : "mkdirat",
324 : "mknodat",
325 : "fchownat",
326 : "futimesat",
327 : "fstatat64",
328 : "unlinkat",
329 : "renameat",
330 : "linkat",
331 : "symlinkat",
332 : "readlinkat",
333 : "fchmodat",
334 : "faccessat",
335 : "pselect6",
336 : "ppoll",
337 : "unshare",
338 : "set_robust_list",
339 : "get_robust_list",
340 : "splice",
341 : "sync_file_range",
342 : "tee",
343 : "vmsplice",
344 : "move_pages",
345 : "getcpu",
346 : "epoll_pwait",
347 : "kexec_load",
348 : "utimensat",
349 : "signalfd",
350 : "timerfd_create",
351 : "eventfd",
352 : "fallocate",
353 : "timerfd_settime",
354 : "timerfd_gettime",
355 : "signalfd4",
356 : "eventfd2",
357 : "epoll_create1",
358 : "dup3",
359 : "pipe2",
360 : "inotify_init1",
361 : "preadv",
362 : "pwritev",
363 : "rt_tgsigqueueinfo",
364 : "perf_event_open",
365 : "recvmmsg",
366 : "accept4",
367 : "fanotify_init",
368 : "fanotify_mark",
369 : "prlimit64",
370 : "name_to_handle_at",
371 : "open_by_handle_at",
372 : "clock_adjtime",
373 : "syncfs",
374 : "sendmmsg",
375 : "setns",
376 : "process_vm_readv",
377 : "process_vm_writev",
378 : "kcmp",
379 : "finit_module",
380 : "sched_setattr",
381 : "sched_getattr",
382 : "renameat2",
383 : "seccomp",
384 : "getrandom",
385 : "memfd_create",
386 : "bpf",
387 : "execveat",
388 : "userfaultfd",
389 : "membarrier",
390 : "mlock2",
391 : "copy_file_range",
392 : "preadv2",
393 : "pwritev2",
394 : "pkey_mprotect",
395 : "pkey_alloc",
396 : "pkey_free",
397 : "statx",
983041 : "ARM_breakpoint",
983042 : "ARM_cacheflush",
983043 : "ARM_usr26",
983044 : "ARM_usr32",
983045 : "ARM_set_tls"
}

#########################################

arch = r.cmdj("ij")
#if arch["bin"]["arch"] != "arm":     #or arm32
if arch.get("bin", {}).get("arch") != "arm":  # or arm32
	file_path = binary_path
	if os.path.exists(file_path):
		os.remove(file_path)
		print(f'{file_path} has been deleted')
	else:
		print(f'{file_path} does not exist')
	sys.exit()

#########################################

svcList1 = []
svcList2 = []
svcFinal = []
svcs = r.cmdj('/ad/j svc [0-9a-fA-F]+')   # svcs = r.cmdj('/atj swi')

for s in svcs:
	s1 = s['code'].split(" ")
	s2 = str(s1[1])
	
	if s2.startswith("0x9"):
		svcList1.append(s)
		
	if len(s2) <= 2:
		svcList2.append(s)
		
class Syscall:
	offset = 0
	hex = 0     
	syscall = ""
	funOffset = 0
	
#########################################

for sf in svcList2:
	addm = sf['offset'] - 4
	addm2 = r.cmdj(f"pdj 1 @ {sf['offset']}-4")
	addm3 = str(addm2[0]['disasm'])
	afoj = r.cmdj(f"afoj @ {sf['offset']}")
	funAddr = afoj.get("address")
	if addm3.startswith("mov r7") and funAddr is not None:
		hex1 = addm3.split(" ")[2]
		sysc = Syscall()
		sysc.hex = hex1
		sysc.offset = sf['offset']
		dec = int(hex1, 16)
		syscall = ARM32_SYSCALLS.get(dec, f"Unknown syscall {dec}")   
		sysc.syscall = syscall
		sysc.funOffset = funAddr
		svcFinal.append(sysc)


for s in svcList1:
	hex2 = s['code']
	pos = hex2.find("9")
	if pos != -1:
		result = hex2[pos + 1:]
	else:
		result = ""
	dec = int(result, 16)
	hex3 = hex(dec)
	svcObj = Syscall()
	svcObj.offset = s['offset']
	svcObj.hex = hex3
	sysname = ARM32_SYSCALLS.get(dec, f"Unknown syscall {dec}")
	svcObj.syscall = sysname
	afoj = r.cmdj(f"afoj @ {s['offset']}")
	funAddr = afoj.get("address")
	svcObj.funOffset = funAddr
	svcFinal.append(svcObj)
	
#########################################

class StrC:
	text = ""
	funAddr = []
	
	#def __repr__(self):
	#	return (self.tolist())
	
	#def	to_json(self):
	#	return {"text": self.text, "funAddr": self.funAddr}

allStrings = []

strings = r.cmdj('izj')
#f = open("stringsFile.txt", "w")
for st in strings:
	funRefs = []
	#f.write(s['string'] + " " + str(s['vaddr']) + "\n")
	refs1 = r.cmdj(f"axtj {st['vaddr']}")
	if refs1:   # and s['string'] not in allString['text']
		strObj = StrC()
		strObj.text = st['string']
		arrAddr = []
		for rf in refs1:
			if rf.get("fcn_addr"):
				if rf['fcn_addr'] not in arrAddr:
					arrAddr.append(rf['fcn_addr'])
		strObj.funAddr = arrAddr
		arrAddr = []
		allStrings.append(strObj)

#function name + offset

#########################################
callgraph = []
filterSvcs = []
contArr = []
strsArray = []

agCj = r.cmdj('agCj')

class Funx:
	name = ""
	offset = 0
	content = []
	stringsUsed = []
	imports = []
	
	def __repr__(self):
		return (self.tolist())
	
	def	to_json(self):
		return {"name": self.name, "offset": self.offset, "content": self.content, "imports": self.imports, "stringsUsed": self.stringsUsed}

#counterX = 0

for f in agCj:
	try:
		agfj = r.cmdj(f"agfj @ {f['name']}")
	except:
		print(f"agfj @ {counterX} has failed!")
		#counterX += 1
		continue
	funObj = Funx()
	funObj.name = agfj[0]['name']
	funObj.offset = agfj[0]['offset']
	for obj in svcFinal:
		if str(obj.funOffset) == str(agfj[0]['offset']):
			filterSvcs.append(obj)
	if len(filterSvcs) != 0:
		for o in filterSvcs:
			contArr.append(o.syscall)
	funObj.content = contArr
	contArr = []
	funObj.imports = f['imports']
	for k in allStrings:
		if agfj[0]['offset'] in k.funAddr:
			strsArray.append(k.text)
	funObj.stringsUsed = strsArray
	strsArray = []
	callgraph.append(funObj.__dict__)
	filterSvcs = []
	#counterX += 1

####################################################
with open(f"json{binary_path}.json", "w") as file:
	json.dump(callgraph, file, indent=4) 

callgraphjson = json.dumps(callgraph)
data = json.loads(callgraphjson)
####################################################

#G = nx.Graph()
G = nx.DiGraph()

for d in data:
	if len(d["content"])>0:
		G.add_node(d["name"], content=d["content"], color='b')  #strings=d["stringsUsed"]
	else:
		G.add_node(d["name"], content="/", color='c')  #strings=d["stringsUsed"]
	if d.get("stringsUsed"):     #if len(d["stringsUsed"]) != 0:
		for p in d["stringsUsed"]:
			rand = random.randint(1, 1000)
			G.add_node(f"{p}"+d["name"]+f"{rand}", content=f"{p}", color='r') 
			G.add_edge(d["name"], f"{p}"+d["name"]+f"{rand}")

for d in data:
	for imp in d.get("imports", []):		#d["imports"]:
		if not G.has_edge(d["name"], imp):
			if G.has_node(imp):
				G.add_edge(d["name"], imp)

plt.figure(figsize=(8, 6))
pos = nx.spring_layout(G)
#nx.draw(G, pos, with_labels=True, labels=nx.get_node_attributes(G, 'content'), node_size=700, node_color='skyblue', font_size=15)
#plt.show()

for node in G.nodes():
    if 'color' not in G.nodes[node]:
        G.nodes[node]['color'] = 'gray'

node_colors = [G.nodes[node].get('color') for node in G.nodes()]
#filtered = [m for m in node_colors if m == 'b' or m == 'r']

subG = G.subgraph([node for node, attrs in G.nodes(data=True) if attrs.get('color') != 'r']).copy()

f1 = open(f"bc--{binary_path}.txt", "a")
centrality = nx.betweenness_centrality(subG, normalized=True)    
for node, centrality_value in centrality.items():
	if subG.nodes[node]['content'] != '/':    ##
	        f1.write(f"Node {subG.nodes[node]['content']}: Betweenness Centrality = {centrality_value} \n")
f1.close()

f2 = open(f"cc--{binary_path}.txt", "a")
closeness = nx.closeness_centrality(subG)    
for node, closeness_value in closeness.items():
	if subG.nodes[node]['content'] != '/':     ##
		f2.write(f"Node {subG.nodes[node]['content']}: Closeness Centrality = {closeness_value} \n")         
f2.close()

f3 = open(f"pr--{binary_path}.txt", "a")
degree = nx.degree_centrality(subG)   
for node, degree_value in degree.items():
	if subG.nodes[node]['content'] != '/':     ##
		f3.write(f"Node {subG.nodes[node]['content']}: Degree Centrality = {degree_value} \n")           
f3.close()

f4 = open(f"dc--{binary_path}.txt", "a")
pagerank_scores = nx.pagerank(subG, alpha=0.9)    
for node, score in pagerank_scores.items():
	if subG.nodes[node]['content'] != '/':     ##
		f4.write(f"Node: {subG.nodes[node]['content']}, PageRank Score: {score} \n")
f4.close()

if subG.number_of_edges() > 0:
	f5 = open(f"lc--{binary_path}.txt", "a")
	communities = nx.community.louvain_communities(subG)
	for iter, community in enumerate(communities):
		f5.write(f"Community {iter+1}: \n")   ##{community}
		for nodeComm in community:
			if subG.nodes[nodeComm]['content'] != '/':
				f5.write(f"Nodes: {subG.nodes[nodeComm]['content']} \n")
	f5.close()

#pos = nx.shell_layout(G)
##nx.draw_networkx(G, pos, with_labels=True, labels=nx.get_node_attributes(G, 'content'), node_size=700, node_color = node_colors, font_size=15)
##plt.title("Functions and Strings Linked in a Graph")
##plt.show()  
#nx.write_gml(G, f"GML{binary_path}.gml")
#write_dot(G, f"DOT{binary_path}.dot")
nld = nx.node_link_data(G)
jsonG = json.dumps(nld)
with open(f"graph{binary_path}.json", 'w') as fg:
    json.dump(jsonG, fg)

print("#############END###############")
sys.exit()

#####################################################
#file_path = binary_path
#if os.path.exists(file_path):
#    os.remove(file_path)
#    print(f'{file_path} has been deleted')
#else:
#    print(f'{file_path} does not exist')
#############################################

