From 096f2b99a4e602e2957df9ab4935e5766f2f7b0c Mon Sep 17 00:00:00 2001 From: Hiroki Ohtani Date: Thu, 19 Mar 2015 08:50:09 +0900 Subject: [PATCH] refactoring. --- websocket/_core.py | 232 ++---------------------------- websocket/_socket.py | 7 +- websocket/tests/test_websocket.py | 35 +++-- 3 files changed, 31 insertions(+), 243 deletions(-) diff --git a/websocket/_core.py b/websocket/_core.py index e7aa71e..c87e7b2 100644 --- a/websocket/_core.py +++ b/websocket/_core.py @@ -50,8 +50,6 @@ else: import os import errno import struct -import uuid -import hashlib import threading # websocket modules @@ -61,6 +59,8 @@ from ._socket import * from ._utils import * from ._url import * from ._logging import * +from ._http import * +from ._handshake import * """ websocket python client. @@ -71,10 +71,6 @@ Please see http://tools.ietf.org/html/rfc6455 for protocol. """ -# websocket supported version. -VERSION = 13 - - def create_connection(url, timeout=None, **options): """ connect to url and return websocket object. @@ -125,17 +121,6 @@ def create_connection(url, timeout=None, **options): return websock -def _create_sec_websocket_key(): - uid = uuid.uuid4() - return base64encode(uid.bytes).decode('utf-8').strip() - - -_HEADERS_TO_CHECK = { - "upgrade": "websocket", - "connection": "upgrade", - } - - class WebSocket(object): """ Low level WebSocket interface. @@ -257,200 +242,17 @@ class WebSocket(object): default is None. """ - - hostname, port, resource, is_secure = parse_url(url) - proxy_host, proxy_port, proxy_auth = get_proxy_info(hostname, is_secure, **options) - if not proxy_host: - addrinfo_list = socket.getaddrinfo(hostname, port, 0, 0, socket.SOL_TCP) - else: - proxy_port = proxy_port and proxy_port or 80 - addrinfo_list = socket.getaddrinfo(proxy_host, proxy_port, 0, 0, socket.SOL_TCP) - - if not addrinfo_list: - raise WebSocketException("Host not found.: " + hostname + ":" + str(port)) - - err = None - for addrinfo in addrinfo_list: - family = addrinfo[0] - self.sock = socket.socket(family) - self.sock.settimeout(self.timeout) - for opts in DEFAULT_SOCKET_OPTION: - self.sock.setsockopt(*opts) - for opts in self.sockopt: - self.sock.setsockopt(*opts) - - address = addrinfo[4] - try: - self.sock.connect(address) - except socket.error as error: - error.remote_ip = str(address[0]) - if error.errno in (errno.ECONNREFUSED, ): - err = error - continue - else: - raise - else: - break - else: - raise err - - if proxy_host: - self._tunnel(hostname, port, proxy_auth) - - if is_secure: - if HAVE_SSL: - sslopt = dict(cert_reqs=ssl.CERT_REQUIRED) - certPath = os.path.join( - os.path.dirname(__file__), "cacert.pem") - if os.path.isfile(certPath): - sslopt['ca_certs'] = certPath - sslopt.update(self.sslopt) - check_hostname = sslopt.pop('check_hostname', True) - self.sock = ssl.wrap_socket(self.sock, **sslopt) - if (sslopt["cert_reqs"] != ssl.CERT_NONE - and check_hostname): - match_hostname(self.sock.getpeercert(), hostname) - else: - raise WebSocketException("SSL not available.") - - self._handshake(hostname, port, resource, **options) - - def _tunnel(self, host, port, auth): - debug("Connecting proxy...") - connect_header = "CONNECT %s:%d HTTP/1.0\r\n" % (host, port) - # TODO: support digest auth. - if auth and auth[0]: - auth_str = auth[0] - if auth[1]: - auth_str += ":" + auth[1] - encoded_str = base64encode(auth_str.encode()).strip().decode() - connect_header += "Proxy-Authorization: Basic %s\r\n" % encoded_str - connect_header += "\r\n" - dump("request header", connect_header) - - self._send(connect_header) + if "sockopt" in options: + del options["sockopt"] + self.sock, addrs = connect(url, self.sockopt, self.sslopt, self.timeout, **options) try: - status, resp_headers = self._read_headers() - except Exepiton as e: - raise WebSocketProxyException(str(e)) - - if status != 200: - raise WebSocketProxyException("failed CONNECT via proxy status: " + str(status)) - - def _get_resp_headers(self, success_status=101): - status, resp_headers = self._read_headers() - if status != success_status: - self.close() - raise WebSocketException("Handshake status %d" % status) - return resp_headers - - def _get_handshake_headers(self, resource, host, port, options): - headers = [] - headers.append("GET %s HTTP/1.1" % resource) - headers.append("Upgrade: websocket") - headers.append("Connection: Upgrade") - if port == 80: - hostport = host - else: - hostport = "%s:%d" % (host, port) - headers.append("Host: %s" % hostport) - - if "origin" in options: - headers.append("Origin: %s" % options["origin"]) - else: - headers.append("Origin: http://%s" % hostport) - - key = _create_sec_websocket_key() - headers.append("Sec-WebSocket-Key: %s" % key) - headers.append("Sec-WebSocket-Version: %s" % VERSION) - - subprotocols = options.get("subprotocols") - if subprotocols: - headers.append("Sec-WebSocket-Protocol: %s" % ",".join(subprotocols)) - - if "header" in options: - headers.extend(options["header"]) - - cookie = options.get("cookie", None) - - if cookie: - headers.append("Cookie: %s" % cookie) - - headers.append("") - headers.append("") - - return headers, key - - def _handshake(self, host, port, resource, **options): - headers, key = self._get_handshake_headers(resource, host, port, options) - - header_str = "\r\n".join(headers) - self._send(header_str) - dump("request header", header_str) - - resp_headers = self._get_resp_headers() - success = self._validate_header(resp_headers, key, options.get("subprotocols")) - if not success: - self.close() - raise WebSocketException("Invalid WebSocket Header") - - self.connected = True - - def _validate_header(self, headers, key, subprotocols): - for k, v in _HEADERS_TO_CHECK.items(): - r = headers.get(k, None) - if not r: - return False - r = r.lower() - if v != r: - return False - - if subprotocols: - subproto = headers.get("sec-websocket-protocol", None) - if not subproto or subproto not in subprotocols: - error("Invalid subprotocol: " + str(subprotocols)) - return False - self.subprotocol = subproto - - result = headers.get("sec-websocket-accept", None) - if not result: - return False - result = result.lower() - - if isinstance(result, six.text_type): - result = result.encode('utf-8') - - value = (key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11").encode('utf-8') - hashed = base64encode(hashlib.sha1(value).digest()).strip().lower() - return hashed == result - - def _read_headers(self): - status = None - headers = {} - trace("--- response header ---") - - while True: - line = self._recv_line() - line = line.decode('utf-8').strip() - if not line: - break - trace(line) - - if not status: - status_info = line.split(" ", 2) - status = int(status_info[1]) - else: - kv = line.split(":", 1) - if len(kv) == 2: - key, value = kv - headers[key.lower()] = value.strip().lower() - else: - raise WebSocketException("Invalid header") - - trace("-----------------------") - - return status, headers + self.subprotocol = handshake(self.sock, *addrs, **options) + self.connected = True + except: + self.sock.close() + self.sock = None + raise def send(self, payload, opcode=ABNF.OPCODE_TEXT): """ @@ -720,15 +522,3 @@ class WebSocket(object): else: self._recv_buffer = [unified[bufsize:]] return unified[:bufsize] - - def _recv_line(self): - try: - return recv_line(self.sock) - except WebSocketConnectionClosedException: - if self.sock: - self.sock.close() - self.sock = None - self.connected = False - raise - except: - raise diff --git a/websocket/_socket.py b/websocket/_socket.py index 1ecf4c6..b835130 100644 --- a/websocket/_socket.py +++ b/websocket/_socket.py @@ -73,7 +73,7 @@ def recv(sock, bufsize): if not bytes: raise WebSocketConnectionClosedException("Connection is already closed.") - + return bytes def recv_line(sock): @@ -89,10 +89,10 @@ def recv_line(sock): def send(sock, data): if isinstance(data, six.text_type): data = data.encode('utf-8') - + if not sock: raise WebSocketConnectionClosedException("socket is already closed.") - + try: return sock.send(data) except socket.timeout as e: @@ -104,4 +104,3 @@ def send(sock, data): raise WebSocketTimeoutException(message) else: raise - diff --git a/websocket/tests/test_websocket.py b/websocket/tests/test_websocket.py index 3d5297d..ab2a881 100644 --- a/websocket/tests/test_websocket.py +++ b/websocket/tests/test_websocket.py @@ -25,13 +25,16 @@ import uuid # websocket-client import websocket as ws -from websocket._core import _create_sec_websocket_key +from websocket._handshake import _create_sec_websocket_key from websocket._url import parse_url, get_proxy_info from websocket._utils import validate_utf8 +from websocket._handshake import _validate as _validate_header +from websocket._http import read_headers # Skip test to access the internet. TEST_WITH_INTERNET = False +# TEST_WITH_INTERNET = True # Skip Secure WebSocket test. TEST_SECURE_WS = False @@ -75,6 +78,7 @@ class HeaderSockMock(SockMock): with open(path, "rb") as f: self.add_packet(f.read()) + class WebSocketTest(unittest.TestCase): def setUp(self): ws.enableTrace(TRACABLE) @@ -178,49 +182,44 @@ class WebSocketTest(unittest.TestCase): self.assertTrue(six.u("¥n") not in key) def testWsUtils(self): - sock = ws.WebSocket() - key = "c6b8hTg4EeGb2gQMztV1/g==" required_header = { "upgrade": "websocket", "connection": "upgrade", "sec-websocket-accept": "Kxep+hNu9n51529fGidYu7a3wO0=", } - self.assertEqual(sock._validate_header(required_header, key, None), True) + self.assertEqual(_validate_header(required_header, key, None), (True, None)) header = required_header.copy() header["upgrade"] = "http" - self.assertEqual(sock._validate_header(header, key, None), False) + self.assertEqual(_validate_header(header, key, None), (False, None)) del header["upgrade"] - self.assertEqual(sock._validate_header(header, key, None), False) + self.assertEqual(_validate_header(header, key, None), (False, None)) header = required_header.copy() header["connection"] = "something" - self.assertEqual(sock._validate_header(header, key, None), False) + self.assertEqual(_validate_header(header, key, None), (False, None)) del header["connection"] - self.assertEqual(sock._validate_header(header, key, None), False) + self.assertEqual(_validate_header(header, key, None), (False, None)) header = required_header.copy() header["sec-websocket-accept"] = "something" - self.assertEqual(sock._validate_header(header, key, None), False) + self.assertEqual(_validate_header(header, key, None), (False, None)) del header["sec-websocket-accept"] - self.assertEqual(sock._validate_header(header, key, None), False) - + self.assertEqual(_validate_header(header, key, None), (False, None)) header = required_header.copy() header["sec-websocket-protocol"] = "sub1" - self.assertEqual(sock._validate_header(header, key, ["sub1", "sub2"]), True) - self.assertEqual(sock._validate_header(header, key, ["sub2", "sub3"]), False) + self.assertEqual(_validate_header(header, key, ["sub1", "sub2"]), (True, "sub1")) + self.assertEqual(_validate_header(header, key, ["sub2", "sub3"]), (False, None)) def testReadHeader(self): - sock = ws.WebSocket() - sock.sock = HeaderSockMock("data/header01.txt") - status, header = sock._read_headers() + status, header = read_headers(HeaderSockMock("data/header01.txt")) self.assertEqual(status, 101) self.assertEqual(header["connection"], "upgrade") - sock.sock = HeaderSockMock("data/header02.txt") - self.assertRaises(ws.WebSocketException, sock._read_headers) + HeaderSockMock("data/header02.txt") + self.assertRaises(ws.WebSocketException, read_headers, HeaderSockMock("data/header02.txt")) def testSend(self): # TODO: add longer frame data