60b5537f3ead4b026b216267839c4931aa792f92
[collecty.git] / src / collecty / ping.py
1 #!/usr/bin/python
2
3 from __future__ import division
4
5 import array
6 import math
7 import os
8 import random
9 import select
10 import socket
11 import struct
12 import sys
13 import time
14
15 ICMP_TYPE_ECHO_REPLY = 0
16 ICMP_TYPE_ECHO_REQUEST = 8
17 ICMP_MAX_RECV = 2048
18
19 MAX_SLEEP = 1000
20
21 class PingError(Exception):
22         msg = None
23
24
25 class PingResolveError(PingError):
26         msg = "Could not resolve hostname"
27
28
29 class Ping(object):
30         def __init__(self, destination, timeout=1000, packet_size=56):
31                 self.destination = self._resolve(destination)
32                 self.timeout = timeout
33                 self.packet_size = packet_size
34
35                 self.id = os.getpid() & 0xffff # XXX ? Is this a good idea?
36
37                 self.seq_number = 0
38
39                 # Number of sent packets.
40                 self.send_count = 0
41
42                 # Save the delay of all responses.
43                 self.times = []
44
45         def run(self, count=None, deadline=None):
46                 while True:
47                         delay = self.do()
48
49                         self.seq_number += 1
50
51                         if count and self.seq_number >= count:
52                                 break
53
54                         if deadline and self.total_time >= deadline:
55                                 break
56
57                         if delay == None:
58                                 delay = 0
59
60                         if MAX_SLEEP > delay:
61                                 time.sleep((MAX_SLEEP - delay) / 1000)
62
63         def do(self):
64                 s = None
65                 try:
66                         # Open a socket for ICMP communication.
67                         s = socket.socket(socket.AF_INET, socket.SOCK_RAW, socket.getprotobyname("icmp"))
68
69                         # Send one packet.
70                         send_time = self.send_icmp_echo_request(s)
71
72                         # Increase number of sent packets (even if it could not be sent).
73                         self.send_count += 1
74
75                         # If the packet could not be sent, we may stop here.
76                         if send_time is None:
77                                 return
78
79                         # Wait for the reply.
80                         receive_time, packet_size, ip, ip_header, icmp_header = self.receive_icmp_echo_reply(s)
81
82                 finally:
83                         # Close the socket.
84                         if s:
85                                 s.close()
86
87                 # If a packet has been received...
88                 if receive_time:
89                         delay = (receive_time - send_time) * 1000
90                         self.times.append(delay)
91
92                         return delay
93
94         def send_icmp_echo_request(self, s):
95                 # Header is type (8), code (8), checksum (16), id (16), sequence (16)
96                 checksum = 0
97
98                 # Create a header with checksum == 0.
99                 header = struct.pack("!BBHHH", ICMP_TYPE_ECHO_REQUEST, 0,
100                         checksum, self.id, self.seq_number)
101
102                 # Get some bytes for padding.
103                 padding = os.urandom(self.packet_size)
104
105                 # Calculate the checksum for header + padding data.
106                 checksum = self._calculate_checksum(header + padding)
107
108                 # Rebuild the header with the new checksum.
109                 header = struct.pack("!BBHHH", ICMP_TYPE_ECHO_REQUEST, 0,
110                         checksum, self.id, self.seq_number)
111
112                 # Build the packet.
113                 packet = header + padding
114
115                 # Save the time when the packet has been sent.
116                 send_time = time.time()
117
118                 # Send the packet.
119                 try:
120                         s.sendto(packet, (self.destination, 0))
121                 except socket.error, (errno, msg):
122                         if errno == 1: # Operation not permitted
123                                 # The packet could not be sent, probably because of
124                                 # wrong firewall settings.
125                                 return
126
127                 return send_time
128
129         def receive_icmp_echo_reply(self, s):
130                 timeout = self.timeout / 1000.0
131
132                 # Wait until the reply packet arrived or until we hit timeout.
133                 while True:
134                         select_start = time.time()
135
136                         inputready, outputready, exceptready = select.select([s], [], [], timeout)
137                         select_duration = (time.time() - select_start)
138
139                         if inputready == []: # Timeout
140                                 return None, 0, 0, 0, 0
141
142                         # Save the time when the packet has been received.
143                         receive_time = time.time()
144
145                         # Read the packet from the socket.
146                         packet_data, address = s.recvfrom(ICMP_MAX_RECV)
147
148                         # Parse the ICMP header.
149                         icmp_header = self._header2dict(
150                                 ["type", "code", "checksum", "packet_id", "seq_number"],
151                                 "!BBHHH", packet_data[20:28]
152                         )
153
154                         # This is the reply to our packet if the ID matches.
155                         if icmp_header["packet_id"] == self.id:
156                                 # Parse the IP header.
157                                 ip_header = self._header2dict(
158                                         ["version", "type", "length", "id", "flags",
159                                         "ttl", "protocol", "checksum", "src_ip", "dst_ip"],
160                                         "!BBHHHBBHII", packet_data[:20]
161                                 )
162
163                                 packet_size = len(packet_data) - 28
164                                 ip = socket.inet_ntoa(struct.pack("!I", ip_header["src_ip"]))
165
166                                 return receive_time, packet_size, ip, ip_header, icmp_header
167
168                         # Check if the timeout has already been hit.
169                         timeout = timeout - select_duration
170                         if timeout <= 0:
171                                 return None, 0, 0, 0, 0
172
173         def _header2dict(self, names, struct_format, data):
174                 """
175                         Unpack tghe raw received IP and ICMP header informations to a dict
176                 """
177                 unpacked_data = struct.unpack(struct_format, data)
178                 return dict(zip(names, unpacked_data))
179
180         def _calculate_checksum(self, source_string):
181                 if len(source_string) % 2:
182                         source_string += "\x00"
183
184                 converted = array.array("H", source_string)
185                 if sys.byteorder == "big":
186                         converted.byteswap()
187
188                 val = sum(converted)
189
190                 # Truncate val to 32 bits (a variance from ping.c, which uses signed
191                 # integers, but overflow is unlinkely in ping).
192                 val &= 0xffffffff
193
194                 # Add high 16 bits to low 16 bits.
195                 val = (val >> 16) + (val & 0xffff)
196
197                 # Add carry from above (if any).
198                 val += (val >> 16)
199
200                 # Invert and truncate to 16 bits.
201                 answer = ~val & 0xffff
202
203                 return socket.htons(answer)
204
205         def _resolve(self, host):
206                 """
207                         Resolve host.
208                 """
209                 if self._is_valid_ipv4_address(host):
210                         return host
211
212                 try:
213                         return socket.gethostbyname(host)
214                 except socket.gaierror as e:
215                         if e.errno == -3:
216                                 raise PingResolveError
217
218                         raise
219
220         def _is_valid_ipv4_address(self, addr):
221                 """
222                         Check addr to be a valid IPv4 address.
223                 """
224                 parts = addr.split(".")
225
226                 if not len(parts) == 4:
227                         return False
228
229                 for part in parts:
230                         try:
231                                 number = int(part)
232                         except ValueError:
233                                 return False
234
235                         if number > 255:
236                                 return False
237
238                 return True
239
240         @property
241         def receive_count(self):
242                 """
243                         The number of received packets.
244                 """
245                 return len(self.times)
246
247         @property
248         def total_time(self):
249                 """
250                         The total time of all roundtrips.
251                 """
252                 try:
253                         return sum(self.times)
254                 except ValueError:
255                         return
256
257         @property
258         def min_time(self):
259                 """
260                         The smallest roundtrip time.
261                 """
262                 try:
263                         return min(self.times)
264                 except ValueError:
265                         return
266
267         @property
268         def max_time(self):
269                 """
270                         The biggest roundtrip time.
271                 """
272                 try:
273                         return max(self.times)
274                 except ValueError:
275                         return
276
277         @property
278         def avg_time(self):
279                 """
280                         Calculate the average response time.
281                 """
282                 try:
283                         return self.total_time / self.receive_count
284                 except ZeroDivisionError:
285                         return
286
287         @property
288         def variance(self):
289                 """
290                         Calculate the variance of all roundtrips.
291                 """
292                 if self.avg_time is None:
293                         return
294
295                 var = 0
296
297                 for t in self.times:
298                         var += (t - self.avg_time) ** 2
299
300                 var /= self.receive_count
301                 return var
302
303         @property
304         def stddev(self):
305                 """
306                         Standard deviation of all roundtrips.
307                 """
308                 return math.sqrt(self.variance)
309
310         @property
311         def loss(self):
312                 """
313                         Outputs the percentage of dropped packets.
314                 """
315                 dropped = self.send_count - self.receive_count
316
317                 return dropped / self.send_count
318
319
320 if __name__ == "__main__":
321         p = Ping("ping.ipfire.org")
322         p.run(count=5)
323
324         print "Min/Avg/Max/Stddev: %.2f/%.2f/%.2f/%.2f" % \
325                 (p.min_time, p.avg_time, p.max_time, p.stddev)
326         print "Sent/Recv/Loss: %d/%d/%.2f" % (p.send_count, p.receive_count, p.loss)