Grasscutters/mitmproxy
3
0
Fork 0
mirror of https://github.com/Grasscutters/mitmproxy.git synced 2024-11-23 00:01:36 +00:00
Code Issues Projects Releases Wiki Activity

Add script hooks, enable new engine for mitmdump.

Browse Source
This commit is contained in:
Aldo Cortesi 2011-08-03 16:36:20 +12:00
parent f7e4e89b12
commit 179cf75862
9 changed files with 117 additions and 110 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
examples/simple_script
View File

@ -1,6 +1,6 @@
#!/usr/bin/env python
from libmproxy import script
f = script.load_flow()
f.request.headers["newheader"] = ["foo"]
script.return_flow(f)
def response(ctx, f):
ctx.log("processing a response")
f.response.headers["newheader"] = ["foo"]

27
libmproxy/dump.py
View File

@ -1,5 +1,5 @@
import sys, os
import flow, filt, utils
import flow, filt, utils, script
class DumpError(Exception): pass
@ -15,8 +15,6 @@ class Options(object):
"kill",
"no_server",
"refresh_server_playback",
"request_script",
"response_script",
"rfile",
"rheaders",
"server_replay",
@ -68,11 +66,6 @@ class DumpMaster(flow.FlowMaster):
else:
self.filt = None
if self.o.response_script:
self.set_response_script(self.o.response_script)
if self.o.request_script:
self.set_request_script(self.o.request_script)
if options.stickycookie:
self.set_stickycookie(options.stickycookie)
@ -109,6 +102,10 @@ class DumpMaster(flow.FlowMaster):
not options.keepserving
)
if options.script:
err = self.load_script(options.script)
if err:
raise DumpError(err)
def _readflow(self, path):
path = os.path.expanduser(path)
@ -119,20 +116,6 @@ class DumpMaster(flow.FlowMaster):
raise DumpError(v.strerror)
return flows
def _runscript(self, f, script):
try:
ret = f.run_script(script)
if self.o.verbosity > 0:
print >> self.outfile, ret
except flow.RunException, e:
if e.errout:
eout = "Script output:\n" + self.indent(4, e.errout) + "\n"
else:
eout = ""
raise DumpError(
"%s: %s\n%s"%(script, e.args[0], eout)
)
def add_event(self, e, level="info"):
if self.eventlog:
print >> self.outfile, e

84
libmproxy/flow.py
View File

