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 # $FreeBSD$
  29 #
  30
  31
  32 from functools import partial
  33 import socket
  34 import select
  35 import argparse
  36 import time
  37
  38
  39 def parse_args():
  40     parser = argparse.ArgumentParser(description='divert socket tester')
  41     parser.add_argument('--sip', type=str, default='', help='IP to listen on')
  42     parser.add_argument('--family', type=str, help='inet/inet6')
  43     parser.add_argument('--ports', type=str, help='packet ports 1,2,3')
  44     parser.add_argument('--match_str', type=str, help='match string to use')
  45     parser.add_argument('--count', type=int, default=1,
  46                         help='Number of messages to receive')
  47     parser.add_argument('--test_name', type=str, required=True,
  48                         help='test name to run')
  49     return parser.parse_args()
  50
  51
  52 def test_listen_tcp(args):
  53     if args.family == 'inet6':
  54         fam = socket.AF_INET6
  55     else:
  56         fam = socket.AF_INET
  57     sockets = []
  58     ports = [int(port) for port in args.ports.split(',')]
  59     for port in ports:
  60         s = socket.socket(fam, socket.SOCK_STREAM)
  61         s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
  62         s.setblocking(0)
  63         s.bind((args.sip, port))
  64         print('binding on {}:{}'.format(args.sip, port))
  65         s.listen(5)
  66         sockets.append(s)
  67     inputs = sockets
  68     count = 0
  69     while count < args.count:
  70         readable, writable, exceptional = select.select(inputs, [], inputs)
  71         for s in readable:
  72             (c, address) = s.accept()
  73             print('C: {}'.format(address))
  74             data = c.recv(9000)
  75             if args.match_str and args.match_str.encode('utf-8') != data:
  76                 raise Exception('Expected "{}" but got "{}"'.format(
  77                     args.match_str, data.decode('utf-8')))
  78             count += 1
  79             c.close()
  80
  81
  82 def test_listen_udp(args):
  83     if args.family == 'inet6':
  84         fam = socket.AF_INET6
  85     else:
  86         fam = socket.AF_INET
  87     sockets = []
  88     ports = [int(port) for port in args.ports.split(',')]
  89     for port in ports:
  90         s = socket.socket(fam, socket.SOCK_DGRAM)
  91         s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
  92         s.setblocking(0)
  93         s.bind((args.sip, port))
  94         print('binding on {}:{}'.format(args.sip, port))
  95         sockets.append(s)
  96     inputs = sockets
  97     count = 0
  98     while count < args.count:
  99         readable, writable, exceptional = select.select(inputs, [], inputs)
 100         for s in readable:
 101             (data, address) = s.recvfrom(9000)
 102             print('C: {}'.format(address))
 103             if args.match_str and args.match_str.encode('utf-8') != data:
 104                 raise Exception('Expected "{}" but got "{}"'.format(
 105                     args.match_str, data.decode('utf-8')))
 106             count += 1
 107
 108
 109 def main():
 110     args = parse_args()
 111     test_ptr = globals()[args.test_name]
 112     test_ptr(args)
 113
 114
 115 if __name__ == '__main__':
 116     main()