2009-11-25 Yuzo Fujishima <yuzo@google.com>
[WebKit-https.git] / WebKitTools / pywebsocket / example / echo_client.py
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
35 This is an example Web Socket client that talks with echo_wsh.py.
36 This may be useful for checking mod_pywebsocket installation.
37
38 Note:
39 This code is far from robust, e.g., we cut corners in handshake.
40 """
41
42
43 import codecs
44 from optparse import OptionParser
45 import socket
46 import 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
65 def _method_line(resource):
66     return 'GET %s HTTP/1.1\r\n' % resource
67
68
69 def _origin_header(origin):
70     return 'Origin: %s\r\n' % origin
71
72
73 class _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
90 class 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
162 def 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
203 if __name__ == '__main__':
204     main()
205
206
207 # vi:sts=4 sw=4 et