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