Python implementation of multicast equivalent of trivial netcat functionality

multicat.py

#!/usr/bin/env python

# http://docs.python.org/lib/socket-example.html

from optparse import OptionParser
import socket
import struct
import sys


def listen(input_socket,
           output_file,
           multicast_address,
           destination_port,
           interface_address,
           receive_block_size):

    input_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

    if sys.platform == "win32":
        input_socket.bind((interface_address, destination_port))
    else:
        input_socket.bind((multicast_address, destination_port))

    ip_mreq = (socket.inet_aton(multicast_address) +
               socket.inet_aton(interface_address))
    input_socket.setsockopt(socket.IPPROTO_IP,
                            socket.IP_ADD_MEMBERSHIP,
                            ip_mreq)

    try:
        while True:
            data = input_socket.recv(receive_block_size)
            if not data:
                pass
            output_file.write(data)
            output_file.flush()
    except KeyboardInterrupt:
        pass
    finally:
        output_file.flush()
        input_socket.setsockopt(socket.IPPROTO_IP,
                                socket.IP_DROP_MEMBERSHIP,
                                ip_mreq)
        input_socket.close()


def transmit(input_file,
             output_socket,
             multicast_address,
             destination_port,
             interface_address,
             transmit_block_size,
             time_to_live):

    if sys.platform != "win32":
        output_socket.setsockopt(socket.IPPROTO_IP,
                                 socket.IP_MULTICAST_IF,
                                 socket.inet_aton(interface_address))
    # disable Path MTU Discovery (PMTUD) as it doesn't address multicast
    # and silently drops any IP packets > a remote network's MTU
    if sys.platform == "linux2":
        try:
            output_socket.setsockopt(socket.IPPROTO_IP,
                                     socket.IP_MTU_DISCOVER,
                                     socket.IP_PMTUDISC_DONT)
        except AttributeError as e:
            # use numeric values instead of symbolic names
            # see /usr/include/linux/in.h
            if 'IP_MTU_DISCOVER' in e.message or \
               'IP_PMTUDISC_DONT' in e.message:
                output_socket.setsockopt(socket.IPPROTO_IP, 10, 0)
            else:
                raise e
    # recommend setting TTL to 31, the maximum recommended for local network
    # multicasts
    output_socket.setsockopt(socket.IPPROTO_IP,
                             socket.IP_MULTICAST_TTL,
                             time_to_live)
    # sent multicast packets should be looped back to local sockets by default,
    # but set IP_MULTICAST_LOOP to true to insure such
    output_socket.setsockopt(socket.IPPROTO_IP,
                             socket.IP_MULTICAST_LOOP,
                             True)

    try:
        write_data = b''

        while True:
            read_data = input_file.read(transmit_block_size - len(write_data))
            write_data += read_data

            while len(write_data) < transmit_block_size and read_data:
                read_data = input_file.read(transmit_block_size -
                                            len(write_data))
                write_data += read_data

            if write_data:
                bytes_sent = output_socket.sendto(write_data,
                                                  (multicast_address,
                                                   destination_port))
                write_data = write_data[bytes_sent:]

            if not read_data:
                # write out remainder of buffer
                while write_data:
                    bytes_sent = output_socket.sendto(write_data,
                                                      (multicast_address,
                                                       destination_port))
                    write_data = write_data[bytes_sent:]
                break
    except KeyboardInterrupt:
        pass
    finally:
        input_file.close()
        output_socket.close()


def main():
    usage = ("%prog [-l] [-t TTL] [-b BLOCKSIZE] [-i INTERFACE]"
             " -a ADDRESS -p PORT")

    parser = OptionParser(usage)

    parser.add_option("-a", "--address",
                      action="store",
                      type="string",
                      dest="address",
                      help="multicast IP address to send to or listen on")
    parser.add_option("-p", "--port",
                      action="store",
                      type="int",
                      dest="port",
                      help="UDP port to send to or listen on")
    parser.add_option("-l", "--listen",
                      action="store_true",
                      dest="listen",
                      default=False,
                      help="listen for packets (default is to send packets)")
    parser.add_option("-t", "--ttl",
                      action="store",
                      type="int",
                      dest="ttl",
                      default=1,
                      help="time to live (TTL) of sent packets")
    parser.add_option("-i", "--interface",
                      action="store",
                      type="string",
                      dest="interface",
                      help="IP address of interface to send or receive on")
    # max udp packet (65535) - udp header (8) - ip header (20) = 65507
    parser.add_option("-b", "--blocksize",
                      action="store",
                      type="int",
                      dest="blocksize",
                      default=65507,
                      help=("size of read or write to IP stack"
                            " (default is 65507)"))

    (options, args) = parser.parse_args()

    if args:
        parser.error("arguments not allowed")
    elif not (options.address and options.port):
        parser.error("address and port required")

    if sys.platform == "win32":
        import os
        import msvcrt
        msvcrt.setmode(sys.stdin.fileno(), os.O_BINARY)
        msvcrt.setmode(sys.stdout.fileno(), os.O_BINARY)

    if options.interface:
        interface = options.interface
    else:
        interface = socket.inet_ntoa(struct.pack('i', socket.INADDR_ANY))

    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)

    # Use binary/bytes, not text/strings, so on Python 3 access underlying
    # stdin/stdout buffer as stdin/stdout are "text files", but buffers are
    # "binary files".  See https://docs.python.org/3/library/sys.html#sys.stdin
    if sys.version_info.major >= 3:
        stdout = sys.stdout.buffer
        stdin = sys.stdin.buffer
    else:
        stdout = sys.stdout
        stdin = sys.stdin

    if options.listen:
        listen(s,
               stdout,
               options.address,
               options.port,
               interface,
               options.blocksize)
    else:
        transmit(stdin,
                 s,
                 options.address,
                 options.port,
                 interface,
                 options.blocksize,
                 options.ttl)


if __name__ == "__main__":
    main()