examples/nc.py
author Tero Marttila <terom@fixme.fi>
Sat, 26 Sep 2009 21:52:06 +0300
changeset 54 e7c136812b0d
parent 48 ee7ade660c0b
child 55 99c4344a35ce
permissions -rw-r--r--
flesh out nc.py on_connect
#!/usr/bin/env python
"""
    Simple (limited) netcat-like example, implemented using the socket/lib.event2 layer.
"""

from qmsk.net.socket import socket, address
from qmsk.net.socket.constants import *

from qmsk.net.lib.event2 import base, event
from qmsk.net.lib.event2.constants import *
from qmsk.net.lib.event2.event import CallbackEvent as cb_event

import sys, os, fcntl, errno
import optparse

# global options
options = None

# global event_base
ev_base = base.event_base()


def parse_argv (argv) :
    global options

    prog = argv.pop(0)

    parser = optparse.OptionParser(prog=prog)

    parser.add_option('-4', "--ipv4", help="Force AF_INET", action='store_true')
    parser.add_option('-6', "--ipv6", help="Force AF_INET6", action='store_true')
    parser.add_option('-v', "--verbose", help="Display status output", action='store_true')
    parser.add_option('-d', "--debug", help="Display extra output", action='store_true')
    parser.add_option('-w', "--timeout", help="Timeout for connect()", type='float')

    options, args = parser.parse_args(argv)

    if options.ipv4 and options.ipv6 :
        raise Exception("-4 and -6 are mutually exclusive!")

    if options.debug :
        # enable both
        options.verbose = True

    return args

def log_msg (prefix, msg, *args) :
    if args :
        msg = msg % args
    
    sys.stderr.write("%s %s\n" % (prefix, msg))

def log_err (msg, *args) :
    log_msg('!!!', msg, *args)

def log_warn (msg, *args) :
    log_msg('+++', msg, *args)

def log_info (msg, *args) :
    if options.verbose :
        log_msg('***', msg, *args)

def log_debug (msg, *args) :
    if options.debug :
        log_msg('---', msg, *args)


def setnonblocking (file) :
    """
        Set the non-blocking IO flag on the given file object
    """

    fcntl.fcntl(file.fileno(), fcntl.F_SETFL, os.O_NONBLOCK)

def sock_connect (host, port, family, socktype) :
    """
        Try and perform a series of non-blocking connects to the given host:port using the given family and socktype,
        yielding a series of socket objects.
    """

    for ai in address.getaddrinfo(host, port, family, socktype) :
        log_info("sock_connect: ai: %s", ai)

        # build socket
        try :
            # construct
            log_debug("sock_connect: socket(%d, %d, %d)", ai.family, ai.socktype, ai.protocol)
            sock = socket.socket(ai.family, ai.socktype, ai.protocol)
            log_debug("sock_connect: socket=%s", sock)
            
            # set nonblock mode
            log_debug("sock_connect: setnonblocking(%s)", sock)
            setnonblocking(sock)
            
            # start connect
            try :
                log_debug("sock_connect: connect(%s)", ai.addr)
                sock.connect(ai.addr)

            except OSError, e :
                if e.errno != errno.EINPROGRESS :
                    raise
            
            else :
                # XXX: wut???
                log_warn("sock_connect: connect: didn't return EINPROGRESS")

        except OSError, e :
            # fsck
            log_warn("sock_connect: %s: %s", ai.addr, e)

        else :
            # yay
            yield sock

def on_connect (ev, events, sock, sock_iter) :
    """
        Outbound connect EV_WRITE callback, i.e. connection failed or was established
    """

    log_debug("on_connect: ev=%r, events=%#x", ev, events)

    # test for timeout
    if events & EV_TIMEOUT :
        log_warn("on_connect: connect failed, timeout")

        # keep trying
        return client_connect_next(sock_iter)


    # test for errno
    err = sock.getsockopt_int(SOL_SOCKET, SO_ERROR)

    if err :
        # fail
        log_warn("on_connect: connect failed, errno=%d", err)
        
        # keep trying
        return client_connect_next(sock_iter)


    # ok, connected
    log_info("on_connect: connected")


def client_connect_next (sock_iter) :
    """
        Attempt to run the given connect operation, on the given iterable of sockets.
    """

    for sock in sock_iter :
        # pend for writing
        log_debug("client_connect_next: cb_event(%d, EV_WRITE, on_connect, %r)", sock.fd, sock)
        ev = cb_event(ev_base, sock.fd, EV_WRITE, on_connect, sock, sock_iter)

        # wait specified timeout
        log_debug("client_connect_next: %r: add(%s)", ev, options.timeout)
        ev.add(options.timeout)

        # ok
        break

    else :
        # fail, ran out of addresses to try
        log_err("client_connect_next: ran out of addresses to try")

def run_client (host, port) :
    """
        Execute in client-mode
    """

    # figure out AF to use
    if options.ipv4 :
        family = AF_INET

    elif options.ipv6 :
        family = AF_INET6

    else :
        family = AF_UNSPEC

    # fixed socktype
    socktype = SOCK_STREAM

    # look up the address and start a non-blocking connect
    sock_iter = sock_connect(host, port, family, socktype)
    
    # start waiting
    client_connect_next(sock_iter)

def main (argv) :
    # parse args
    args = parse_argv(argv)
    
    # XXX: support listen mode
    host, port = args

    run_client(host, port)
        
    # run mainloop
    log_debug("main: entering event loop")

    if ev_base.loop() :
        log_debug("main: event loop done")

    else :
        log_err("main: event loop was idle!")

if __name__ == '__main__' :
    main(sys.argv)