src/scripts/vnet/uri/dummy_app.py

   1 #!/usr/bin/env python3
   2
   3 import socket
   4 import sys
   5 import time
   6 import argparse
   7
   8 # action can be reflect or drop
   9 action = "drop"
  10 test = 0
  11
  12
  13 def test_data(data, n_rcvd):
  14     n_read = len(data)
  15     for i in range(n_read):
  16         expected = (n_rcvd + i) & 0xFF
  17         byte_got = ord(data[i])
  18         if byte_got != expected:
  19             print(
  20                 "Difference at byte {}. Expected {} got {}".format(
  21                     n_rcvd + i, expected, byte_got
  22                 )
  23             )
  24     return n_read
  25
  26
  27 def handle_connection(connection, client_address):
  28     print("Received connection from {}".format(repr(client_address)))
  29     n_rcvd = 0
  30     try:
  31         while True:
  32             data = connection.recv(4096)
  33             if not data:
  34                 break
  35             if test == 1:
  36                 n_rcvd += test_data(data, n_rcvd)
  37             if action != "drop":
  38                 connection.sendall(data)
  39     finally:
  40         connection.close()
  41
  42
  43 def run_tcp_server(ip, port):
  44     print("Starting TCP server {}:{}".format(repr(ip), repr(port)))
  45     sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  46     sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
  47     server_address = (ip, int(port))
  48     sock.bind(server_address)
  49     sock.listen(1)
  50     while True:
  51         connection, client_address = sock.accept()
  52         handle_connection(connection, client_address)
  53
  54
  55 def run_udp_server(ip, port):
  56     print("Starting UDP server {}:{}".format(repr(ip), repr(port)))
  57     sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
  58     sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
  59     server_address = (ip, int(port))
  60     sock.bind(server_address)
  61     while True:
  62         data, addr = sock.recvfrom(4096)
  63         if action != "drop":
  64             # snd_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
  65             sock.sendto(data, addr)
  66
  67
  68 def run_server(ip, port, proto):
  69     if proto == "tcp":
  70         run_tcp_server(ip, port)
  71     elif proto == "udp":
  72         run_udp_server(ip, port)
  73
  74
  75 def prepare_data(power):
  76     buf = []
  77     for i in range(0, pow(2, power)):
  78         buf.append(i & 0xFF)
  79     return bytearray(buf)
  80
  81
  82 def run_tcp_client(ip, port):
  83     print("Starting TCP client {}:{}".format(repr(ip), repr(port)))
  84     sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  85     server_address = (ip, int(port))
  86     sock.connect(server_address)
  87
  88     data = prepare_data(16)
  89     n_rcvd = 0
  90     n_sent = len(data)
  91     try:
  92         sock.sendall(data)
  93
  94         timeout = time.time() + 2
  95         while n_rcvd < n_sent and time.time() < timeout:
  96             tmp = sock.recv(1500)
  97             tmp = bytearray(tmp)
  98             n_read = len(tmp)
  99             for i in range(n_read):
 100                 if data[n_rcvd + i] != tmp[i]:
 101                     print(
 102                         "Difference at byte {}. Sent {} got {}".format(
 103                             n_rcvd + i, data[n_rcvd + i], tmp[i]
 104                         )
 105                     )
 106             n_rcvd += n_read
 107
 108         if n_rcvd < n_sent or n_rcvd > n_sent:
 109             print("Sent {} and got back {}".format(n_sent, n_rcvd))
 110         else:
 111             print("Got back what we've sent!!")
 112
 113     finally:
 114         sock.close()
 115
 116
 117 def run_udp_client(ip, port):
 118     print("Starting UDP client {}:{}".format(repr(ip), repr(port)))
 119     n_packets = 100
 120     sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
 121     server_address = (ip, int(port))
 122     data = prepare_data(10)
 123     try:
 124         for i in range(0, n_packets):
 125             sock.sendto(data, server_address)
 126     finally:
 127         sock.close()
 128
 129
 130 def run_client(ip, port, proto):
 131     if proto == "tcp":
 132         run_tcp_client(ip, port)
 133     elif proto == "udp":
 134         run_udp_client(ip, port)
 135
 136
 137 def run(mode, ip, port, proto):
 138     if mode == "server":
 139         run_server(ip, port, proto)
 140     elif mode == "client":
 141         run_client(ip, port, proto)
 142     else:
 143         raise Exception("Unknown mode. Only client and server supported")
 144
 145
 146 if __name__ == "__main__":
 147     parser = argparse.ArgumentParser()
 148     parser.add_argument("-m", action="store", dest="mode")
 149     parser.add_argument("-i", action="store", dest="ip")
 150     parser.add_argument("-p", action="store", dest="port")
 151     parser.add_argument("-proto", action="store", dest="proto")
 152     parser.add_argument("-a", action="store", dest="action")
 153     parser.add_argument("-t", action="store", dest="test")
 154     results = parser.parse_args()
 155     action = results.action
 156     test = results.test
 157     run(results.mode, results.ip, results.port, results.proto)
 158     # if (len(sys.argv)) < 4:
 159     #    raise Exception("Usage: ./dummy_app <mode> <ip> <port> [<action> <test>]")
 160     # if (len(sys.argv) == 6):
 161     #    action = sys.argv[4]
 162     #    test = int(sys.argv[5])
 163     # run (sys.argv[1], sys.argv[2], int(sys.argv[3]))