#!/usr/bin/env python
import argparse
import sys

from scapy.all import IP, UDP, Packet, XIntField, XByteField, bind_layers, Ether, sendp, sniff


class EasyTrace(Packet):
    PROTOCOL = 0xFD
    PORT = 0x6901
    name = "EasyTrace"
    fields_desc = [XIntField("num_valid", 0x0)]


class EasyPort(Packet):
    name = "EasyPort"
    fields_desc = [XByteField("port", 0x1)]


bind_layers(IP, UDP, proto=EasyTrace.PROTOCOL)
bind_layers(UDP, EasyTrace, dport=EasyTrace.PORT)


arp_table = {"10.0.0.1": "00:00:00:00:00:01", "10.0.0.2": "00:00:00:00:00:02", "10.0.0.3": "00:00:00:00:00:03"}


def client(args):
    p = (Ether(dst=arp_table[args.destination]) /
         IP(dst=args.destination, proto=17) /
         UDP(chksum=0, sport=12345, dport=args.port) /
         args.message)
    p.show()
    sendp(p, iface=args.interface)


def handle(pkt):
    pkt.show()
    if IP in pkt and pkt.proto == EasyTrace.PROTOCOL:
        num_valid = pkt[EasyTrace].num_valid
        print("EasyTrace of length {} (num_valid)".format(num_valid))
        if hasattr(pkt, 'load'):
            for i, port in enumerate(reversed(pkt.load[:num_valid])):
                print("{}. port = {}".format(i, str(ord(port))))
        elif num_valid:
            print("Malformed packet: num_valid={} but no ports present".format(num_valid))
        else:
            print "No ports (num_valid={})".format(num_valid)



def server(iface):
    sniff(iface=iface, prn=lambda pkt: handle(pkt))


def main():
    parser = argparse.ArgumentParser(description='receiver and sender to test P4 program')
    parser.add_argument("-s", "--server", help="run as server", action="store_true")
    parser.add_argument("-c", "--client", help="run as client", action="store_true")
    parser.add_argument("-i", "--interface", default='eth0', help="bind to specified interface")
    parser.add_argument("-p", "--port", type=int, default=EasyTrace.PORT, help="UDP destination Port")
    parser.add_argument("-d", "--destination", help="IP address of the destination")
    parser.add_argument("-m", "--message", default="", help="Packet payload")
    args = parser.parse_args()

    if args.server:
        server(args.interface)
    elif args.client:
        if args.destination is None:
            print "Missing destination IP address"
            parser.print_help()
            sys.exit(1)

        client(args)
    else:
        parser.print_help()

if __name__ == '__main__':
    main()