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