61b129c50d59b1c3e5ebb7e004a65fc0ce6ccd6d
[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 _DEFAULT_PORT = 80
50 _DEFAULT_SECURE_PORT = 443
51 _UNDEFINED_PORT = -1
52
53 _UPGRADE_HEADER = 'Upgrade: WebSocket\r\n'
54 _CONNECTION_HEADER = 'Connection: Upgrade\r\n'
55 _EXPECTED_RESPONSE = (
56         'HTTP/1.1 101 Web Socket Protocol Handshake\r\n' +
57         _UPGRADE_HEADER +
58         _CONNECTION_HEADER)
59
60
61 def _method_line(resource):
62     return 'GET %s HTTP/1.1\r\n' % resource
63
64
65 def _origin_header(origin):
66     return 'Origin: %s\r\n' % origin
67
68
69 class _TLSSocket(object):
70     """Wrapper for a TLS connection."""
71
72     def __init__(self, raw_socket):
73         self._ssl = socket.ssl(raw_socket)
74
75     def send(self, bytes):
76         return self._ssl.write(bytes)
77
78     def recv(self, size=-1):
79         return self._ssl.read(size)
80
81     def close(self):
82         # Nothing to do.
83         pass
84
85
86 class EchoClient(object):
87     """Web Socket echo client."""
88
89     def __init__(self, options):
90         self._options = options
91         self._socket = None
92
93     def run(self):
94         """Run the client.
95
96         Shake hands and then repeat sending message and receiving its echo.
97         """
98         self._socket = socket.socket()
99         try:
100             self._socket.connect((self._options.server_host,
101                                   self._options.server_port))
102             if self._options.use_tls:
103                 self._socket = _TLSSocket(self._socket)
104             self._handshake()
105             for line in self._options.message.split(','):
106                 frame = '\x00' + line.encode('utf-8') + '\xff'
107                 self._socket.send(frame)
108                 if self._options.verbose:
109                     print 'Send: %s' % line
110                 received = self._socket.recv(len(frame))
111                 if received != frame:
112                     raise Exception('Incorrect echo: %r' % received)
113                 if self._options.verbose:
114                     print 'Recv: %s' % received[1:-1].decode('utf-8')
115         finally:
116             self._socket.close()
117
118     def _handshake(self):
119         self._socket.send(_method_line(self._options.resource))
120         self._socket.send(_UPGRADE_HEADER)
121         self._socket.send(_CONNECTION_HEADER)
122         self._socket.send(self._format_host_header())
123         self._socket.send(_origin_header(self._options.origin))
124         self._socket.send('\r\n')
125
126         for expected_char in _EXPECTED_RESPONSE:
127             received = self._socket.recv(1)[0]
128             if expected_char != received:
129                 raise Exception('Handshake failure')
130         # We cut corners and skip other headers.
131         self._skip_headers()
132
133     def _skip_headers(self):
134         terminator = '\r\n\r\n'
135         pos = 0
136         while pos < len(terminator):
137             received = self._socket.recv(1)[0]
138             if received == terminator[pos]:
139                 pos += 1
140             elif received == terminator[0]:
141                 pos = 1
142             else:
143                 pos = 0
144
145     def _format_host_header(self):
146         host = 'Host: ' + self._options.server_host
147         if ((not self._options.use_tls and
148              self._options.server_port != _DEFAULT_PORT) or
149             (self._options.use_tls and
150              self._options.server_port != _DEFAULT_SECURE_PORT)):
151             host += ':' + str(self._options.server_port)
152         host += '\r\n'
153         return host
154
155
156 def main():
157     sys.stdout = codecs.getwriter('utf-8')(sys.stdout)
158
159     parser = OptionParser()
160     parser.add_option('-s', '--server_host', dest='server_host', type='string',
161                       default='localhost', help='server host')
162     parser.add_option('-p', '--server_port', dest='server_port', type='int',
163                       default=_UNDEFINED_PORT, help='server port')
164     parser.add_option('-o', '--origin', dest='origin', type='string',
165                       default='http://localhost/', help='origin')
166     parser.add_option('-r', '--resource', dest='resource', type='string',
167                       default='/echo', help='resource path')
168     parser.add_option('-m', '--message', dest='message', type='string',
169                       help='comma-separated messages to send')
170     parser.add_option('-q', '--quiet', dest='verbose', action='store_false',
171                       default=True, help='suppress messages')
172     parser.add_option('-t', '--tls', dest='use_tls', action='store_true',
173                       default=False, help='use TLS (wss://)')
174     (options, unused_args) = parser.parse_args()
175
176     # Default port number depends on whether TLS is used.
177     if options.server_port == _UNDEFINED_PORT:
178         if options.use_tls:
179             options.server_port = _DEFAULT_SECURE_PORT
180         else:
181             options.server_port = _DEFAULT_PORT
182
183     # optparse doesn't seem to handle non-ascii default values.
184     # Set default message here.
185     if not options.message:
186         options.message = u'Hello,\u65e5\u672c'   # "Japan" in Japanese
187
188     EchoClient(options).run()
189
190
191 if __name__ == '__main__':
192     main()
193
194
195 # vi:sts=4 sw=4 et