Simple server to test sockets.

tsrv.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# author: kyle isom <[email protected]>
# license: ISC / public domain
"""
Simple server to test sockets.
"""

import argparse
import os
import random
import socket


DEFAULT_PORT = 4141                             # default port to listen on
DEFAULT_RECV = 4096                             # default buffer size
BANNER = None


def setup(port):
    """
    Set up socket and prepare for servery stuff.
    """
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.bind(('0.0.0.0', port))
    sock.listen(1)

    return sock


def client(conn, addr, echo_mode, recv_size):
    """
    Print client connections. If echo_mode is on, send back.
    """
    print 'client %s:%d ->' % addr, 

    data = ''
    if BANNER: 
        conn.send(BANNER)
    rcv = conn.recv(recv_size).strip()
    conn.settimeout(1)
    while rcv:
        data += rcv
        try:
            rcv = conn.recv(recv_size).strip()
        except socket.timeout:
            break

    print '%d bytes received.' % (len(data, ))

    if echo_mode:
        sent = len(data)
        while sent > 0:
            conn.send(data[:recv_size])
            data = data[recv_size:]
            sent -= recv_size
    conn.close()
    os._exit(os.EX_OK)


def server(echo_mode, port, recv_size):
    """Run the server. Forks for new connections."""
    sock = setup(port)

    while True:
        try:
            conn, addr = sock.accept()
            if 0 == os.fork():
                client(conn, addr, echo_mode, recv_size)
                os._exit(os.EX_OK)
        except KeyboardInterrupt:
            sock.shutdown(socket.SHUT_RDWR)
            exit(0)


def main():
    """
    Parse arguments and run the server.
    """
    global BANNER

    parser = argparse.ArgumentParser(description='test socket server')
    parser.add_argument('-b', '--banner', action='store',
                        help='send banner to clients on connect')
    parser.add_argument('-e', '--echo', action='store_true',
                        help='echo data sent back to sender')
    parser.add_argument('-p', '--port', help='port to listen on')
    parser.add_argument('-r', '--random', action='store_true',
                        help='randomise port (overrides -p)')
    parser.add_argument('-s', '--size', help='network buffer size')

    args = parser.parse_args()
    port = DEFAULT_PORT
    recv_size = DEFAULT_RECV

    if args.port:
        port = int(args.port)
    if args.random:
        port = random.randint(1024, 32768)
    if args.echo:
        echo = True
    if args.banner:
        BANNER = args.banner
    else:
        echo = False
    if args.size:
        recv_size = int(args.size)

    print 'listening on %d' % (port, )
    server(echo, port, recv_size)

if '__main__' == __name__:
    main()