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(len(os.listdir(self.base_dir)), 1)
self.assertTrue(os.path.exists(os.path.join(self.base_dir, '.seleniumwire', 'teststorage')))
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))
# 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.get_event_loop()
if self._event_loop.is_closed():
# The event loop may be closed if the server had previously
# been shutdown and then spun up again
self._event_loop = asyncio.new_event_loop()
asyncio.set_event_loop(self._event_loop)
# mitmproxy specific options
mitmproxy_opts = Options(
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(
confdir=DEFAULT_CONFDIR,
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 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 proxy server
raise ValueError('Cannot specify both http AND https '
'proxy 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