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