From b66e1b6e8c504f9866593621f755f64b262624b0 Mon Sep 17 00:00:00 2001 From: Krassimir Valev Date: Sun, 26 Jan 2020 19:55:06 +0100 Subject: [PATCH 1/2] Basic websocket support --- pook/interceptors/__init__.py | 5 +- pook/interceptors/websocket.py | 115 ++++++++++++++++++++++ pook/matchers/url.py | 4 +- requirements-dev.txt | 1 + tests/unit/interceptors/websocket_test.py | 13 +++ tests/unit/matchers/url_test.py | 3 + 6 files changed, 138 insertions(+), 3 deletions(-) create mode 100644 pook/interceptors/websocket.py create mode 100644 tests/unit/interceptors/websocket_test.py diff --git a/pook/interceptors/__init__.py b/pook/interceptors/__init__.py index ffa9000..089477b 100644 --- a/pook/interceptors/__init__.py +++ b/pook/interceptors/__init__.py @@ -1,4 +1,5 @@ import sys +from .websocket import WebSocketInterceptor from .urllib3 import Urllib3Interceptor from .http import HTTPClientInterceptor from .base import BaseInterceptor @@ -7,6 +8,7 @@ __all__ = ( 'interceptors', 'add', 'get', 'BaseInterceptor', + 'WebSocketInterceptor', 'Urllib3Interceptor', 'HTTPClientInterceptor', 'AIOHTTPInterceptor', @@ -15,7 +17,8 @@ # Store built-in interceptors in pook. interceptors = [ Urllib3Interceptor, - HTTPClientInterceptor + HTTPClientInterceptor, + WebSocketInterceptor ] # Import aiohttp in modern Python runtimes diff --git a/pook/interceptors/websocket.py b/pook/interceptors/websocket.py new file mode 100644 index 0000000..a3d48e2 --- /dev/null +++ b/pook/interceptors/websocket.py @@ -0,0 +1,115 @@ +from ..request import Request +from .base import BaseInterceptor + +# Support Python 2/3 +try: + import mock +except Exception: + from unittest import mock + + +class MockFrame(object): + def __init__(self, opcode, data, fin=1): + self.opcode = opcode + self.data = data + self.fin = fin + + +class MockSocket(object): + def fileno(self): + return 0 + + def gettimeout(self): + return 0 + + +class MockHandshakeResponse(object): + + def __init__(self, status, headers, subprotocol): + self.status = status + self.headers = headers + self.subprotocol = subprotocol + + +class WebSocketInterceptor(BaseInterceptor): + + def _handshake(self, sock, hostname, port, resource, **options): + return MockHandshakeResponse(200, {}, "") + + def _connect(self, url, options, proxy, socket): + req = Request() + req.headers = {} # TODO + req.url = url + + # TODO does this work multithreaded?!? + self._mock = self.engine.match(req) + if not self._mock: + # we cannot forward, as we have mocked away the connection + raise ValueError( + "Request to '%s' could not be matched or forwarded" % url + ) + + # make the body always a list to simplify our lives + body = self._mock._response._body + if not isinstance(body, list): + self._mock._response._body = [body] + + sock = MockSocket() + addr = ("hostname", "port", "resource") + + return sock, addr + + def _send(self, data): + # mock as if all data has been sent + return len(data) + + def _recv_frame(self): + # alias + body = self._mock._response._body + + idx = getattr(self, "_data_index", 0) + if len(body) <= idx: + # close frame + return MockFrame(0x8, None) + + # data frame + self._data_index = idx + 1 + return MockFrame(0x1, body[idx].encode("utf-8"), 1) + + def _patch(self, path, handler): + try: + # Create a new patcher for Urllib3 urlopen function + # used as entry point for all the HTTP communications + patcher = mock.patch(path, handler) + # Retrieve original patched function that we might need for real + # networking + # request = patcher.get_original()[0] + # Start patching function calls + patcher.start() + except Exception: + # Exceptions may accur due to missing package + # Ignore all the exceptions for now + pass + else: + self.patchers.append(patcher) + + def activate(self): + """ + Activates the traffic interceptor. + This method must be implemented by any interceptor. + """ + patches = [ + ('websocket._core.connect', self._connect), + ('websocket._core.handshake', self._handshake), + ('websocket.WebSocket._send', self._send), + ('websocket.WebSocket.recv_frame', self._recv_frame), + ] + + [self._patch(path, handler) for path, handler in patches] + + def disable(self): + """ + Disables the traffic interceptor. + This method must be implemented by any interceptor. + """ + [patch.stop() for patch in self.patchers] diff --git a/pook/matchers/url.py b/pook/matchers/url.py index 8939a55..155a469 100644 --- a/pook/matchers/url.py +++ b/pook/matchers/url.py @@ -10,8 +10,8 @@ else: # Python 3 from urllib.parse import urlparse -# URI protocol test regular expression -protoregex = re.compile('^http[s]?://', re.IGNORECASE) +# URI protocol test regular expression (without group capture) +protoregex = re.compile('^(?:ws|http)[s]?://', re.IGNORECASE) class URLMatcher(BaseMatcher): diff --git a/requirements-dev.txt b/requirements-dev.txt index dffee53..3cf40e7 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -12,3 +12,4 @@ urllib3>=1.24.2 bumpversion~=0.5.3 aiohttp~=1.1.5 ; python_version >= '3.4.2' mocket~=1.6.0 +websocket-client~=0.57.0 diff --git a/tests/unit/interceptors/websocket_test.py b/tests/unit/interceptors/websocket_test.py new file mode 100644 index 0000000..642df4c --- /dev/null +++ b/tests/unit/interceptors/websocket_test.py @@ -0,0 +1,13 @@ +from websocket import WebSocket +import pook + + +@pook.on +def test_websocket(): + (pook.get('ws://some-non-existing.org') + .reply(204) + .body('test')) + + socket = WebSocket() + socket.connect('ws://some-non-existing.org/', header=['x-custom: header']) + socket.send('test') diff --git a/tests/unit/matchers/url_test.py b/tests/unit/matchers/url_test.py index 8d0823d..c63cea8 100644 --- a/tests/unit/matchers/url_test.py +++ b/tests/unit/matchers/url_test.py @@ -32,6 +32,8 @@ def test_url_matcher_urlparse(): ('http://foo.com/foo/bar', 'http://foo.com/foo/bar', True), ('http://foo.com/foo/bar/baz', 'http://foo.com/foo/bar/baz', True), ('http://foo.com/foo?x=y&z=w', 'http://foo.com/foo?x=y&z=w', True), + ('ws://foo.com', 'ws://foo.com', True), + ('wss://foo.com', 'wss://foo.com', True), # Invalid cases ('http://foo.com', 'http://bar.com', False), @@ -41,6 +43,7 @@ def test_url_matcher_urlparse(): ('http://foo.com/foo/bar', 'http://foo.com/bar/foo', False), ('http://foo.com/foo/bar/baz', 'http://foo.com/baz/bar/foo', False), ('http://foo.com/foo?x=y&z=w', 'http://foo.com/foo?x=x&y=y', False), + ('ws://foo.com', 'wss://foo.com', False), )) From 63fef0005c78b8c036610fc0bd3813aeb3ad3952 Mon Sep 17 00:00:00 2001 From: Krassimir Valev Date: Wed, 29 Jan 2020 14:59:14 +0100 Subject: [PATCH 2/2] Return a valid file descriptor as a socket --- pook/interceptors/websocket.py | 13 ++++--------- 1 file changed, 4 insertions(+), 9 deletions(-) diff --git a/pook/interceptors/websocket.py b/pook/interceptors/websocket.py index a3d48e2..8855f29 100644 --- a/pook/interceptors/websocket.py +++ b/pook/interceptors/websocket.py @@ -1,3 +1,4 @@ +import tempfile from ..request import Request from .base import BaseInterceptor @@ -15,14 +16,6 @@ def __init__(self, opcode, data, fin=1): self.fin = fin -class MockSocket(object): - def fileno(self): - return 0 - - def gettimeout(self): - return 0 - - class MockHandshakeResponse(object): def __init__(self, status, headers, subprotocol): @@ -54,7 +47,9 @@ def _connect(self, url, options, proxy, socket): if not isinstance(body, list): self._mock._response._body = [body] - sock = MockSocket() + # We have to return a valid file descriptor, not just a random integer. + # Source: https://docs.python.org/3/library/select.html#select.select + sock = tempfile.TemporaryFile() addr = ("hostname", "port", "resource") return sock, addr