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