1

#!/usr/bin/env python

     2

"""

     3

    Simple (limited) netcat-like example, implemented using the socket/lib.event2 layer.

     4

"""

     5

     6

from qmsk.net.socket import socket, address

     7

from qmsk.net.socket.constants import *

     8

     9

from qmsk.net.lib.event2 import base, event

    10

from qmsk.net.lib.event2.constants import *

    11

from qmsk.net.lib.event2.event import CallbackEvent as cb_event

    12

    13

import sys, os, fcntl, errno

    14

import optparse

    15

    16

# global options

    17

options = None

    18

    19

# global event_base

    20

ev_base = base.event_base()

    21

    22

    23

def parse_argv (argv) :

    24

    global options

    25

    26

    prog = argv.pop(0)

    27

    28

    parser = optparse.OptionParser(prog=prog)

    29

    30

    parser.add_option('-4', "--ipv4", help="Force AF_INET", action='store_true')

    31

    parser.add_option('-6', "--ipv6", help="Force AF_INET6", action='store_true')

    32

    parser.add_option('-v', "--verbose", help="Display status output", action='store_true')

    33

    parser.add_option('-d', "--debug", help="Display extra output", action='store_true')

    34

    parser.add_option('-w', "--timeout", help="Timeout for connect()", type='float')

    35

    36

    options, args = parser.parse_args(argv)

    37

    38

    if options.ipv4 and options.ipv6 :

    39

        raise Exception("-4 and -6 are mutually exclusive!")

    40

    41

    if options.debug :

    42

        # enable both

    43

        options.verbose = True

    44

    45

    return args

    46

    47

def log_msg (prefix, msg, *args) :

    48

    if args :

    49

        msg = msg % args

    50

    51

    sys.stderr.write("%s %s\n" % (prefix, msg))

    52

    53

def log_err (msg, *args) :

    54

    log_msg('!!!', msg, *args)

    55

    56

def log_warn (msg, *args) :

    57

    log_msg('+++', msg, *args)

    58

    59

def log_info (msg, *args) :

    60

    if options.verbose :

    61

        log_msg('***', msg, *args)

    62

    63

def log_debug (msg, *args) :

    64

    if options.debug :

    65

        log_msg('---', msg, *args)

    66

    67

    68

def setnonblocking (file) :

    69

    """

    70

        Set the non-blocking IO flag on the given file object

    71

    """

    72

    73

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

    74

    75

def sock_connect (host, port, family, socktype) :

    76

    """

    77

        Try and perform a series of non-blocking connects to the given host:port using the given family and socktype,

    78

        yielding a series of socket objects.

    79

    """

    80

    81

    for ai in address.getaddrinfo(host, port, family, socktype) :

    83

    84

        # build socket

    85

        try :

    86

            # construct

    87

            log_debug("sock_connect: socket(%d, %d, %d)", ai.family, ai.socktype, ai.protocol)

    88

            sock = socket.socket(ai.family, ai.socktype, ai.protocol)

    89

            log_debug("sock_connect: socket=%s", sock)

    90

    91

            # set nonblock mode

    92

            log_debug("sock_connect: setnonblocking(%s)", sock)

    93

            setnonblocking(sock)

    94

    95

            # start connect

    96

            try :

    97

                log_debug("sock_connect: connect(%s)", ai.addr)

    98

                sock.connect(ai.addr)

    99

   100

            except OSError, e :

   101

                if e.errno != errno.EINPROGRESS :

   102

                    raise

   103

   104

            else :

   105

                # XXX: wut???

   106

                log_warn("sock_connect: connect: didn't return EINPROGRESS")

   107

   108

        except OSError, e :

   109

            # fsck

   110

            log_warn("sock_connect: %s: %s", ai.addr, e)

   111

   112

        else :

   113

            # yay

   114

            yield sock

   115

   117

    """

   118

        Outbound connect EV_WRITE callback, i.e. connection failed or was established

   119

    """

   120

   145

   146

def client_connect_next (sock_iter) :

   147

    """

   148

        Attempt to run the given connect operation, on the given iterable of sockets.

   149

    """

   150

   151

    for sock in sock_iter :

   152

        # pend for writing

   153

        log_debug("client_connect_next: cb_event(%d, EV_WRITE, on_connect, %r)", sock.fd, sock)

   155

   156

        # wait specified timeout

   157

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

   158

        ev.add(options.timeout)

   159

   160

        # ok

   161

        break

   162

   163

    else :

   164

        # fail, ran out of addresses to try

   165

        log_err("client_connect_next: ran out of addresses to try")

   166

   167

def run_client (host, port) :

   168

    """

   169

        Execute in client-mode

   170

    """

   171

   172

    # figure out AF to use

   173

    if options.ipv4 :

   174

        family = AF_INET

   175

   176

    elif options.ipv6 :

   177

        family = AF_INET6

   178

   179

    else :

   180

        family = AF_UNSPEC

   181

   182

    # fixed socktype

   183

    socktype = SOCK_STREAM

   184

   185

    # look up the address and start a non-blocking connect

   186

    sock_iter = sock_connect(host, port, family, socktype)

   187

   188

    # start waiting

   189

    client_connect_next(sock_iter)

   190

   191

def main (argv) :

   192

    # parse args

   193

    args = parse_argv(argv)

   194

   195

    # XXX: support listen mode

   196

    host, port = args

   197

   198

    run_client(host, port)

   199

   200

    # run mainloop

   201

    log_debug("main: entering event loop")

   202

   203

    if ev_base.loop() :

   204

        log_debug("main: event loop done")

   205

   206

    else :

   207

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

   208

   209

if __name__ == '__main__' :

   210

    main(sys.argv)

   211