def test_cleanup_removes_storage(self):
storage = RequestStorage(base_dir=self.base_dir)
storage.cleanup()
# The 'seleniumwire' parent folder should have been cleaned up
# when there is nothing left inside of it.
self.assertFalse(os.listdir(self.base_dir))
def test_cleanup_does_not_remove_parent_folder(self):
# There is an existing storage folder
os.makedirs(os.path.join(self.base_dir, '.seleniumwire', 'teststorage'))
storage = RequestStorage(base_dir=self.base_dir)
storage.cleanup()
# The existing storage folder is not cleaned up
self.assertEqual(1, len(os.listdir(self.base_dir)))
self.assertTrue(os.path.exists(os.path.join(self.base_dir, '.seleniumwire', 'teststorage')))
def test_load_last_request(self):
request_1 = self._create_request()
request_2 = self._create_request()
storage = RequestStorage(base_dir=self.base_dir)
storage.save_request(request_1)
storage.save_request(request_2)
last_request = storage.load_last_request()
self.assertEqual(request_2.id, last_request.id)
def __init__(self, host, port, options):
self.options = options
# Used to stored captured requests
self.storage = RequestStorage(
base_dir=options.pop('request_storage_base_dir', None))
extract_cert_and_key(self.storage.home_dir)
# Used to modify requests/responses passing through the server
# DEPRECATED. Will be superceded by request/response interceptors.
self.modifier = RequestModifier()
# The scope of requests we're interested in capturing.
self.scopes = []
self.request_interceptor = None
self.response_interceptor = None
self._event_loop = asyncio.new_event_loop()
# mitmproxy specific options
mitmproxy_opts = Options(
confdir=self.storage.home_dir,
listen_host=host,
listen_port=port,
)
# Create an instance of the mitmproxy server
self._master = Master(self._event_loop, mitmproxy_opts)
self._master.server = ProxyServer(ProxyConfig(mitmproxy_opts))
self._master.addons.add(*addons.default_addons())
self._master.addons.add(SendToLogger())
self._master.addons.add(InterceptRequestHandler(self))
# Update the options now all addons have been added
mitmproxy_opts.update(
ssl_insecure=options.get('verify_ssl', True),
upstream_cert=DEFAULT_UPSTREAM_CERT,
stream_websockets=DEFAULT_STREAM_WEBSOCKETS,
suppress_connection_errors=options.get(
'suppress_connection_errors', True),
**self._get_upstream_proxy_args(),
)
if options.get('disable_capture', False):
self.scopes = ['$^']
mitmproxy_opts.update(ignore_hosts=['.*'])
# Options that are prefixed mitm_ are passed through to mitmproxy
mitmproxy_opts.update(
**{k[5:]: v
for k, v in options.items() if k.startswith('mitm_')})
def test_save_request_with_body(self):
body = b'test request body'
request = self._create_request(body=body)
storage = RequestStorage(base_dir=self.base_dir)
storage.save_request(request)
request_file_path = self._get_stored_path(request.id, 'request')
with open(request_file_path[0], 'rb') as loaded:
loaded_request = pickle.load(loaded)
self.assertEqual(body, loaded_request.body)
def test_iter_requests(self):
request_1 = self._create_request()
request_2 = self._create_request()
storage = RequestStorage(base_dir=self.base_dir)
storage.save_request(request_1)
storage.save_request(request_2)
requests = storage.iter_requests()
self.assertIsInstance(requests, Iterator)
requests = list(requests)
self.assertEqual(2, len(requests))
self.assertEqual(request_1.id, requests[0].id)
self.assertEqual(request_2.id, requests[1].id)
def test_load_requests(self):
request_1 = self._create_request()
request_2 = self._create_request()
storage = RequestStorage(base_dir=self.base_dir)
storage.save_request(request_1)
storage.save_request(request_2)
requests = storage.load_requests()
self.assertEqual(2, len(requests))
self.assertEqual(request_1.id, requests[0].id)
self.assertEqual(request_2.id, requests[1].id)
self.assertIsNone(requests[0].response)
self.assertIsNone(requests[1].response)
def test_save_har_entry_no_request(self):
request = self._create_request()
storage = RequestStorage(base_dir=self.base_dir)
storage.save_request(request)
storage.clear_requests()
storage.save_har_entry(request.id, {'name': 'test_har_entry'})
har_file_path = self._get_stored_path(request.id, 'har_entry')
self.assertFalse(har_file_path)
def test_save_response_no_request(self):
request = self._create_request()
storage = RequestStorage(base_dir=self.base_dir)
storage.save_request(request)
response = self._create_response()
storage.clear_requests()
storage.save_response(request.id, response)
response_file_path = self._get_stored_path(request.id, 'response')
self.assertFalse(response_file_path)
def test_save_request(self):
request = self._create_request()
storage = RequestStorage(base_dir=self.base_dir)
storage.save_request(request)
request_file_path = self._get_stored_path(request.id, 'request')
with open(request_file_path[0], 'rb') as loaded:
loaded_request = pickle.load(loaded)
self.assertEqual(request.id, loaded_request.id)
self.assertEqual('http://www.example.com/test/path/', loaded_request.url)
self.assertEqual('GET', loaded_request.method)
self.assertEqual({
'Host': 'www.example.com',
'Accept': '*/*'
}, dict(loaded_request.headers))
self.assertIsNone(loaded_request.response)
def test_load_response(self):
request = self._create_request()
storage = RequestStorage(base_dir=self.base_dir)
storage.save_request(request)
response = self._create_response()
storage.save_response(request.id, response)
requests = storage.load_requests()
self.assertIsNotNone(requests[0].response)
def test_initialise_clears_old_folders(self):
old_dir = os.path.join(self.base_dir, '.seleniumwire', 'storage-test1')
new_dir = os.path.join(self.base_dir, '.seleniumwire', 'storage-test2')
os.makedirs(old_dir)
os.makedirs(new_dir)
two_days_ago = (datetime.now() - timedelta(days=2)).timestamp()
os.utime(old_dir, times=(two_days_ago, two_days_ago))
RequestStorage(base_dir=self.base_dir)
self.assertFalse(os.path.exists(old_dir))
self.assertTrue(os.path.exists(new_dir))
def test_load_request_cert_data(self):
storage = RequestStorage(base_dir=self.base_dir)
request = self._create_request()
storage.save_request(request)
response = self._create_response()
response.cert = {'subject': 'test_cert'}
storage.save_response(request.id, response)
requests = storage.load_requests()
self.assertEqual({'subject': 'test_cert'}, requests[0].cert)
def test_clear_requests(self):
request_1 = self._create_request()
request_2 = self._create_request()
storage = RequestStorage(base_dir=self.base_dir)
storage.save_request(request_1)
storage.save_request(request_2)
storage.clear_requests()
requests = storage.load_requests()
self.assertFalse(requests)
self.assertFalse(glob.glob(os.path.join(self.base_dir, '.seleniumwire', 'storage-*', '*')))
def test_load_response_encoded_body_error(self):
body = b'test response body'
request = self._create_request()
storage = RequestStorage(base_dir=self.base_dir)
storage.save_request(request)
response = self._create_response(body=body)
response.headers['Content-Encoding'] = 'gzip'
storage.save_response(request.id, response)
response_body = storage.load_requests()[0].response.body
self.assertEqual(body, response_body)
def test_load_response_encoded_body(self):
body = b'test response body'
io = BytesIO()
with gzip.GzipFile(fileobj=io, mode='wb') as f:
f.write(body)
request = self._create_request()
storage = RequestStorage(base_dir=self.base_dir)
storage.save_request(request)
response = self._create_response(body=io.getvalue())
response.headers['Content-Encoding'] = 'gzip'
storage.save_response(request.id, response)
response_body = storage.load_requests()[0].response.body
self.assertEqual(body, response_body)
def test_save_har_entry(self):
request = self._create_request()
storage = RequestStorage(base_dir=self.base_dir)
storage.save_request(request)
storage.save_har_entry(request.id, {'name': 'test_har_entry'})
har_file_path = self._get_stored_path(request.id, 'har_entry')
with open(har_file_path[0], 'rb') as loaded:
loaded_har = pickle.load(loaded)
self.assertEqual(loaded_har['name'], 'test_har_entry')
def test_save_response_with_body(self):
body = b'some response body'
request = self._create_request()
storage = RequestStorage(base_dir=self.base_dir)
storage.save_request(request)
response = self._create_response(body=body)
storage.save_response(request.id, response)
response_file_path = self._get_stored_path(request.id, 'response')
with open(response_file_path[0], 'rb') as loaded:
loaded_response = pickle.load(loaded)
self.assertEqual(b'some response body', loaded_response.body)
def test_load_request_with_ws_messages(self):
storage = RequestStorage(base_dir=self.base_dir)
request_1 = self._create_request() # Websocket handshake request
request_2 = self._create_request()
storage.save_request(request_1)
storage.save_request(request_2)
storage.save_ws_message(
request_1.id,
WebSocketMessage(
from_client=True,
content='websocket test message',
date=datetime.now(),
))
requests = storage.load_requests()
self.assertTrue(len(requests[0].ws_messages) > 0)
self.assertEqual('websocket test message',
requests[0].ws_messages[0].content)
self.assertTrue(len(requests[1].ws_messages) == 0)
def test_save_response(self):
request = self._create_request()
storage = RequestStorage(base_dir=self.base_dir)
storage.save_request(request)
response = self._create_response()
storage.save_response(request.id, response)
response_file_path = self._get_stored_path(request.id, 'response')
with open(response_file_path[0], 'rb') as loaded:
loaded_response = pickle.load(loaded)
self.assertEqual(200, loaded_response.status_code)
self.assertEqual('OK', loaded_response.reason)
self.assertEqual({
'Content-Type': 'application/json',
'Content-Length': '500'
}, dict(loaded_response.headers))
class MitmProxy:
"""Run and manage a mitmproxy server instance.
DEPRECATED.
"""
def __init__(self, host, port, options):
self.options = options
# Used to stored captured requests
self.storage = RequestStorage(
base_dir=options.pop('request_storage_base_dir', None))
# Used to modify requests/responses passing through the server
# DEPRECATED. Will be superceded by request/response interceptors.
self.modifier = RequestModifier()
# The scope of requests we're interested in capturing.
self.scopes = []
self.request_interceptor = None
self.response_interceptor = None
self._event_loop = self._get_event_loop()
# mitmproxy specific options
mitmproxy_opts = Options(
confdir=options.pop('mitm_confdir', DEFAULT_CONFDIR),
listen_host=host,
listen_port=port,
)
# Create an instance of the mitmproxy server
self._master = Master(mitmproxy_opts)
self._master.server = ProxyServer(ProxyConfig(mitmproxy_opts))
self._master.addons.add(*addons.default_addons())
self._master.addons.add(SendToLogger())
self._master.addons.add(MitmProxyRequestHandler(self))
# Update the options now all addons have been added
mitmproxy_opts.update(
ssl_insecure=options.get('verify_ssl', True),
upstream_cert=DEFAULT_UPSTREAM_CERT,
stream_websockets=DEFAULT_STREAM_WEBSOCKETS,
**self._get_upstream_proxy_args(),
)
# Options that are prefixed mitm_ are passed through to mitmproxy
mitmproxy_opts.update(
**{k[5:]: v
for k, v in options.items() if k.startswith('mitm_')})
def _get_event_loop(self):
try:
event_loop = asyncio.get_event_loop()
if event_loop.is_closed():
# The event loop may be closed if the server had previously
# been shutdown and then spun up again
event_loop = asyncio.new_event_loop()
asyncio.set_event_loop(event_loop)
except Exception:
event_loop = asyncio.new_event_loop()
asyncio.set_event_loop(event_loop)
return event_loop
def serve_forever(self):
"""Run the server."""
asyncio.set_event_loop(self._event_loop)
self._master.run_loop(self._event_loop.run_forever)
def address(self):
"""Get a tuple of the address and port the mitmproxy server
is listening on.
"""
return self._master.server.address
def shutdown(self):
"""Shutdown the server and perform any cleanup."""
try:
# Wait for any active requests to finish. This reduces the
# probability of seeing shutdown errors in the console.
self._master.server.wait_for_silence()
except Timeout:
pass
self._master.shutdown()
self.storage.cleanup()
def _get_upstream_proxy_args(self):
proxy_config = get_upstream_proxy(self.options)
http_proxy = proxy_config.get('http')
https_proxy = proxy_config.get('https')
conf = None
if http_proxy and https_proxy:
if http_proxy.hostport != https_proxy.hostport:
# We only support a single upstream mitmproxy server
raise ValueError('Cannot specify both http AND https '
'mitmproxy settings with mitmproxy backend')
conf = https_proxy
elif http_proxy:
conf = http_proxy
elif https_proxy:
conf = https_proxy
args = {}
if conf:
scheme, username, password, hostport = conf
args['mode'] = 'upstream:{}://{}'.format(scheme, hostport)
if username and password:
args['upstream_auth'] = '{}:{}'.format(username, password)
return args
def test_find_similar_urls(self):
request_1 = self._create_request('https://192.168.1.1/redfish/v1')
request_2 = self._create_request('https://192.168.1.1/redfish')
mock_response = self._create_response()
storage = RequestStorage(base_dir=self.base_dir)
storage.save_request(request_1)
storage.save_response(request_1.id, mock_response)
storage.save_request(request_2)
storage.save_response(request_2.id, mock_response)
self.assertEqual(request_1.id, storage.find('.*v1').id)
self.assertEqual(request_2.id,
storage.find('https://192.168.1.1/redfish$').id)
def test_find(self):
request_1 = self._create_request(
'http://www.example.com/test/path/?foo=bar')
request_2 = self._create_request(
'http://www.stackoverflow.com/other/path/?x=y')
mock_response = self._create_response()
storage = RequestStorage(base_dir=self.base_dir)
storage.save_request(request_1)
storage.save_response(request_1.id, mock_response)
storage.save_request(request_2)
self.assertEqual(request_1.id, storage.find('/test/path/').id)
self.assertEqual(request_1.id,
storage.find(r'/test/path/\?foo=bar').id)
self.assertEqual(
request_1.id,
storage.find(r'http://www.example.com/test/path/\?foo=bar').id)
self.assertEqual(request_1.id,
storage.find(r'http://www.example.com/test/path/').id)
self.assertIsNone(storage.find('/different/path'))
self.assertIsNone(storage.find('/test/path/?x=y'))
self.assertIsNone(
storage.find(r'http://www.example.com/different/path/\?foo=bar'))
self.assertIsNone(
storage.find(r'http://www.different.com/test/path/\?foo=bar'))
self.assertIsNone(
storage.find(r'http://www.example.com/test/path/\?x=y'))
def test_get_session_dir(self):
storage = RequestStorage(base_dir=self.base_dir)
self.assertTrue(
fnmatch(storage.session_dir,
os.path.join(self.base_dir, '.seleniumwire', 'storage-*')))
def test_get_home_dir(self):
storage = RequestStorage(base_dir=self.base_dir)
self.assertEqual(os.path.join(self.base_dir, '.seleniumwire'),
storage.home_dir)
def test_load_last_request_none(self):
storage = RequestStorage(base_dir=self.base_dir)
last_request = storage.load_last_request()
self.assertIsNone(last_request)
def setUp(self):
self.base_dir = os.path.join(os.path.dirname(__file__), 'data')
self.storage = RequestStorage(base_dir=self.base_dir)
class MitmProxy:
"""Run and manage a mitmproxy server instance."""
def __init__(self, host, port, options):
self.options = options
# Used to stored captured requests
self.storage = RequestStorage(
base_dir=options.pop('request_storage_base_dir', None))
extract_cert_and_key(self.storage.home_dir)
# Used to modify requests/responses passing through the server
# DEPRECATED. Will be superceded by request/response interceptors.
self.modifier = RequestModifier()
# The scope of requests we're interested in capturing.
self.scopes = []
self.request_interceptor = None
self.response_interceptor = None
self._event_loop = asyncio.new_event_loop()
# mitmproxy specific options
mitmproxy_opts = Options(
confdir=self.storage.home_dir,
listen_host=host,
listen_port=port,
)
# Create an instance of the mitmproxy server
self._master = Master(self._event_loop, mitmproxy_opts)
self._master.server = ProxyServer(ProxyConfig(mitmproxy_opts))
self._master.addons.add(*addons.default_addons())
self._master.addons.add(SendToLogger())
self._master.addons.add(MitmProxyRequestHandler(self))
# Update the options now all addons have been added
mitmproxy_opts.update(
ssl_insecure=options.get('verify_ssl', True),
upstream_cert=DEFAULT_UPSTREAM_CERT,
stream_websockets=DEFAULT_STREAM_WEBSOCKETS,
suppress_connection_errors=options.get(
'suppress_connection_errors', True),
**self._get_upstream_proxy_args(),
)
# Options that are prefixed mitm_ are passed through to mitmproxy
mitmproxy_opts.update(
**{k[5:]: v
for k, v in options.items() if k.startswith('mitm_')})
def serve_forever(self):
"""Run the server."""
asyncio.set_event_loop(self._event_loop)
self._master.run()
def address(self):
"""Get a tuple of the address and port the proxy server
is listening on.
"""
return self._master.server.address
def shutdown(self):
"""Shutdown the server and perform any cleanup."""
self._master.shutdown()
self.storage.cleanup()
def _get_upstream_proxy_args(self):
proxy_config = get_upstream_proxy(self.options)
http_proxy = proxy_config.get('http')
https_proxy = proxy_config.get('https')
conf = None
if http_proxy and https_proxy:
if http_proxy.hostport != https_proxy.hostport:
# We only support a single upstream proxy server
raise ValueError(
'Different settings for http and https proxy servers not supported'
)
conf = https_proxy
elif http_proxy:
conf = http_proxy
elif https_proxy:
conf = https_proxy
args = {}
if conf:
scheme, username, password, hostport = conf
args['mode'] = 'upstream:{}://{}'.format(scheme, hostport)
if username:
args['upstream_auth'] = '{}:{}'.format(username, password)
custom_auth = proxy_config.get('custom_authorization')
if custom_auth:
args['upstream_custom_auth'] = custom_auth
return args
def test_initialise(self):
RequestStorage(base_dir=self.base_dir)
storage_dir = glob.glob(
os.path.join(self.base_dir, '.seleniumwire', 'storage-*'))
self.assertEqual(1, len(storage_dir))
