78 lines
2.3 KiB
Python
78 lines
2.3 KiB
Python
from argparse import ArgumentParser
|
|
from socket import socket, AF_INET, SOCK_DGRAM, IPPROTO_IP, IP_TOS
|
|
from time import sleep
|
|
import re
|
|
import time
|
|
import sys
|
|
from loguru import logger
|
|
|
|
|
|
def parse_bandwidth(input_str):
|
|
pattern = r'^(\d+(?:\.\d+)?)\s*([kKmMgGtT])bps$'
|
|
match = re.match(pattern, input_str)
|
|
|
|
if match:
|
|
value = float(match.group(1))
|
|
unit = match.group(2).lower()
|
|
|
|
multipliers = {
|
|
'k': 1e3,
|
|
'm': 1e6,
|
|
'g': 1e9,
|
|
't': 1e12
|
|
}
|
|
|
|
multiplier = multipliers.get(unit, 1.0)
|
|
return value * multiplier
|
|
else:
|
|
raise ValueError(f"Invalid bandwidth: {input_str}")
|
|
|
|
|
|
def main():
|
|
parser = ArgumentParser()
|
|
parser.add_argument('--host', type=str, default='localhost')
|
|
parser.add_argument('--port', type=int, default=12345)
|
|
parser.add_argument("--log-level", type=str, default="INFO")
|
|
parser.add_argument("--bandwidth", type=str, default="1Mbps")
|
|
parser.add_argument("--duration", type=int, default=10)
|
|
parser.add_argument("--dscp", action='append', default=[])
|
|
args = parser.parse_args()
|
|
logger.remove()
|
|
logger.add(sys.stderr, level=args.log_level)
|
|
PACKET_SIZE = 1280
|
|
bandwidth = parse_bandwidth(args.bandwidth)
|
|
time_between_packets = PACKET_SIZE / bandwidth
|
|
packet_cnt = 0
|
|
logger.info(f"Connecting to {args.host}:{args.port}")
|
|
|
|
|
|
sockets = []
|
|
dscp = [int(x) for x in args.dscp]
|
|
if dscp:
|
|
for value in dscp:
|
|
sock = socket(AF_INET, SOCK_DGRAM)
|
|
sock.setsockopt(IPPROTO_IP, IP_TOS, value << 2)
|
|
sockets.append(sock)
|
|
else:
|
|
sock = socket(AF_INET, SOCK_DGRAM)
|
|
sockets.append(sock)
|
|
CODE_WORD = "CAFEABBA"
|
|
payload = (CODE_WORD + "A" * (PACKET_SIZE - len(CODE_WORD))).encode()
|
|
t0 = time.time()
|
|
while True:
|
|
logger.debug("Sending message:")
|
|
for sock in sockets:
|
|
sock.sendto(payload, (args.host, args.port))
|
|
packet_cnt += 1
|
|
t1 = time.time() - t0
|
|
if t1 > args.duration:
|
|
break
|
|
time_to_sleep = time_between_packets * packet_cnt - t1
|
|
if time_to_sleep > 0:
|
|
sleep(time_between_packets * packet_cnt - t1)
|
|
else:
|
|
logger.warning("We are too slow: {}s", -time_to_sleep)
|
|
|
|
if __name__ == '__main__':
|
|
main()
|