Python UDP

udp_socket.py

import socket
import sys


def socket_client(host: str, port: int):
	with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as udp_client:
		dest = (host, port)
		print(f"Connected to: {host}:{port}")
		msg = input()
		while msg != '\x18':
			udp_client.sendto(str.encode(msg), dest)
			msg = input()


def socket_server(host: str, port: int):
	with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as udp_server:
		orig = (host, port)
		udp_server.bind(orig)
		print(f"Server started at: {host}:{port}")
		while True:
			msg, client = udp_server.recvfrom(1024)
			msg_arr = str(msg.decode()).split(',')
			if len(msg_arr) == 3:
				host_msg, port_msg, msg_msg = msg_arr
				print(f'Forwarding "{msg_msg}" to {host_msg}:{port_msg}')
				with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as udp_client:
					udp_client.sendto(str.encode(msg_msg), (host_msg, int(port_msg)))
			else:
				print(f"{client}:{msg}")


def index_of_arg(args: list, param) -> int:
	try:
		return args.index(param[0])
	except:
		pass
	try:
		return args.index(param[1])
	except:
		pass
	return -1


def get_params(args: list, param, count: int) -> list:
	pos = index_of_arg(args, param)
	if pos > 0:
		return args[pos + 1:pos + count + 1]
	return []


if __name__ == "__main__":
	client_params = get_params(sys.argv, ("-c", "--client"), 2)
	server_params = get_params(sys.argv, ("-s", "--server"), 2)
	if client_params:
		socket_client(client_params[0], int(client_params[1]))
	elif server_params:
		socket_server(server_params[0], int(server_params[1]))
	else:
		print("No option selected")