Files
dscp-toy/sender.py
2025-01-29 12:13:15 +01:00

69 lines
2.1 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", type=str, 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}")
sock = socket(AF_INET, SOCK_DGRAM)
dscp = int(args.dscp) if args.dscp else 0
if dscp:
sock.setsockopt(IPPROTO_IP, IP_TOS, dscp << 2)
CODE_WORD = "CAFEABBA"
payload = (CODE_WORD + "A" * (PACKET_SIZE - len(CODE_WORD))).encode()
t0 = time.time()
while True:
logger.debug("Sending message:")
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()