diff --git a/pingplot.py b/pingplot.py index c590e0d..e6da8c1 100644 --- a/pingplot.py +++ b/pingplot.py @@ -1,23 +1,20 @@ -import datetime -import os -import re +import signal import sys import time -from optparse import OptionParser - +import os import numpy as np +import datetime +import re +import argparse +from optparse import OptionParser -# software version __version__ = "1.1.0" ping_flag = 'n' if sys.platform != 'win32': ping_flag = 'c' - -# ping def pinger(host, n): - """Executes the PCs ping command""" if sys.platform != 'win32': proc = os.popen(f"ping -{ping_flag} {n} {host}") else: @@ -26,50 +23,37 @@ def pinger(host, n): proc.close() return result - -# wrapper for ping def call_pinger(host, n, ping, loss, t): - """Calls the pinger function and returns results as arrays""" out = pinger(host, n) try: if sys.platform == 'win32': - loss_idx = float(re.search("\d+(?=% loss)", out).group(0)) - ping_idx = float(re.search("(?<=Average =) \d+", out).group(0)) + loss_idx = float(re.search(r"\d+(?=% loss)", out).group(0)) + ping_idx = float(re.search(r"(?<=Average =) \d+", out).group(0)) + + else: - # the next two lines assume this format: - # 4 packets transmitted, 4 received, 0% packet loss, time 3002ms - # rtt min/avg/max/mdev = 24.146/63.155/128.436/42.823 ms - loss_idx = float(re.search("\d+\.\d+(?=% packet loss)", out).group(0)) + loss_idx = float(re.search(r"\d+\.\d+(?=% packet loss)", out).group(0)) ping_idx = float(out.split('/')[-3]) except: - ping_idx = np.nan # bad connection + ping_idx = np.nan loss_idx = 100. - - # append data ping = np.append(ping, ping_idx) loss = np.append(loss, loss_idx) t = np.append(t, time.time()) return ping, loss, t, out - -# writes out to the log file def write_log(log_file, log_body): - """Writes results to a log file""" log_file.write(f"TIME: {datetime.datetime.now().ctime()}\n{log_body}\n\n") - -# produces ping vs time plot def plot_gen(ping, now, nans, host, interactive=False, size="1280x640"): - """Generates ping vs time plot""" if not interactive: import matplotlib - matplotlib.use("Agg") # no need to load gui toolkit, can run headless + matplotlib.use("Agg") import matplotlib.pyplot as plt - size = [int(dim) for dim in size.split('x')] datestr = now[0].ctime().split() datestr = datestr[0] + " " + datestr[1] + " " + datestr[2] + " " + datestr[-1] - plt.figure(figsize=(size[0] / 80., size[1] / 80.)) # dpi is 80 + plt.figure(figsize=(size[0] / 80., size[1] / 80.)) plt.plot(now[~nans], ping[~nans], drawstyle='steps') plt.title(f"Ping Results for {host}") plt.ylabel("Latency [ms]") @@ -77,8 +61,6 @@ def plot_gen(ping, now, nans, host, interactive=False, size="1280x640"): plt.xticks(size=10) plt.yticks(size=10) plt.ylim(ping[~nans].min() - 5, ping[~nans].max() + 5) - - # plot packet losses start = [] finish = [] for i in range(len(nans)): @@ -87,46 +69,45 @@ def plot_gen(ping, now, nans, host, interactive=False, size="1280x640"): start.append(i) if i == len(nans) - 1 or nans[i + 1] != nans[i]: finish.append(i) - - # add the red bars for bad pings for i in range(len(start)): plt.axvspan(now[start[i]], now[finish[i]], color='red') return plt +def signal_handler(signum, frame, log_file=None, ping=None, now=None, nans=None, opts=None): + print("\nSIGTERM received, generating output files before exiting...") + if opts.log and log_file is not None: + log_file.write(f"\nGraceful shutdown at {datetime.datetime.now().ctime()}\n") + log_file.close() + print(f"Saved log file {log_file.name}") + if opts.fsave and ping is not None: + if len(ping[~nans]) == 0: + print("Error: cannot generate plot; no data collected.") + else: + plt = plot_gen(ping, now, nans, opts.host, opts.plot, opts.size) + plot_name = f"pingplot_v{__version__}_{opts.host}_{datetime.datetime.now().isoformat()}.png" + plt.savefig(plot_name) + print(f"Saved plot {plot_name}") + sys.exit(0) -# main def main(argv=None): - # for interactive mode if not argv: argv = sys.argv[1:] - # handle cmd line arguments - parser = OptionParser() - parser.add_option("-p", "--plot", dest="plot", action="store_true", - help="generates plot after data collection is finished") - parser.add_option("-f", "--file", dest="fsave", action="store_true", - help="save plot to file in the current directory") - parser.add_option("-H", "--host", dest="host", default="google.com", - help="the url or ip address to ping [default: %default]") - parser.add_option("-n", "--num", dest="n", default=1, type="int", - help="the number of packets to send on each ping iteration [default: %default]") - parser.add_option("-t", "--dt", dest="dt", default=0.5, type="float", - help="the time interval (seconds) in which successive pings are sent [default: %default s]") - parser.add_option("-l", "--log", dest="log", action="store_true", - help="save a log file of the event in the current directory") - parser.add_option("-s", "--size", dest="size", default="1280x640", - help="If plotting/saving a plot, this is the plot's dimensions" \ - "in pixels (at 80 DPI) in the format XxY [default: 1280x640]") + parser = argparse.ArgumentParser(description="PingPlot - Ping host and generate ping vs time plot.") + parser.add_argument("-p", "--plot", action="store_true", help="Generate plot after data collection is finished.") + parser.add_argument("-f", "--file", dest="fsave", action="store_true", help="Save plot to file in the current directory.") + parser.add_argument("-H", "--host", default="google.com", help="Hostname to ping.") + parser.add_argument("-n", "--num", default=1, type=int, help="Number of packets to send.") + parser.add_argument("-t", "--dt", default=0.5, type=float, help="Delay between pings in seconds.") + parser.add_argument("-l", "--log", action="store_true", help="Log the output to a file.") + parser.add_argument("-s", "--size", default="1280x640", help="Size of the plot.") - # unpack and initialize data - opts, args = parser.parse_args(argv) + opts = parser.parse_args(argv) ping = np.array([]) loss = np.array([]) t = np.array([]) now = np.array([]) cnt = 0 - - # write log if specified if opts.log or opts.fsave: now_time = datetime.datetime.now() date_str = now_time.isoformat()[:-7][:10] @@ -137,73 +118,53 @@ def main(argv=None): if opts.log: log_file = open(log_name, 'w') log_file.write(f"PingPlot Version {__version__} - Log File\n\n\n") - - # start the main loop + signal.signal(signal.SIGTERM, lambda signum, frame: signal_handler( + signum, frame, log_file=log_file if opts.log else None, ping=ping, now=now, nans=np.isnan(ping), opts=opts)) print(f"PingPlot Version {__version__} -- by ccampo\n") print("{0:^23}\n=======================".format("Run Parameters")) print("{0:>17} {1}".format("Hostname:", opts.host)) print("{0:>17} {1}".format("Ping interval:", str(opts.dt) + " s")) - print("{0:>17} {1}".format("Packets per ping:", opts.n)) + print("{0:>17} {1}".format("Packets per ping:", opts.num)) print("\n\nPress CTRL+C to quit...\n") print("{0:^15} {1:^15} {2:^15} {3:^15} {4:^15}\n".format("AVG. PING", "PACKET LOSS", "NUM. PINGS", "NUM. TIMEOUTS", "TIME ELAPSED")) - while True: - # quit on ctrl+c try: - ping, loss, t, out = call_pinger(opts.host, opts.n, ping, loss, t) + ping, loss, t, out = call_pinger(opts.host, opts.num, ping, loss, t) now = np.append(now, datetime.datetime.now()) cnt += 1 - - # get ping data mean_loss = loss.mean() nans = np.isnan(ping) if len(ping[~nans]) > 0: mean_ping = ping[~np.isnan(ping)].mean() else: mean_ping = np.nan - if opts.log: write_log(log_file, out) - - # only ping after time dt time.sleep(opts.dt) - delta_t = datetime.timedelta(seconds=(round(time.time() - t[0], 0))) sys.stdout.write("\r{0:^15.8} {1:^15.10} {2:^15} {3:^15} {4:^15}".format(str(round(mean_ping, 2)) + " ms", str(round(mean_loss, 2)) + " %", - cnt * opts.n, len(ping[nans]), + cnt * opts.num, len(ping[nans]), str(delta_t))) sys.stdout.flush() except KeyboardInterrupt: break - print("\n") - - # close log file if opts.log: print(f"Saved log file {log_name}") log_file.close() - - # make plot to save if opts.fsave or opts.plot: - # check if any data was collected if len(ping[~nans]) == 0: print("Error: cannot generate plot; no data collected. Please check your connection.") return 2 plt = plot_gen(ping, now, nans, opts.host, opts.plot, opts.size) - - # save if applicable if opts.fsave: print(f"Saved plot {plot_name}") plt.savefig(plot_name) - - # show plot if specified if opts.plot: plt.show() - - return 2 # exit - + return 2 if __name__ == "__main__": - sys.exit(main()) + main()