-
Notifications
You must be signed in to change notification settings - Fork 0
/
controller.py
450 lines (412 loc) · 21.4 KB
/
controller.py
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
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
import os
import argparse
import time
import random
import json
import os.path as osp
from pathlib import Path
from typing import Tuple, List
from client import setup_client, stop_client, recv_message, sniff
from proto_gen.sniffed_info_pb2 import FlowArray, Flow
from flask import Flask, request
app = Flask(__name__)
app.config["DEBUG"] = True
nodes = {}
edges = {}
def create_missing_directories(output_path):
dirs = osp.split(output_path)[0]
if not osp.exists(dirs):
Path(dirs).mkdir(parents=True, exist_ok=True)
class node: #attributes of a node
def __init__(self, ptype, pname, paddress, pPort, pcolor):
self.type = ptype
self.name = pname
self.address = paddress
self.port = pPort
self.color = pcolor
def __str__(self):
return self.type + ", " + self.name + ", " + self.address + "\n"
def __eq__(self, other):
if self.type == other.type and self.address == other.address:
return True
else:
return False
class edge: #attributes of an edge
def __init__(self, ptype, pTH, pRST, pC, pdomain, prange):
self.type = ptype
self.TH = pTH
self.RST = pRST
self.C = pC
self.domain = pdomain
self.range = prange
def __str__(self):
return self.type + ", " + self.TH + ", " + str(self.RST) + ", " + self.domain + ", " + self.range + "\n"
def __eq__(self, other):
if self.type == other.type and self.range == other.range and self.domain == other.domain:
return True
else:
return False
def reset_global_vars():
global nodes; global edges
nodes.clear(); edges.clear()
def randomColor(): #generate random light colors
return "#"+''.join([(str)(hex(random.randint(80,255))).split("x")[1], (str)(hex(random.randint(80,255))).split("x")[1], (str)(hex(random.randint(80,255))).split("x")[1]])
def generate_graph_from_file(log:str):
data = open(log, "r")#input
newNode1 = None
newNode2 = None
id1 = None
id2 = None
for line in data:#read edge information
if len(line.strip().split(" ")) == 4:
node1, node2, service_type, weight = line.strip().split(" ")#node1: the first node; node2: the second node; weight: information
node12=node1.split(":")[1]#seperate port and ip address
node22=node2.split(":")[1]
if service_type[-1] == 'C':#node 1 is client and node2 is server
newNode1 = node("owl:Class", "Client", node1.split(":")[0], node12, randomColor())
newNode2 = node("owl:equivalentClass", service_type[0:-2], node2.split(":")[0], node22,randomColor())
elif service_type[-1] == 'S':#node 1 is server and node2 is client
newNode1 = node("owl:equivalentClass", service_type[0:-2], (node1.split(":")[0]), node12,randomColor())
newNode2 = node("owl:Class", "Client", node2.split(":")[0], node22, randomColor())
else:#both nodes are not server
newNode1 = node("owl:Class", "Unknown", node1.split(":")[0], node12,randomColor())
newNode2 = node("owl:Class", "Unknown", node2.split(":")[0], node22,randomColor())
if newNode1 not in nodes.values():#if node 1 is not included in graph
id1 = len(nodes)
nodes[len(nodes)] = newNode1#expand node dictionary
for key in nodes:#assign same color to nodes with same IPaddress; connect nodes with same IPaddress with an edged named "same address"
if nodes[key].address == newNode1.address and nodes[key].type == "owl:Class" and newNode1.type == "owl:equivalentClass":
newNode1.color = nodes[key].color
edges[len(edges)] = edge("owl:DatatypeProperty", "same address",0,0, str(key), str(id1))
elif nodes[key].address == newNode1.address and nodes[key].type == "owl:equivalentClass" and newNode1.type == "owl:Class":
newNode1.color = nodes[key].color
edges[len(edges)] = edge("owl:DatatypeProperty", "same address",0,0, str(id1), str(key))
else:
for key in nodes:
if nodes[key] == newNode1:
id1 = key
if nodes[key].name == "Unknown" and newNode1.name == "Client":
nodes[key].name = "Client"
break
if newNode2 not in nodes.values():#if node 2 is not included in graph
id2 = len(nodes)
nodes[len(nodes)] = newNode2#expand node dictionary
for key in nodes:#assign same color to nodes with same IPaddress
if nodes[key].address == newNode2.address and nodes[key].type == "owl:Class" and newNode2.type == "owl:equivalentClass":
newNode2.color = nodes[key].color
edges[len(edges)] = edge("owl:DatatypeProperty", "same address",0,0, str(key), str(id2))
elif nodes[key].address == newNode2.address and nodes[key].type == "owl:equivalentClass" and newNode2.type == "owl:Class":
newNode2.color = nodes[key].color
edges[len(edges)] = edge("owl:DatatypeProperty", "same address",0,0, str(id2), str(key))
else:
for key in nodes:
if nodes[key] == newNode2:
id2 = key
if nodes[key].name == "Unknown" and newNode2.name == "Client":
nodes[key].name = "Client"
break
if "-" in weight:#if edge information contains more than throughput, add more information to the edge
th, rst = weight.split("-", 1)
newEdge = edge("owl:ObjectProperty", th, str(round(float(rst), 3)),1, str(id1), str(id2))
else:#only throughput
th = weight
newEdge = edge("owl:ObjectProperty", th,0,1, str(id1), str(id2))
if newEdge not in edges.values():#add new edge
edges[len(edges)] = newEdge
else:
for key in edges:#if edge in graph, add throughput to existed edge and increment connection
if edges[key] == newEdge:
edges[key].TH = str(int(edges[key].TH) + int(th))
edges[key].C = str(int(edges[key].C) + 1)
if rst:
oRST = float(edges[key].RST)
oC = int(edges[key].C)
edges[key].RST = str(round((oRST*(oC-1) + float(rst))/oC, 3))
rst = 0
newNode1 = None
newNode2 = None
id1 = None
id2 = None
data.close()
def generate_graph(flow_array:FlowArray):
newNode1 = None
newNode2 = None
id1 = None
id2 = None
for flow in flow_array.flows:
if flow.service_type == "Unknown": # neither node is client or server
newNode1 = node("owl:Class", flow.service_type, flow.s_addr, flow.s_port, randomColor())
newNode2 = node("owl:Class", flow.service_type, flow.d_addr, flow.d_port, randomColor())
elif not flow.is_server:#node 1 is client and node2 is server
newNode1 = node("owl:Class", "Client", flow.s_addr, flow.s_port, randomColor())
newNode2 = node("owl:equivalentClass", flow.service_type, flow.d_addr, flow.d_port, randomColor())
else:#node 1 is server and node2 is client
newNode1 = node("owl:equivalentClass", flow.service_type, flow.s_addr, flow.s_port, randomColor())
newNode2 = node("owl:Class", "Client", flow.d_addr, flow.d_port, randomColor())
if newNode1 not in nodes.values():#if node 1 is not included in graph
id1 = len(nodes)
nodes[len(nodes)] = newNode1#expand node dictionary
for key in nodes:#assign same color to nodes with same IPaddress; connect nodes with same IPaddress with an edged named "same address"
if nodes[key].address == newNode1.address and nodes[key].type == "owl:Class" and newNode1.type == "owl:equivalentClass":
newNode1.color = nodes[key].color
edges[len(edges)] = edge("owl:DatatypeProperty", "same address",0,0, str(key), str(id1))
elif nodes[key].address == newNode1.address and nodes[key].type == "owl:equivalentClass" and newNode1.type == "owl:Class":
newNode1.color = nodes[key].color
edges[len(edges)] = edge("owl:DatatypeProperty", "same address",0,0, str(id1), str(key))
else:
for key in nodes:
if nodes[key] == newNode1:
id1 = key
if nodes[key].name == "Unknown" and newNode1.name == "Client":
nodes[key].name = "Client"
break
if newNode2 not in nodes.values():#if node 2 is not included in graph
id2 = len(nodes)
nodes[len(nodes)] = newNode2#expand node dictionary
for key in nodes:#assign same color to nodes with same IPaddress
if nodes[key].address == newNode2.address and nodes[key].type == "owl:Class" and newNode2.type == "owl:equivalentClass":
newNode2.color = nodes[key].color
edges[len(edges)] = edge("owl:DatatypeProperty", "same address",0,0, str(key), str(id2))
elif nodes[key].address == newNode2.address and nodes[key].type == "owl:equivalentClass" and newNode2.type == "owl:Class":
newNode2.color = nodes[key].color
edges[len(edges)] = edge("owl:DatatypeProperty", "same address",0,0, str(id2), str(key))
else:
for key in nodes:
if nodes[key] == newNode2:
id2 = key
if nodes[key].name == "Unknown" and newNode2.name == "Client":
nodes[key].name = "Client"
break
newEdge = edge("owl:ObjectProperty", str(flow.num_bytes), str(round(flow.rst, 3)), 1, str(id1), str(id2))
if newEdge not in edges.values():#add new edge
edges[len(edges)] = newEdge
else:
for key in edges:#if edge in graph, add throughput to existed edge and increment connection
if edges[key] == newEdge:
edges[key].TH = str(int(edges[key].TH) + int(flow.num_bytes))
edges[key].C = str(int(edges[key].C) + 1)
if flow.rst:
oRST = float(edges[key].RST)
oC = int(edges[key].C)
edges[key].RST = str(round((oRST*(oC-1) + float(flow.rst))/oC, 3))
newNode1 = None
newNode2 = None
id1 = None
id2 = None
def write_json_output(fname:str, cmd_mode:bool=False):
print("Producing call graph data")
json_dict = {}
json_dict["class"] = [{"id":str(key), "type":nodes[key].type} for key in nodes]
json_dict["classAttribute"] = [{"id":str(key), "label":nodes[key].name, "comment":{"undefined":nodes[key].address + ", Port: " + nodes[key].port}, "attributes":[nodes[key].color]} for key in nodes]
json_dict["property"] = [{"id":str(key), "type":edges[key].type} for key in edges]
json_dict["propertyAttribute"] = [{"id":str(key), "label":edges[key].TH if edges[key].TH == "same address" else "TH: " + render_readable(int(edges[key].TH)) + ", ", "domain":edges[key].domain, "range":edges[key].range} for key in edges]
for propAtt in json_dict["propertyAttribute"]:
if propAtt["label"] != "same address":
key = int(propAtt["id"])
if float(edges[key].RST):
propAtt["label"] += "RST: " + edges[key].RST
else:
propAtt["label"] += "C: " + str(edges[key].C)
reset_global_vars()
out_path = osp.join("json", fname)
create_missing_directories(out_path)
with open(out_path, "w") as f:
json.dump(json_dict, f, indent=4)
if not cmd_mode:
print("Data sent to local machine")
return json_dict
def render_readable(num:int) -> str:
if num < 10000:
return str(num)
elif num < 100000:
return str(round(num/1000, 1)) + "K"
elif num < 1000000:
return str(int(num/1000)) + "K"
elif num < 10000000:
return str(round(num/1000000, 1)) + "M"
elif num < 1000000000:
return str(int(num/1000000)) + "M"
else:
return str(round(num/1000000000, 0)) + "B"
def equal_flows(new_flow:Flow, flow:Flow) -> bool:
if new_flow.s_addr == flow.s_addr and new_flow.s_port == flow.s_port and \
new_flow.d_addr == flow.d_addr and new_flow.d_port == flow.d_port and \
new_flow.is_server == flow.is_server and new_flow.service_type == flow.service_type:
return True
return False
def next_hop_extractor(new_flows_container, ip:str, visited:List[str], blacklisted_ips:List[str]=[]) -> Tuple[List[str], List[str]]:
ips = []
if type(new_flows_container) is not str:
for flow in new_flows_container.flows:
if flow.s_addr == ip:
new_ip = flow.d_addr
if new_ip not in visited and new_ip not in blacklisted_ips:
ips.append(new_ip)
visited.append(new_ip)
else:
with open(new_flows_container, "r") as f:
for line in f:
if line.split(':')[0] == ip: # if flow has current ip as saddr
new_ip = line.split(' ')[1].split(':')[0]
if new_ip not in visited and new_ip not in blacklisted_ips:
ips.append(new_ip)
visited.append(new_ip)
return (ips, visited)
def run_main(mode:str, host:str, log:str, arg:str, sniff_time:int, out:str, cmd_mode:bool=False):
"""
When monitoring an application, iteratively send IP address to agent and collect flows.
When reading a pcap file, send filename and collect flows. Then in both cases, generate call graph.
Args:
mode (str): The framework usage mode (i for application monitoring or f for reading from file), not to be confused with TCP vs logging mode.
host (str): The address of the machine hosting the agent.
log (str): Indicates the mode in which to send flows (* for TCP mode, else a filename for logging mode).
arg (str): IP address of component to be monitored or name of pcap file to be read.
sniff_time (int): Number of seconds that component at IP address specified or capture file specified in arg will be monitored.
out (str): Name of JSON output file containing call graph information.
cmd_mode (bool): True if using framework in command line mode, False otherwise.
Returns:
dict: JSON object containing call graph information.
"""
# INITIAL START
setup_client(host)
print("Connected")
total_time=0.0
t = time.perf_counter()
f = FlowArray()
log_orig = log
if log != "*":
log_orig = log
log = os.path.join("logs", log)
create_missing_directories(log)
open(log, "w").close()
if mode == "f": # reading from pcap file
sniff(mode, log_orig, arg, sniff_time)
if log == "*":
print("Waiting for flows from agent...")
# STOP - FILE/TCP
response = recv_message()
# START - FILE/TCP
while response is not None:
f.flows.extend(response.flows)
response = recv_message()
print("Received flows from agent")
tg_start = time.perf_counter()
generate_graph(f)
else: # reading from log
print("Waiting for flows to be recorded by agent...")
# Proceed to read from logfile only when sniffer closes connection and sends a
# blank message, indicating it is done writing to logfile
# STOP - FILE/LOG
recv_message()
# START - FILE/LOG
print("Reading flows from file")
tg_start = time.perf_counter()
generate_graph_from_file(log)
else: # sniffing network interface
q, visited, ips = [arg], [arg], []
exists = False
if log == "*": # if using tcp
l = FlowArray()
while len(q) > 0:
print("IP address(es) in queue: {}".format(q))
ip = q.pop(0)
sniff(mode, log_orig, ip, sniff_time)
# STOP - APP/TCP
print("Waiting for flows from agent...")
response = recv_message()#sniffed_info_pb2.FlowArray) # uses protobuf
# START - APP/TCP
while response is not None:
f.flows.extend(response.flows)
response = recv_message()#sniffed_info_pb2.FlowArray) # uses protobuf
print("Received flows from agent")
if len(l.flows) == 0:
l.flows.extend(f.flows)
else:
for new_flow in f.flows:
for flow in l.flows:
exists = False
if equal_flows(new_flow, flow):
exists = True
break
if not exists:
l.flows.append(new_flow)
t1_start = time.perf_counter()
ips, visited = next_hop_extractor(f, ip, visited) # can also specify ips to omit (on blacklist)
end_time=time.perf_counter()-t1_start
total_time=total_time+end_time
q.extend(ips)
del f.flows[:]
if cmd_mode:
stop_client()
tg_start = time.perf_counter()
generate_graph(l)
else: # if using log
temp_log = os.path.join("logs", "temp-log.txt")
lines_to_write = []
while len(q) > 0:
print("IP address(es) in queue: {}".format(q))
ip = q.pop(0)
sniff(mode, log_orig, ip, sniff_time)
print("Waiting for flows to be recorded by agent...")
# STOP - APP/LOG
recv_message()
# START - APP/LOG
with open(log, "r") as l, open(temp_log, "r") as f:
print("Reading flows from file")
f.seek(0)
if os.stat(log).st_size == 0:
lines_to_write.extend(f)
else:
for new_line in f:
for line in l:
exists = False
if ' '.join(new_line.split(' ')[0:3]) in line:
exists = True
break
if not exists:
lines_to_write.append(new_line)
l.seek(0)
with open(log, "a") as l:
l.writelines(lines_to_write)
t1_start = time.perf_counter()
ips, visited = next_hop_extractor(temp_log, ip, visited) # can also specify ips to omit (on blacklist)
end_time=time.perf_counter()-t1_start
total_time=total_time+end_time
q.extend(ips)
lines_to_write.clear()
if cmd_mode:
stop_client()
os.remove(temp_log)
tg_start = time.perf_counter()
generate_graph_from_file(log)
endg_time=time.perf_counter()-tg_start
total_time=total_time+endg_time
#print("Controller time: {} seconds".format(round(total_time, 5)))
print("Elapsed time since controller started: {} seconds".format(round(time.perf_counter() - t, 5)))
# FINAL STOP - this excludes the execution time for write_json_output()
return write_json_output(out, cmd_mode)
@app.route("/run", methods=['GET'])
def run():
mode, log, host, arg, time, out = request.args.values()
return run_main(mode, host, log, arg, int(time), out)
def run_parser():
parser = argparse.ArgumentParser()
group1 = parser.add_mutually_exclusive_group(required=True)
group2 = parser.add_mutually_exclusive_group()
group1.add_argument("-f", "--file", help="read capture file containing flows to be sniffed")
group1.add_argument("-i", "--IP", help="perform live sniffing starting with provided IP")
group2.add_argument("-H", "--host", nargs="?", const='127.0.0.1', help="address for sniffer host. Defaults to localhost")
group2.add_argument("-l", "--log", nargs="?", const="log.txt", default="*", help="send results from sniffer using log file (uses log.txt if no arg). Defaults to sending flows via TCP and omitting a log")
parser.add_argument("-t", "--time", type=int, choices=range(7,1000), metavar="7-1000", default=8, help="sniffing time for each component")
parser.add_argument("-o", "--output", default="out.json", help="name of json output file. Defaults to out.json")
args = parser.parse_args()
mode, arg = None, None
if args.IP is not None:
mode = "i"
arg = args.IP
elif args.file is not None:
mode = "f"
arg = args.file
run_main(mode, args.host, args.log, arg, args.time, args.output, cmd_mode=True)
if __name__ == '__main__':
run_parser()