@ -3,7 +3,7 @@
with their responses, and provide filtering and interception facilities.
"""
import subprocess, sys, json, hashlib, Cookie, cookielib
import proxy, threading, netstring, filt
import proxy, threading, netstring, filt, script
import controller, version
class RunException(Exception):
@ -191,10 +191,6 @@ class Flow:
f.load_state(state)
return f
def script_serialize(self):
data = self.get_state()
return json.dumps(data)
@classmethod
def script_deserialize(klass, data):
try:
@ -203,41 +199,6 @@ class Flow:
return None
return klass.from_state(data)
def run_script(self, path):
"""
Run a script on a flow.
Returns a (flow, stderr output) tuple, or raises RunException if
there's an error.
"""
self.backup()
data = self.script_serialize()
try:
p = subprocess.Popen(
[path],
stdout=subprocess.PIPE,
stdin=subprocess.PIPE,
stderr=subprocess.PIPE,
)
except OSError, e:
raise RunException(e.args[1], None, None)
so, se = p.communicate(data)
if p.returncode:
raise RunException(
"Script returned error code %s"%p.returncode,
p.returncode,
se
)
f = Flow.script_deserialize(so)
if not f:
raise RunException(
"Invalid response from script.",
p.returncode,
se
)
self.load_state(f.get_state())
return se
def get_state(self, nobackup=False):
d = dict(
request = self.request.get_state() if self.request else None,
@ -463,7 +424,7 @@ class FlowMaster(controller.Master):
self.server_playback = None
self.client_playback = None
self.kill_nonreplay = False
self.plugin = None
self.script = None
self.stickycookie_state = False
self.stickycookie_txt = None
@ -476,19 +437,26 @@ class FlowMaster(controller.Master):
self.autodecode = False
self.refresh_server_playback = False
def _runscript(self, f, script):
#begin nocover
raise NotImplementedError
#end nocover
def add_event(self, e, level="info"):
"""
level: info, error
"""
pass
def set_plugin(self, p):
self.plugin = p
def load_script(self, path):
"""
Loads a script. Returns an error description if something went
wrong.
"""
s = script.Script(path, self)
try:
s.load()
except script.ScriptError, v:
return v.args[0]
ret = s.run("start")
if not ret[0] and ret[1]:
return "Error in script start:\n\n" + ret[1][1]
self.script = s
def set_stickycookie(self, txt):
if txt:
@ -620,11 +588,20 @@ class FlowMaster(controller.Master):
rt.start()
#end nocover
def handle_clientconnect(self, r):
self.add_event("Connect from: %s:%s"%r.address)
r.ack()
def run_script(self, name, *args, **kwargs):
if self.script:
ret = self.script.run(name, *args, **kwargs)
if not ret[0] and ret[1]:
e = "Script error:\n" + ret[1][1]
self.add_event(e, "error")
def handle_clientconnect(self, cc):
self.run_script("clientconnect", cc)
self.add_event("Connect from: %s:%s"%cc.address)
cc.ack()
def handle_clientdisconnect(self, r):
self.run_script("clientdisconnect", r)
s = "Disconnect from: %s:%s"%r.client_conn.address
self.add_event(s)
if r.client_conn.requestcount:
@ -638,6 +615,8 @@ class FlowMaster(controller.Master):
def handle_error(self, r):
f = self.state.add_error(r)
if f:
self.run_script("error", f)
if self.client_playback:
self.client_playback.clear(f)
r.ack()
@ -645,11 +624,14 @@ class FlowMaster(controller.Master):
def handle_request(self, r):
f = self.state.add_request(r)
self.run_script("request", f)
self.process_new_request(f)
return f
def handle_response(self, r):
f = self.state.add_response(r)
if f:
self.run_script("response", f)
if self.client_playback:
self.client_playback.clear(f)
if not f:

9
libmproxy/script.py
View File

@ -8,7 +8,7 @@ class Context:
self.master, self.state = master, state
def log(self, *args, **kwargs):
self.master.log(*args, **kwargs)
self.master.add_event(*args, **kwargs)
class Script:
@ -32,9 +32,14 @@ class Script:
Raises ScriptError on failure, with argument equal to an error
message that may be a formatted traceback.
"""
path = os.path.expanduser(self.path)
if not os.path.exists(path):
raise ScriptError("No such file: %s"%self.path)
if not os.path.isfile(path):
raise ScriptError("Not a file: %s"%self.path)
ns = {}
try:
self.mod = execfile(os.path.expanduser(self.path), {}, ns)
self.mod = execfile(path, ns, ns)
except Exception, v:
raise ScriptError(traceback.format_exc(v))
self.ns = ns

20
test/scripts/all.py Normal file
View File

@ -0,0 +1,20 @@
log = []
def clientconnect(ctx, cc):
ctx.log("XCLIENTCONNECT")
log.append("clientconnect")
def request(ctx, r):
ctx.log("XREQUEST")
log.append("request")
def response(ctx, r):
ctx.log("XRESPONSE")
log.append("response")
def clientdisconnect(ctx, cc):
ctx.log("XCLIENTDISCONNECT")
log.append("clientdisconnect")
def error(ctx, cc):
ctx.log("XERROR")
log.append("error")

3
test/scripts/starterr.py Normal file
View File

@ -0,0 +1,3 @@
def start(ctx):
raise ValueError

30
test/test_dump.py
View File

