#!/usr/bin/env python

import argparse
import os
import sys

from scapy.all import *
import matplotlib.pyplot as plt

def process_pcap(file_name):
    print('Opening {}...'.format(file_name))

    x_axis_time = [0]
    y_axis_com_data = [0]
    
    # Get times
    packets = rdpcap(file_name)
    for pkt in packets:
        x_axis_time.append(pkt.time)
    time_start = x_axis_time[1]
    time_end = x_axis_time[-1]
    # format time axis
    for i in range(1, len(x_axis_time)):
        x_axis_time[i] = x_axis_time[i] - time_start
    
    # get payloads
    pcap_file = RawPcapReader(file_name)
    for (pkt_data, pkt_metadata,) in pcap_file:
        y_axis_com_data.append(y_axis_com_data[-1] + pkt_metadata.wirelen)
    
    print('{} contains {} packets.\nStart time: {}\tEnd time: {}\t Transfered: {}kB\tDuration: {}'.format(file_name, len(packets), time_start, time_end, (y_axis_com_data[-1] / 8) / 1000, time_end - time_start))
    bitrate = (y_axis_com_data[-1] / (time_end - time_start)) / 1000000 # MBit/s
    print('Bitrate: ' + str(math.floor(bitrate * 100)/100) + 'MBit/s\n')

    return (x_axis_time, y_axis_com_data)    

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Plot two PCAP files')
    parser.add_argument('--client', metavar='<pcap file name>', help='pcap file to parse', required=True)
    parser.add_argument('--server', metavar='<pcap file name>', help='pcap file to parse', required=True)
    args = parser.parse_args()
    
    file_name1 = args.server
    file_name2 = args.client
    
    if not os.path.isfile(file_name1):
        print('"{}" does not exist'.format(file_name), file=sys.stderr)
        sys.exit(-1)
    if not os.path.isfile(file_name2):
        print('"{}" does not exist'.format(file_name), file=sys.stderr)
        sys.exit(-1)

    pcap1_x, pcap1_y = process_pcap(file_name1)
    pcap2_x, pcap2_y = process_pcap(file_name2)
    
    plt.plot(pcap1_x, pcap1_y)
    plt.plot(pcap2_x, pcap2_y, color='green')
    plt.xlabel('Time [ms]')
    plt.ylabel('Transfered Data [bit]')
    plt.legend(["Server (Laptop)", "Client (Pi)"])
    plt.show()
    sys.exit(0)