tcp_udp_payload_sender

#!/usr/bin/env python3

"""
This program connects to a ip and port via udp or tcp
and sends the datafiles to the destination.
"""

import socket
import sys
from pathlib import Path


class ArgumentError(Exception):
    pass


def usage():
    exe = Path(sys.executable).name
    prg = sys.argv[0]
    descr = __doc__
    cmd_usage = f'Usage: {exe} {prg} <ip> <port> <udp|tcp> <datafile> <...>'
    print(descr)
    print()
    print(cmd_usage)


def get_payloads(datafiles):
    return [Path(file).read_bytes() for file in datafiles]


def get_arguments():
    if len(sys.argv) < 5:
        usage()
        raise ArgumentError('Wrong number of arguments')
    ip, port, proto, *datafiles = sys.argv[1:]
    if proto == 'udp':
        sock_args = (socket.AF_INET, socket.SOCK_DGRAM)
    elif proto == 'tcp':
        sock_args = (socket.AF_INET, socket.SOCK_STREAM)
    else:
        usage()
        raise ArgumentError(f'Wrong protocol: {proto}')
    try:
        addr = (ip, int(port))
    except ValueError:
        usage()
        raise ArgumentError('Port must be a valid number')
    return addr, proto, sock_args, datafiles


def send_payloads_udp(connection, dest, payloads):
    for n, payload in enumerate(payloads, start=1):
        connection.sendto(payload, dest)
        print(f'[*] Payload {n} sent')


def send_payloads_tcp(connection, payloads):
    for n, payload in enumerate(payloads, start=1):
        connection.sendall(payload)
        print(f'[*] Payload {n} sent')


def main():
    try:
        addr, proto, sock_args, datafiles = get_arguments()
    except ArgumentError as e:
        print('Error:', e)
        return 1
    try:
        payloads = get_payloads(datafiles)
    except Exception as e:
        print('Error:', e)
        return 10
    with socket.socket(*sock_args) as sock:
        if proto == 'tcp':
            try:
                sock.connect(addr)
            except Exception as e:
                print('Error:', e)
                return 2
            print("[*] Connected")
            try:
                send_payloads_tcp(sock, payloads)
            except Exception as e:
                print('Error:', e)
                return 20
        elif proto == 'udp':
            try:
                send_payloads_udp(sock, addr, payloads)
            except Exception as e:
                print('Error:', e)
                return 10


if __name__ == '__main__':
    retval = main()
    sys.exit(retval)