tests/sys/common/net_receiver.py

   1 #!/usr/bin/env python
   2 # -
   3 # SPDX-License-Identifier: BSD-2-Clause
   4 #
   5 # Copyright (c) 2020 Alexander V. Chernikov
   6 #
   7 # Redistribution and use in source and binary forms, with or without
   8 # modification, are permitted provided that the following conditions
   9 # are met:
  10 # 1. Redistributions of source code must retain the above copyright
  11 #    notice, this list of conditions and the following disclaimer.
  12 # 2. Redistributions in binary form must reproduce the above copyright
  13 #    notice, this list of conditions and the following disclaimer in the
  14 #    documentation and/or other materials provided with the distribution.
  15 #
  16 # THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
  17 # ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
  18 # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
  19 # ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
  20 # FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
  21 # DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
  22 # OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
  23 # HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
  24 # LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
  25 # OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
  26 # SUCH DAMAGE.
  27 #
  28 #
  29
  30
  31 from functools import partial
  32 import socket
  33 import select
  34 import argparse
  35 import time
  36
  37
  38 def parse_args():
  39     parser = argparse.ArgumentParser(description='divert socket tester')
  40     parser.add_argument('--sip', type=str, default='', help='IP to listen on')
  41     parser.add_argument('--family', type=str, help='inet/inet6')
  42     parser.add_argument('--ports', type=str, help='packet ports 1,2,3')
  43     parser.add_argument('--match_str', type=str, help='match string to use')
  44     parser.add_argument('--count', type=int, default=1,
  45                         help='Number of messages to receive')
  46     parser.add_argument('--test_name', type=str, required=True,
  47                         help='test name to run')
  48     return parser.parse_args()
  49
  50
  51 def test_listen_tcp(args):
  52     if args.family == 'inet6':
  53         fam = socket.AF_INET6
  54     else:
  55         fam = socket.AF_INET
  56     sockets = []
  57     ports = [int(port) for port in args.ports.split(',')]
  58     for port in ports:
  59         s = socket.socket(fam, socket.SOCK_STREAM)
  60         s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
  61         s.setblocking(0)
  62         s.bind((args.sip, port))
  63         print('binding on {}:{}'.format(args.sip, port))
  64         s.listen(5)
  65         sockets.append(s)
  66     inputs = sockets
  67     count = 0
  68     while count < args.count:
  69         readable, writable, exceptional = select.select(inputs, [], inputs)
  70         for s in readable:
  71             (c, address) = s.accept()
  72             print('C: {}'.format(address))
  73             data = c.recv(9000)
  74             if args.match_str and args.match_str.encode('utf-8') != data:
  75                 raise Exception('Expected "{}" but got "{}"'.format(
  76                     args.match_str, data.decode('utf-8')))
  77             count += 1
  78             c.close()
  79
  80
  81 def test_listen_udp(args):
  82     if args.family == 'inet6':
  83         fam = socket.AF_INET6
  84     else:
  85         fam = socket.AF_INET
  86     sockets = []
  87     ports = [int(port) for port in args.ports.split(',')]
  88     for port in ports:
  89         s = socket.socket(fam, socket.SOCK_DGRAM)
  90         s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
  91         s.setblocking(0)
  92         s.bind((args.sip, port))
  93         print('binding on {}:{}'.format(args.sip, port))
  94         sockets.append(s)
  95     inputs = sockets
  96     count = 0
  97     while count < args.count:
  98         readable, writable, exceptional = select.select(inputs, [], inputs)
  99         for s in readable:
 100             (data, address) = s.recvfrom(9000)
 101             print('C: {}'.format(address))
 102             if args.match_str and args.match_str.encode('utf-8') != data:
 103                 raise Exception('Expected "{}" but got "{}"'.format(
 104                     args.match_str, data.decode('utf-8')))
 105             count += 1
 106
 107
 108 def main():
 109     args = parse_args()
 110     test_ptr = globals()[args.test_name]
 111     test_ptr(args)
 112
 113
 114 if __name__ == '__main__':
 115     main()