1
+ from scapy .all import *
2
+ import psutil
3
+ from collections import defaultdict
4
+ import os
5
+ from threading import Thread
6
+ import pandas as pd
7
+
8
+ # get the all network adapter's MAC addresses
9
+ all_macs = {iface .mac for iface in ifaces .values ()}
10
+ # A dictionary to map each connection to its correponding process ID (PID)
11
+ connection2pid = {}
12
+ # A dictionary to map each process ID (PID) to total Upload (0) and Download (1) traffic
13
+ pid2traffic = defaultdict (lambda : [0 , 0 ])
14
+ # the global Pandas DataFrame that's used to track previous traffic stats
15
+ global_df = None
16
+ # global boolean for status of the program
17
+ is_program_running = True
18
+
19
+ def get_size (bytes ):
20
+ """
21
+ Returns size of bytes in a nice format
22
+ """
23
+ for unit in ['' , 'K' , 'M' , 'G' , 'T' , 'P' ]:
24
+ if bytes < 1024 :
25
+ return f"{ bytes :.2f} { unit } B"
26
+ bytes /= 1024
27
+
28
+
29
+ def process_packet (packet ):
30
+ global pid2traffic
31
+ try :
32
+ # get the packet source & destination IP addresses and ports
33
+ packet_connection = (packet .sport , packet .dport )
34
+ except (AttributeError , IndexError ):
35
+ # sometimes the packet does not have TCP/UDP layers, we just ignore these packets
36
+ pass
37
+ else :
38
+ # get the PID responsible for this connection from our `connection2pid` global dictionary
39
+ packet_pid = connection2pid .get (packet_connection )
40
+ if packet_pid :
41
+ if packet .src in all_macs :
42
+ # the source MAC address of the packet is our MAC address
43
+ # so it's an outgoing packet, meaning it's upload
44
+ pid2traffic [packet_pid ][0 ] += len (packet )
45
+ else :
46
+ # incoming packet, download
47
+ pid2traffic [packet_pid ][1 ] += len (packet )
48
+
49
+
50
+ def get_connections ():
51
+ """A function that keeps listening for connections on this machine
52
+ and adds them to `connection2pid` global variable"""
53
+ global connection2pid
54
+ while is_program_running :
55
+ # using psutil, we can grab each connection's source and destination ports
56
+ # and their process ID
57
+ for c in psutil .net_connections ():
58
+ if c .laddr and c .raddr and c .pid :
59
+ # if local address, remote address and PID are in the connection
60
+ # add them to our global dictionary
61
+ connection2pid [(c .laddr .port , c .raddr .port )] = c .pid
62
+ connection2pid [(c .raddr .port , c .laddr .port )] = c .pid
63
+ # sleep for a second, feel free to adjust this
64
+ time .sleep (1 )
65
+
66
+
67
+ def print_pid2traffic ():
68
+ global global_df
69
+ # initialize the list of processes
70
+ processes = []
71
+ for pid , traffic in pid2traffic .items ():
72
+ # `pid` is an integer that represents the process ID
73
+ # `traffic` is a list of two values: total Upload and Download size in bytes
74
+ try :
75
+ # get the process object from psutil
76
+ p = psutil .Process (pid )
77
+ except psutil .NoSuchProcess :
78
+ # if process is not found, simply continue to the next PID for now
79
+ continue
80
+ # get the name of the process, such as chrome.exe, etc.
81
+ name = p .name ()
82
+ # get the time the process was spawned
83
+ try :
84
+ create_time = datetime .fromtimestamp (p .create_time ())
85
+ except OSError :
86
+ # system processes, using boot time instead
87
+ create_time = datetime .fromtimestamp (psutil .boot_time ())
88
+ # construct our dictionary that stores process info
89
+ process = {
90
+ "pid" : pid , "name" : name , "create_time" : create_time , "Upload" : traffic [0 ],
91
+ "Download" : traffic [1 ],
92
+ }
93
+ try :
94
+ # calculate the upload and download speeds by simply subtracting the old stats from the new stats
95
+ process ["Upload Speed" ] = traffic [0 ] - global_df .at [pid , "Upload" ]
96
+ process ["Download Speed" ] = traffic [1 ] - global_df .at [pid , "Download" ]
97
+ except (KeyError , AttributeError ):
98
+ # If it's the first time running this function, then the speed is the current traffic
99
+ # You can think of it as if old traffic is 0
100
+ process ["Upload Speed" ] = traffic [0 ]
101
+ process ["Download Speed" ] = traffic [1 ]
102
+ # append the process to our processes list
103
+ processes .append (process )
104
+ # construct our Pandas DataFrame
105
+ df = pd .DataFrame (processes )
106
+ try :
107
+ # set the PID as the index of the dataframe
108
+ df = df .set_index ("pid" )
109
+ # sort by column, feel free to edit this column
110
+ df .sort_values ("Download" , inplace = True , ascending = False )
111
+ except KeyError as e :
112
+ # when dataframe is empty
113
+ pass
114
+ # make another copy of the dataframe just for fancy printing
115
+ printing_df = df .copy ()
116
+ try :
117
+ # apply the function get_size to scale the stats like '532.6KB/s', etc.
118
+ printing_df ["Download" ] = printing_df ["Download" ].apply (get_size )
119
+ printing_df ["Upload" ] = printing_df ["Upload" ].apply (get_size )
120
+ printing_df ["Download Speed" ] = printing_df ["Download Speed" ].apply (get_size ).apply (lambda s : f"{ s } /s" )
121
+ printing_df ["Upload Speed" ] = printing_df ["Upload Speed" ].apply (get_size ).apply (lambda s : f"{ s } /s" )
122
+ except KeyError as e :
123
+ # when dataframe is empty again
124
+ pass
125
+ # clear the screen based on your OS
126
+ os .system ("cls" ) if "nt" in os .name else os .system ("clear" )
127
+ # print our dataframe
128
+ print (printing_df .to_string ())
129
+ # update the global df to our dataframe
130
+ global_df = df
131
+
132
+
133
+ def print_stats ():
134
+ """Simple function that keeps printing the stats"""
135
+ while is_program_running :
136
+ time .sleep (1 )
137
+ print_pid2traffic ()
138
+
139
+
140
+
141
+ if __name__ == "__main__" :
142
+ # start the printing thread
143
+ printing_thread = Thread (target = print_stats )
144
+ printing_thread .start ()
145
+ # start the get_connections() function to update the current connections of this machine
146
+ connections_thread = Thread (target = get_connections )
147
+ connections_thread .start ()
148
+ # start sniffing
149
+ print ("Started sniffing" )
150
+ sniff (prn = process_packet , store = False )
151
+ # setting the global variable to False to exit the program
152
+ is_program_running = False
153
+
154
+
0 commit comments