• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/usr/bin/env python
2#
3# Copyright 2009, Google Inc.
4# All rights reserved.
5#
6# Redistribution and use in source and binary forms, with or without
7# modification, are permitted provided that the following conditions are
8# met:
9#
10#     * Redistributions of source code must retain the above copyright
11# notice, this list of conditions and the following disclaimer.
12#     * Redistributions in binary form must reproduce the above
13# copyright notice, this list of conditions and the following disclaimer
14# in the documentation and/or other materials provided with the
15# distribution.
16#     * Neither the name of Google Inc. nor the names of its
17# contributors may be used to endorse or promote products derived from
18# this software without specific prior written permission.
19#
20# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
21# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
22# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
23# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
24# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
25# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
26# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
27# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
28# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
29# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
30# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
31
32
33"""Web Socket Echo client.
34
35This is an example Web Socket client that talks with echo_wsh.py.
36This may be useful for checking mod_pywebsocket installation.
37
38Note:
39This code is far from robust, e.g., we cut corners in handshake.
40"""
41
42
43import codecs
44from optparse import OptionParser
45import socket
46import sys
47
48
49_TIMEOUT_SEC = 10
50
51_DEFAULT_PORT = 80
52_DEFAULT_SECURE_PORT = 443
53_UNDEFINED_PORT = -1
54
55_UPGRADE_HEADER = 'Upgrade: WebSocket\r\n'
56_CONNECTION_HEADER = 'Connection: Upgrade\r\n'
57_EXPECTED_RESPONSE = (
58        'HTTP/1.1 101 Web Socket Protocol Handshake\r\n' +
59        _UPGRADE_HEADER +
60        _CONNECTION_HEADER)
61
62_GOODBYE_MESSAGE = 'Goodbye'
63
64
65def _method_line(resource):
66    return 'GET %s HTTP/1.1\r\n' % resource
67
68
69def _origin_header(origin):
70    return 'Origin: %s\r\n' % origin
71
72
73class _TLSSocket(object):
74    """Wrapper for a TLS connection."""
75
76    def __init__(self, raw_socket):
77        self._ssl = socket.ssl(raw_socket)
78
79    def send(self, bytes):
80        return self._ssl.write(bytes)
81
82    def recv(self, size=-1):
83        return self._ssl.read(size)
84
85    def close(self):
86        # Nothing to do.
87        pass
88
89
90class EchoClient(object):
91    """Web Socket echo client."""
92
93    def __init__(self, options):
94        self._options = options
95        self._socket = None
96
97    def run(self):
98        """Run the client.
99
100        Shake hands and then repeat sending message and receiving its echo.
101        """
102        self._socket = socket.socket()
103        self._socket.settimeout(self._options.socket_timeout)
104        try:
105            self._socket.connect((self._options.server_host,
106                                  self._options.server_port))
107            if self._options.use_tls:
108                self._socket = _TLSSocket(self._socket)
109            self._handshake()
110            for line in self._options.message.split(',') + [_GOODBYE_MESSAGE]:
111                frame = '\x00' + line.encode('utf-8') + '\xff'
112                self._socket.send(frame)
113                if self._options.verbose:
114                    print 'Send: %s' % line
115                received = self._socket.recv(len(frame))
116                if received != frame:
117                    raise Exception('Incorrect echo: %r' % received)
118                if self._options.verbose:
119                    print 'Recv: %s' % received[1:-1].decode('utf-8',
120                                                             'replace')
121        finally:
122            self._socket.close()
123
124    def _handshake(self):
125        self._socket.send(_method_line(self._options.resource))
126        self._socket.send(_UPGRADE_HEADER)
127        self._socket.send(_CONNECTION_HEADER)
128        self._socket.send(self._format_host_header())
129        self._socket.send(_origin_header(self._options.origin))
130        self._socket.send('\r\n')
131
132        for expected_char in _EXPECTED_RESPONSE:
133            received = self._socket.recv(1)[0]
134            if expected_char != received:
135                raise Exception('Handshake failure')
136        # We cut corners and skip other headers.
137        self._skip_headers()
138
139    def _skip_headers(self):
140        terminator = '\r\n\r\n'
141        pos = 0
142        while pos < len(terminator):
143            received = self._socket.recv(1)[0]
144            if received == terminator[pos]:
145                pos += 1
146            elif received == terminator[0]:
147                pos = 1
148            else:
149                pos = 0
150
151    def _format_host_header(self):
152        host = 'Host: ' + self._options.server_host
153        if ((not self._options.use_tls and
154             self._options.server_port != _DEFAULT_PORT) or
155            (self._options.use_tls and
156             self._options.server_port != _DEFAULT_SECURE_PORT)):
157            host += ':' + str(self._options.server_port)
158        host += '\r\n'
159        return host
160
161
162def main():
163    sys.stdout = codecs.getwriter('utf-8')(sys.stdout)
164
165    parser = OptionParser()
166    parser.add_option('-s', '--server_host', dest='server_host', type='string',
167                      default='localhost', help='server host')
168    parser.add_option('-p', '--server_port', dest='server_port', type='int',
169                      default=_UNDEFINED_PORT, help='server port')
170    parser.add_option('-o', '--origin', dest='origin', type='string',
171                      default='http://localhost/', help='origin')
172    parser.add_option('-r', '--resource', dest='resource', type='string',
173                      default='/echo', help='resource path')
174    parser.add_option('-m', '--message', dest='message', type='string',
175                      help=('comma-separated messages to send excluding "%s" '
176                            'that is always sent at the end' %
177                            _GOODBYE_MESSAGE))
178    parser.add_option('-q', '--quiet', dest='verbose', action='store_false',
179                      default=True, help='suppress messages')
180    parser.add_option('-t', '--tls', dest='use_tls', action='store_true',
181                      default=False, help='use TLS (wss://)')
182    parser.add_option('-k', '--socket_timeout', dest='socket_timeout',
183                      type='int', default=_TIMEOUT_SEC,
184                      help='Timeout(sec) for sockets')
185
186    (options, unused_args) = parser.parse_args()
187
188    # Default port number depends on whether TLS is used.
189    if options.server_port == _UNDEFINED_PORT:
190        if options.use_tls:
191            options.server_port = _DEFAULT_SECURE_PORT
192        else:
193            options.server_port = _DEFAULT_PORT
194
195    # optparse doesn't seem to handle non-ascii default values.
196    # Set default message here.
197    if not options.message:
198        options.message = u'Hello,\u65e5\u672c'   # "Japan" in Japanese
199
200    EchoClient(options).run()
201
202
203if __name__ == '__main__':
204    main()
205
206
207# vi:sts=4 sw=4 et
208