Rooba
/
mitmproxy
mirror of https://github.com/mitmproxy/mitmproxy.git
1
1
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #1132 from mhils/improve-script-loading 

Improve script loading
This commit is contained in:
Thomas Kriechbaumer 2016-05-11 11:58:44 -05:00
parent 4c8e334642 67a37e6d1f
commit e34b2d7e65
9 changed files with 76 additions and 54 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
mitmproxy/console/__init__.py
View File

@ -19,7 +19,7 @@ from netlib import tcp
from .. import flow, script, contentviews
from . import flowlist, flowview, help, window, signals, options
from . import grideditor, palettes, statusbar, palettepicker
from ..exceptions import FlowReadException
from ..exceptions import FlowReadException, ScriptException
EVENTLOG_SIZE = 500
@ -229,9 +229,10 @@ class ConsoleMaster(flow.FlowMaster):
if options.scripts:
for i in options.scripts:
err = self.load_script(i)
if err:
print("Script load error: {}".format(err), file=sys.stderr)
try:
self.load_script(i)
except ScriptException as e:
print("Script load error: {}".format(e), file=sys.stderr)
sys.exit(1)
if options.outfile:
@ -320,11 +321,11 @@ class ConsoleMaster(flow.FlowMaster):
try:
s = script.Script(command, script.ScriptContext(self))
s.load()
except script.ScriptException as v:
except script.ScriptException as e:
signals.status_message.send(
message = "Error loading script."
message='Error loading "{}".'.format(command)
)
signals.add_event("Error loading script:\n%s" % v.args[0], "error")
signals.add_event('Error loading "{}":\n{}'.format(command, e), "error")
return
if f.request:
@ -336,13 +337,6 @@ class ConsoleMaster(flow.FlowMaster):
s.unload()
signals.flow_change.send(self, flow = f)
def set_script(self, command):
if not command:
return
ret = self.load_script(command)
if ret:
signals.status_message.send(message=ret)
def toggle_eventlog(self):
self.eventlog = not self.eventlog
signals.pop_view_state.send(self)
@ -670,7 +664,13 @@ class ConsoleMaster(flow.FlowMaster):
self.unload_scripts()
for command in commands:
self.load_script(command)
try:
self.load_script(command)
except ScriptException as e:
signals.status_message.send(
message='Error loading "{}".'.format(command)
)
signals.add_event('Error loading "{}":\n{}'.format(command, e), "error")
signals.update_settings.send(self)
def stop_client_playback_prompt(self, a):

4
mitmproxy/console/grideditor.py
View File

@ -642,8 +642,8 @@ class ScriptEditor(GridEditor):
def is_error(self, col, val):
try:
script.Script.parse_command(val)
except script.ScriptException as v:
return str(v)
except script.ScriptException as e:
return str(e)
class HostPatternEditor(GridEditor):

9
mitmproxy/dump.py
View File

@ -7,7 +7,7 @@ import itertools
from netlib import tcp
import netlib.utils
from . import flow, filt, contentviews
from .exceptions import ContentViewException, FlowReadException
from .exceptions import ContentViewException, FlowReadException, ScriptException
class DumpError(Exception):
@ -125,9 +125,10 @@ class DumpMaster(flow.FlowMaster):
scripts = options.scripts or []
for command in scripts:
err = self.load_script(command, use_reloader=True)
if err:
raise DumpError(err)
try:
self.load_script(command, use_reloader=True)
except ScriptException as e:
raise DumpError(str(e))
if options.rfile:
try:

22
mitmproxy/exceptions.py
View File

@ -7,6 +7,10 @@ See also: http://lucumr.pocoo.org/2014/10/16/on-error-handling/
"""
from __future__ import (absolute_import, print_function, division)
import traceback
import sys
class ProxyException(Exception):
"""
@ -59,7 +63,23 @@ class ReplayException(ProxyException):
class ScriptException(ProxyException):
pass
@classmethod
def from_exception_context(cls, cut_tb=1):
"""
Must be called while the current stack handles an exception.
Args:
cut_tb: remove N frames from the stack trace to hide internal calls.
"""
exc_type, exc_value, exc_traceback = sys.exc_info()
while cut_tb > 0:
exc_traceback = exc_traceback.tb_next
cut_tb -= 1
tb = "".join(traceback.format_exception(exc_type, exc_value, exc_traceback))
return cls(tb)
class FlowReadException(ProxyException):

19
mitmproxy/flow.py
View File

@ -695,14 +695,13 @@ class FlowMaster(controller.ServerMaster):
def load_script(self, command, use_reloader=False):
"""
Loads a script. Returns an error description if something went
wrong.
Loads a script.
Raises:
ScriptException
"""
try:
s = script.Script(command, script.ScriptContext(self))
s.load()
except script.ScriptException as e:
return traceback.format_exc(e)
s = script.Script(command, script.ScriptContext(self))
s.load()
if use_reloader:
script.reloader.watch(s, lambda: self.event_queue.put(("script_change", s)))
self.scripts.append(s)
@ -712,7 +711,7 @@ class FlowMaster(controller.ServerMaster):
try:
script_obj.run(name, *args, **kwargs)
except script.ScriptException as e:
self.add_event("Script error:\n" + str(e), "error")
self.add_event("Script error:\n{}".format(e), "error")
def run_script_hook(self, name, *args, **kwargs):
for script_obj in self.scripts:
@ -1069,12 +1068,12 @@ class FlowMaster(controller.ServerMaster):
s.unload()
except script.ScriptException as e:
ok = False
self.add_event('Error reloading "{}": {}'.format(s.filename, str(e)), 'error')
self.add_event('Error reloading "{}":\n{}'.format(s.filename, e), 'error')
try:
s.load()
except script.ScriptException as e:
ok = False
self.add_event('Error reloading "{}": {}'.format(s.filename, str(e)), 'error')
self.add_event('Error reloading "{}":\n{}'.format(s.filename, e), 'error')
else:
self.add_event('"{}" reloaded.'.format(s.filename), 'info')
return ok

