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) :

    82

        log_info("sock_connect: ai=%s", ai)

    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

   116

def on_connect (fd, events, sock) :

   117

    """

   118

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

   119

    """

   120

   121

    log_info("on_connect: %x", events)

   122

   123

def client_connect_next (sock_iter) :

   124

    """

   125

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

   126

    """

   127

   128

    for sock in sock_iter :

   129

        # pend for writing

   130

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

   131

        ev = cb_event(ev_base, sock.fd, EV_WRITE, on_connect, sock)

   132

   133

        # wait specified timeout

   134

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

   135

        ev.add(options.timeout)

   136

   137

        # ok

   138

        break

   139

   140

    else :

   141

        # fail, ran out of addresses to try

   142

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

   143

   144

def run_client (host, port) :

   145

    """

   146

        Execute in client-mode

   147

    """

   148

   149

    # figure out AF to use

   150

    if options.ipv4 :

   151

        family = AF_INET

   152

   153

    elif options.ipv6 :

   154

        family = AF_INET6

   155

   156

    else :

   157

        family = AF_UNSPEC

   158

   159

    # fixed socktype

   160

    socktype = SOCK_STREAM

   161

   162

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

   163

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

   164

   165

    # start waiting

   166

    client_connect_next(sock_iter)

   167

   168

def main (argv) :

   169

    # parse args

   170

    args = parse_argv(argv)

   171

   172

    # XXX: support listen mode

   173

    host, port = args

   174

   175

    run_client(host, port)

   176

   177

    # run mainloop

   178

    log_debug("main: entering event loop")

   179

   180

    if ev_base.loop() :

   181

        log_debug("main: event loop done")

   182

   183

    else :

   184

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

   185

   186

if __name__ == '__main__' :

   187

    main(sys.argv)

   188