diff --git a/ChangeLog b/ChangeLog index 1aaf7ce..758ddb2 100644 --- a/ChangeLog +++ b/ChangeLog @@ -5,6 +5,7 @@ ChangeLog - remove unittest2 requirements for python 2.6 (#156) - fixed subprotocol case during header validation (#158) + - get response status and headers (#160) - refactoring. - 0.26.0 diff --git a/websocket/_core.py b/websocket/_core.py index c87e7b2..23d9520 100644 --- a/websocket/_core.py +++ b/websocket/_core.py @@ -155,15 +155,11 @@ class WebSocket(object): """ Initalize WebSocket object. """ - if sockopt is None: - sockopt = [] - if sslopt is None: - sslopt = {} - self.connected = False + self.sock_opt = sock_opt(sockopt, sslopt) + self.handshake_response = None self.sock = None - self._timeout = None - self.sockopt = sockopt - self.sslopt = sslopt + + self.connected = False self.get_mask_key = get_mask_key self.fire_cont_frame = fire_cont_frame self.skip_utf8_validation = skip_utf8_validation @@ -174,13 +170,12 @@ class WebSocket(object): self._frame_buffer = FrameBuffer() self._cont_data = None self._recving_frames = None + if enable_multithread: self.lock = threading.Lock() else: self.lock = NoLock() - self.subprotocol = None - def fileno(self): return self.sock.fileno() @@ -200,7 +195,7 @@ class WebSocket(object): """ Get the websocket timeout(second). """ - return self._timeout + return self.sock_opt.timeout def settimeout(self, timeout): """ @@ -208,12 +203,45 @@ class WebSocket(object): timeout: timeout time(second). """ - self._timeout = timeout + self.sock_opt.timeout = timeout if self.sock: self.sock.settimeout(timeout) timeout = property(gettimeout, settimeout) + def getsubprotocol(self): + """ + get subprotocol + """ + if self.handshake_response: + return self.handshake_response.subprotocol + else: + return None + + subprotocol = property(getsubprotocol) + + def getstatus(self): + """ + get handshake status + """ + if self.handshake_response: + return self.handshake_response.status + else: + return None + + status = property(getstatus) + + def getheaders(self): + """ + get handshake response header + """ + if self.handshake_response: + return self.handshake_response.headers + else: + return None + + headers = property(getheaders) + def connect(self, url, **options): """ Connect to url. url is websocket url scheme. @@ -232,6 +260,7 @@ class WebSocket(object): options: "header" -> custom http header list. "cookie" -> cookie value. + "origin" -> custom origin url. "http_proxy_host" - http proxy host name. "http_proxy_port" - http proxy port. If not set, set to 80. "http_no_proxy" - host names, which doesn't use proxy. @@ -242,12 +271,10 @@ class WebSocket(object): default is None. """ - if "sockopt" in options: - del options["sockopt"] - self.sock, addrs = connect(url, self.sockopt, self.sslopt, self.timeout, **options) + self.sock, addrs = connect(url, self.sock_opt, proxy_info(**options)) try: - self.subprotocol = handshake(self.sock, *addrs, **options) + self.handshake_response = handshake(self.sock, *addrs, **options) self.connected = True except: self.sock.close() diff --git a/websocket/_handshake.py b/websocket/_handshake.py index ddcd4d3..1fbef04 100644 --- a/websocket/_handshake.py +++ b/websocket/_handshake.py @@ -35,12 +35,18 @@ from ._socket import* from ._http import * from ._exceptions import * -__all__ = ["handshake"] +__all__ = ["handshake_response", "handshake"] # websocket supported version. VERSION = 13 +class handshake_response(object): + def __init__(self, status, headers, subprotocol): + self.status = status + self.headers = headers + self.subprotocol = subprotocol + def handshake(sock, host, port, resource, **options): headers, key = _get_handshake_headers(resource, host, port, options) @@ -48,12 +54,12 @@ def handshake(sock, host, port, resource, **options): send(sock, header_str) dump("request header", header_str) - resp = _get_resp_headers(sock) + status, resp = _get_resp_headers(sock) success, subproto = _validate(resp, key, options.get("subprotocols")) if not success: raise WebSocketException("Invalid WebSocket Header") - return subproto + return handshake_response(status, resp, subproto) def _get_handshake_headers(resource, host, port, options): @@ -98,7 +104,7 @@ def _get_resp_headers(sock, success_status=101): status, resp_headers = read_headers(sock) if status != success_status: raise WebSocketException("Handshake status %d" % status) - return resp_headers + return status, resp_headers _HEADERS_TO_CHECK = { "upgrade": "websocket", diff --git a/websocket/_http.py b/websocket/_http.py index 1cf5f1a..383d610 100644 --- a/websocket/_http.py +++ b/websocket/_http.py @@ -49,25 +49,36 @@ from ._url import * from ._socket import* from ._exceptions import * -__all__ = ["connect", "read_headers"] +__all__ = ["proxy_info", "connect", "read_headers"] +class proxy_info(object): + def __init__(self, **options): + self.host = options.get("http_proxy_host", None) + if self.host: + self.port = options.get("http_proxy_port", 0) + self.auth = options.get("http_proxy_auth", None) + self.no_proxy = options.get("http_no_proxy", None) + else: + self.port = 0 + self.auth = None + self.no_proxy = None -def connect(url, sockopt, sslopt, timeout, **options): +def connect(url, options, proxy): hostname, port, resource, is_secure = parse_url(url) - addrinfo_list, need_tunnel, auth = _get_addrinfo_list(hostname, port, is_secure, **options) + addrinfo_list, need_tunnel, auth = _get_addrinfo_list(hostname, port, is_secure, proxy) if not addrinfo_list: raise WebSocketException( "Host not found.: " + hostname + ":" + str(port)) sock = None try: - sock = _open_socket(addrinfo_list, sockopt, timeout) + sock = _open_socket(addrinfo_list, options.sockopt, options.timeout) if need_tunnel: sock = _tunnel(sock, hostname, port, auth) if is_secure: if HAVE_SSL: - sock = _ssl_socket(sock, sslopt) + sock = _ssl_socket(sock, options.sslopt) else: raise WebSocketException("SSL not available.") @@ -78,8 +89,9 @@ def connect(url, sockopt, sslopt, timeout, **options): raise -def _get_addrinfo_list(hostname, port, is_secure, **options): - phost, pport, pauth = get_proxy_info(hostname, is_secure, **options) +def _get_addrinfo_list(hostname, port, is_secure, proxy): + phost, pport, pauth = get_proxy_info(hostname, is_secure, + proxy.host, proxy.port, proxy.auth, proxy.no_proxy) if not phost: addrinfo_list = socket.getaddrinfo(hostname, port, 0, 0, socket.SOL_TCP) return addrinfo_list, False, None diff --git a/websocket/_socket.py b/websocket/_socket.py index 1abf834..233d076 100644 --- a/websocket/_socket.py +++ b/websocket/_socket.py @@ -38,9 +38,18 @@ if hasattr(socket, "TCP_KEEPCNT"): _default_timeout = None -__all__ = ["DEFAULT_SOCKET_OPTION", "setdefaulttimeout", "getdefaulttimeout", +__all__ = ["DEFAULT_SOCKET_OPTION", "sock_opt", "setdefaulttimeout", "getdefaulttimeout", "recv", "recv_line", "send"] +class sock_opt(object): + def __init__(self, sockopt, sslopt): + if sockopt is None: + sockopt = [] + if sslopt is None: + sslopt = {} + self.sockopt = sockopt + self.sslopt = sslopt + self.timeout = None def setdefaulttimeout(timeout): """ diff --git a/websocket/_url.py b/websocket/_url.py index a8a5127..703b715 100644 --- a/websocket/_url.py +++ b/websocket/_url.py @@ -82,7 +82,8 @@ def _is_no_proxy_host(hostname, no_proxy): return hostname in no_proxy -def get_proxy_info(hostname, is_secure, **options): +def get_proxy_info(hostname, is_secure, + proxy_host=None, proxy_port=0, proxy_auth=None, no_proxy=None): """ try to retrieve proxy host and port from environment if not provided in options. @@ -103,14 +104,13 @@ def get_proxy_info(hostname, is_secure, **options): tuple of username and password. defualt is None """ - if _is_no_proxy_host(hostname, options.get("http_no_proxy", None)): + if _is_no_proxy_host(hostname, no_proxy): return None, 0, None - http_proxy_host = options.get("http_proxy_host", None) - if http_proxy_host: - port = options.get("http_proxy_port", 0) - auth = options.get("http_proxy_auth", None) - return http_proxy_host, port, auth + if proxy_host: + port = proxy_port + auth = proxy_auth + return proxy_host, port, auth env_keys = ["http_proxy"] if is_secure: diff --git a/websocket/tests/test_websocket.py b/websocket/tests/test_websocket.py index 983ccf1..2fcb78f 100644 --- a/websocket/tests/test_websocket.py +++ b/websocket/tests/test_websocket.py @@ -565,23 +565,23 @@ class ProxyInfoTest(unittest.TestCase): def testProxyFromArgs(self): - self.assertEqual(get_proxy_info("echo.websocket.org", False, http_proxy_host="localhost"), ("localhost", 0, None)) - self.assertEqual(get_proxy_info("echo.websocket.org", False, http_proxy_host="localhost", http_proxy_port=3128), ("localhost", 3128, None)) - self.assertEqual(get_proxy_info("echo.websocket.org", True, http_proxy_host="localhost"), ("localhost", 0, None)) - self.assertEqual(get_proxy_info("echo.websocket.org", True, http_proxy_host="localhost", http_proxy_port=3128), ("localhost", 3128, None)) + self.assertEqual(get_proxy_info("echo.websocket.org", False, proxy_host="localhost"), ("localhost", 0, None)) + self.assertEqual(get_proxy_info("echo.websocket.org", False, proxy_host="localhost", proxy_port=3128), ("localhost", 3128, None)) + self.assertEqual(get_proxy_info("echo.websocket.org", True, proxy_host="localhost"), ("localhost", 0, None)) + self.assertEqual(get_proxy_info("echo.websocket.org", True, proxy_host="localhost", proxy_port=3128), ("localhost", 3128, None)) - self.assertEqual(get_proxy_info("echo.websocket.org", False, http_proxy_host="localhost", http_proxy_auth=("a", "b")), + self.assertEqual(get_proxy_info("echo.websocket.org", False, proxy_host="localhost", proxy_auth=("a", "b")), ("localhost", 0, ("a", "b"))) - self.assertEqual(get_proxy_info("echo.websocket.org", False, http_proxy_host="localhost", http_proxy_port=3128, http_proxy_auth=("a", "b")), + self.assertEqual(get_proxy_info("echo.websocket.org", False, proxy_host="localhost", proxy_port=3128, proxy_auth=("a", "b")), ("localhost", 3128, ("a", "b"))) - self.assertEqual(get_proxy_info("echo.websocket.org", True, http_proxy_host="localhost", http_proxy_auth=("a", "b")), + self.assertEqual(get_proxy_info("echo.websocket.org", True, proxy_host="localhost", proxy_auth=("a", "b")), ("localhost", 0, ("a", "b"))) - self.assertEqual(get_proxy_info("echo.websocket.org", True, http_proxy_host="localhost", http_proxy_port=3128, http_proxy_auth=("a", "b")), + self.assertEqual(get_proxy_info("echo.websocket.org", True, proxy_host="localhost", proxy_port=3128, proxy_auth=("a", "b")), ("localhost", 3128, ("a", "b"))) - self.assertEqual(get_proxy_info("echo.websocket.org", True, http_proxy_host="localhost", http_proxy_port=3128, http_no_proxy=["example.com"], http_proxy_auth=("a", "b")), + self.assertEqual(get_proxy_info("echo.websocket.org", True, proxy_host="localhost", proxy_port=3128, no_proxy=["example.com"], proxy_auth=("a", "b")), ("localhost", 3128, ("a", "b"))) - self.assertEqual(get_proxy_info("echo.websocket.org", True, http_proxy_host="localhost", http_proxy_port=3128, http_no_proxy=["echo.websocket.org"], http_proxy_auth=("a", "b")), + self.assertEqual(get_proxy_info("echo.websocket.org", True, proxy_host="localhost", proxy_port=3128, no_proxy=["echo.websocket.org"], proxy_auth=("a", "b")), (None, 0, None))