8
mitmproxy/script/script.py
View File

@ -79,10 +79,10 @@ class Script(object):
with open(self.filename) as f:
code = compile(f.read(), self.filename, 'exec')
exec (code, self.ns, self.ns)
except Exception as e:
except Exception:
six.reraise(
ScriptException,
ScriptException(str(e)),
ScriptException.from_exception_context(),
sys.exc_info()[2]
)
finally:
@ -113,10 +113,10 @@ class Script(object):
if f:
try:
return f(self.ctx, *args, **kwargs)
except Exception as e:
except Exception:
six.reraise(
ScriptException,
ScriptException(str(e)),
ScriptException.from_exception_context(),
sys.exc_info()[2]
)
else:

6
test/mitmproxy/test_examples.py
View File

@ -106,8 +106,8 @@ def test_modify_querystring():
def test_modify_response_body():
with tutils.raises(script.ScriptException):
with example("modify_response_body.py") as ex:
pass
with example("modify_response_body.py"):
assert True
flow = tutils.tflow(resp=netutils.tresp(content="I <3 mitmproxy"))
with example("modify_response_body.py mitmproxy rocks") as ex:
@ -125,7 +125,7 @@ def test_redirect_requests():
def test_har_extractor():
with tutils.raises(script.ScriptException):
with example("har_extractor.py") as ex:
with example("har_extractor.py"):
pass
times = dict(

26
test/mitmproxy/test_flow.py
View File

@ -7,7 +7,7 @@ import netlib.utils
from netlib import odict
from netlib.http import Headers
from mitmproxy import filt, controller, tnetstring, flow
from mitmproxy.exceptions import FlowReadException
from mitmproxy.exceptions import FlowReadException, ScriptException
from mitmproxy.models import Error
from mitmproxy.models import Flow
from mitmproxy.models import HTTPFlow
@ -747,12 +747,16 @@ class TestFlowMaster:
def test_load_script(self):
s = flow.State()
fm = flow.FlowMaster(None, s)
assert not fm.load_script(tutils.test_data.path("scripts/a.py"))
assert not fm.load_script(tutils.test_data.path("scripts/a.py"))
assert not fm.unload_scripts()
assert fm.load_script("nonexistent")
assert "ValueError" in fm.load_script(
tutils.test_data.path("scripts/starterr.py"))
fm.load_script(tutils.test_data.path("scripts/a.py"))
fm.load_script(tutils.test_data.path("scripts/a.py"))
fm.unload_scripts()
with tutils.raises(ScriptException):
fm.load_script("nonexistent")
try:
fm.load_script(tutils.test_data.path("scripts/starterr.py"))
except ScriptException as e:
assert "ValueError" in str(e)
assert len(fm.scripts) == 0
def test_getset_ignore(self):
@ -779,7 +783,7 @@ class TestFlowMaster:
def test_script_reqerr(self):
s = flow.State()
fm = flow.FlowMaster(None, s)
assert not fm.load_script(tutils.test_data.path("scripts/reqerr.py"))
fm.load_script(tutils.test_data.path("scripts/reqerr.py"))
f = tutils.tflow()
fm.handle_clientconnect(f.client_conn)
assert fm.handle_request(f)
@ -787,7 +791,7 @@ class TestFlowMaster:
def test_script(self):
s = flow.State()
fm = flow.FlowMaster(None, s)
assert not fm.load_script(tutils.test_data.path("scripts/all.py"))
fm.load_script(tutils.test_data.path("scripts/all.py"))
f = tutils.tflow(resp=True)
fm.handle_clientconnect(f.client_conn)
@ -799,7 +803,7 @@ class TestFlowMaster:
fm.handle_response(f)
assert fm.scripts[0].ns["log"][-1] == "response"
# load second script
assert not fm.load_script(tutils.test_data.path("scripts/all.py"))
fm.load_script(tutils.test_data.path("scripts/all.py"))
assert len(fm.scripts) == 2
fm.handle_clientdisconnect(f.server_conn)
assert fm.scripts[0].ns["log"][-1] == "clientdisconnect"
@ -808,7 +812,7 @@ class TestFlowMaster:
# unload first script
fm.unload_scripts()
assert len(fm.scripts) == 0
assert not fm.load_script(tutils.test_data.path("scripts/all.py"))
fm.load_script(tutils.test_data.path("scripts/all.py"))
f.error = tutils.terr()
fm.handle_error(f)

6
test/mitmproxy/test_server.py
View File

@ -285,8 +285,7 @@ class TestHTTP(tservers.HTTPProxyTest, CommonMixin, AppMixin):
self.master.set_stream_large_bodies(None)
def test_stream_modify(self):
self.master.load_script(
tutils.test_data.path("scripts/stream_modify.py"))
self.master.load_script(tutils.test_data.path("scripts/stream_modify.py"))
d = self.pathod('200:b"foo"')
assert d.content == "bar"
self.master.unload_scripts()
@ -511,8 +510,7 @@ class TestTransparent(tservers.TransparentProxyTest, CommonMixin, TcpMixin):
ssl = False
def test_tcp_stream_modify(self):
self.master.load_script(
tutils.test_data.path("scripts/tcp_stream_modify.py"))
self.master.load_script(tutils.test_data.path("scripts/tcp_stream_modify.py"))
self._tcpproxy_on()
d = self.pathod('200:b"foo"')