@ -111,30 +111,22 @@ class uDumpMaster(libpry.AutoTree):
wfile = "nonexistentdir/foo"
)
def test_request_script(self):
ret = self._dummy_cycle(1, None, "", request_script="scripts/a", verbosity=1)
assert "TESTOK" in ret
assert "DEBUG" in ret
def test_script(self):
ret = self._dummy_cycle(
1, None, "",
script="scripts/all.py", verbosity=0, eventlog=True
)
assert "XCLIENTCONNECT" in ret
assert "XREQUEST" in ret
assert "XRESPONSE" in ret
assert "XCLIENTDISCONNECT" in ret
libpry.raises(
dump.DumpError,
self._dummy_cycle, 1, None, "", request_script="nonexistent"
self._dummy_cycle, 1, None, "", script="nonexistent"
)
libpry.raises(
dump.DumpError,
self._dummy_cycle, 1, None, "", request_script="scripts/err_return"
)
def test_response_script(self):
ret = self._dummy_cycle(1, None, "", response_script="scripts/a", verbosity=1)
assert "TESTOK" in ret
assert "DEBUG" in ret
libpry.raises(
dump.DumpError,
self._dummy_cycle, 1, None, "", response_script="nonexistent"
)
libpry.raises(
dump.DumpError,
self._dummy_cycle, 1, None, "", response_script="scripts/err_return"
self._dummy_cycle, 1, None, "", script="starterr.py"
)
def test_stickycookie(self):

36
test/test_flow.py
View File

@ -138,15 +138,6 @@ class uFlow(libpry.AutoTree):
assert "DEBUG" == se.strip()
assert f.request.host == "TESTOK"
def test_run_script_err(self):
f = tutils.tflow()
f.response = tutils.tresp()
f.request = f.response.request
libpry.raises("returned error", f.run_script,"scripts/err_return")
libpry.raises("invalid response", f.run_script,"scripts/err_data")
libpry.raises("no such file", f.run_script,"nonexistent")
libpry.raises("permission denied", f.run_script,"scripts/nonexecutable")
def test_match(self):
f = tutils.tflow()
f.response = tutils.tresp()
@ -449,13 +440,38 @@ class uSerialize(libpry.AutoTree):
class uFlowMaster(libpry.AutoTree):
def test_load_script(self):
s = flow.State()
fm = flow.FlowMaster(None, s)
assert not fm.load_script("scripts/a.py")
assert fm.load_script("nonexistent")
assert "ValueError" in fm.load_script("scripts/starterr.py")
def test_script(self):
s = flow.State()
fm = flow.FlowMaster(None, s)
assert not fm.load_script("scripts/all.py")
req = tutils.treq()
fm.handle_clientconnect(req.client_conn)
assert fm.script.ns["log"][-1] == "clientconnect"
f = fm.handle_request(req)
assert fm.script.ns["log"][-1] == "request"
resp = tutils.tresp(req)
fm.handle_response(resp)
assert fm.script.ns["log"][-1] == "response"
dc = proxy.ClientDisconnect(req.client_conn)
fm.handle_clientdisconnect(dc)
assert fm.script.ns["log"][-1] == "clientdisconnect"
err = proxy.Error(f.request, "msg")
fm.handle_error(err)
assert fm.script.ns["log"][-1] == "error"
def test_all(self):
s = flow.State()
fm = flow.FlowMaster(None, s)
fm.anticache = True
fm.anticomp = True
req = tutils.treq()
fm.handle_clientconnect(req.client_conn)
f = fm.handle_request(req)

12
test/test_script.py
View File

@ -27,17 +27,23 @@ class uScript(libpry.AutoTree):
s = script.Script("nonexistent", fm)
libpry.raises(
script.ScriptError,
"no such file",
s.load
)
s = script.Script(os.path.join("scripts", "syntaxerr.py"), fm)
s = script.Script("scripts", fm)
libpry.raises(
"not a file",
s.load
)
s = script.Script("scripts/syntaxerr.py", fm)
libpry.raises(
script.ScriptError,
s.load
)
s = script.Script(os.path.join("scripts", "loaderr.py"), fm)
s = script.Script("scripts/loaderr.py", fm)
libpry.raises(
script.ScriptError,
s.load