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 }
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 } )
121+ printing_df ["Upload Speed" ] = printing_df ["Upload Speed" ].apply (get_size ).apply (lambda s : f"{ 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