def test_err(self):
c = flow.State()
f = tutils.tflow()
c.add_request(f)
f.error = Error("message")
assert c.add_error(f)
c = flow.State()
f = tutils.tflow()
c.add_request(f)
c.set_limit("~e")
assert not c.view
f.error = tutils.terr()
assert c.add_error(f)
assert c.view
def test_flow(self):
"""
normal flow:
connect -> request -> response
"""
bc = tutils.tclient_conn()
c = flow.State()
req = tutils.treq(bc)
f = c.add_request(req)
assert f
assert c.flow_count() == 1
assert c.active_flow_count() == 1
newreq = tutils.treq()
assert c.add_request(newreq)
assert c.active_flow_count() == 2
resp = tutils.tresp(req)
assert c.add_response(resp)
assert c.flow_count() == 2
assert c.active_flow_count() == 1
unseen_resp = tutils.tresp()
unseen_resp.flow = None
assert not c.add_response(unseen_resp)
assert c.active_flow_count() == 1
resp = tutils.tresp(newreq)
assert c.add_response(resp)
assert c.active_flow_count() == 0
def test_stream(self):
with tutils.tmpdir() as tdir:
p = os.path.join(tdir, "foo")
def r():
r = flow.FlowReader(open(p, "rb"))
return list(r.stream())
s = flow.State()
fm = flow.FlowMaster(None, s)
f = tutils.tflow(resp=True)
fm.start_stream(file(p, "ab"), None)
fm.handle_request(f)
fm.handle_response(f)
fm.stop_stream()
assert r()[0].response
f = tutils.tflow()
fm.start_stream(file(p, "ab"), None)
fm.handle_request(f)
fm.shutdown()
assert not r()[1].response
def test_script(self):
s = flow.State()
fm = flow.FlowMaster(None, s)
assert not fm.load_script(tutils.test_data.path("scripts/all.py"))
f = tutils.tflow(resp=True)
fm.handle_clientconnect(f.client_conn)
assert fm.scripts[0].ns["log"][-1] == "clientconnect"
fm.handle_serverconnect(f.server_conn)
assert fm.scripts[0].ns["log"][-1] == "serverconnect"
fm.handle_request(f)
assert fm.scripts[0].ns["log"][-1] == "request"
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"))
assert len(fm.scripts) == 2
fm.handle_clientdisconnect(f.server_conn)
assert fm.scripts[0].ns["log"][-1] == "clientdisconnect"
assert fm.scripts[1].ns["log"][-1] == "clientdisconnect"
#unload first script
fm.unload_scripts()
assert len(fm.scripts) == 0
assert not fm.load_script(tutils.test_data.path("scripts/all.py"))
f.error = tutils.terr()
fm.handle_error(f)
assert fm.scripts[0].ns["log"][-1] == "error"
def test_getset_ignore(self):
p = mock.Mock()
p.config.ignore = []
fm = flow.FlowMaster(p, flow.State())
assert not fm.get_ignore()
fm.set_ignore(["^apple\.com:", ":443$"])
assert fm.get_ignore()
def test_set_intercept(self):
c = flow.State()
assert not c.set_intercept("~q")
assert c.intercept_txt == "~q"
assert "Invalid" in c.set_intercept("~")
assert not c.set_intercept(None)
assert c.intercept_txt == None
def test_getset_ignore(self):
p = mock.Mock()
p.config.check_ignore = HostMatcher()
fm = flow.FlowMaster(p, flow.State())
assert not fm.get_ignore_filter()
fm.set_ignore_filter(["^apple\.com:", ":443$"])
assert fm.get_ignore_filter()
def test_load_flows_reverse(self):
r = self._treader()
s = flow.State()
conf = ProxyConfig(mode="reverse", upstream_server=[True,True,"use-this-domain",80])
fm = flow.FlowMaster(DummyServer(conf), s)
fm.load_flows(r)
assert s._flow_list[0].request.host == "use-this-domain"
def test_err():
s = flow.State()
fm = flow.FlowMaster(None, s)
sc = script.ScriptContext(fm)
tutils.raises(
"not found",
script.Script, "nonexistent", sc
)
tutils.raises(
"not a file",
script.Script, tutils.test_data.path("scripts"), sc
)
tutils.raises(
script.ScriptException,
script.Script, tutils.test_data.path("scripts/syntaxerr.py"), sc
)
tutils.raises(
script.ScriptException,
script.Script, tutils.test_data.path("scripts/loaderr.py"), sc
)
scr = script.Script(tutils.test_data.path("scripts/unloaderr.py"), sc)
tutils.raises(script.ScriptException, scr.unload)
def test_script(self):
s = flow.State()
fm = flow.FlowMaster(None, s)
assert not fm.load_script([tutils.test_data.path("scripts/all.py")])
req = tutils.treq()
fm.handle_clientconnect(req.client_conn)
assert fm.scripts[0].ns["log"][-1] == "clientconnect"
f = fm.handle_request(req)
assert fm.scripts[0].ns["log"][-1] == "request"
resp = tutils.tresp(req)
fm.handle_response(resp)
assert fm.scripts[0].ns["log"][-1] == "response"
#load second script
assert not fm.load_script([tutils.test_data.path("scripts/all.py")])
assert len(fm.scripts) == 2
dc = flow.ClientDisconnect(req.client_conn)
dc.reply = controller.DummyReply()
fm.handle_clientdisconnect(dc)
assert fm.scripts[0].ns["log"][-1] == "clientdisconnect"
assert fm.scripts[1].ns["log"][-1] == "clientdisconnect"
#unload first script
fm.unload_script(fm.scripts[0])
assert len(fm.scripts) == 1
err = flow.Error(f.request, "msg")
err.reply = controller.DummyReply()
fm.handle_error(err)
assert fm.scripts[0].ns["log"][-1] == "error"
def test_concurrent_err():
s = flow.State()
fm = flow.FlowMaster(None, s)
tutils.raises("Concurrent decorator not supported for 'start' method",
script.Script,
tutils.test_data.path("scripts/concurrent_decorator_err.py"),
fm)
def test_client_playback(self):
s = flow.State()
f = tutils.tflow(resp=True)
pb = [tutils.tflow(resp=True), f]
fm = flow.FlowMaster(DummyServer(ProxyConfig()), s)
assert not fm.start_server_playback(
pb,
False,
[],
False,
False,
None,
False,
None,
False)
assert not fm.start_client_playback(pb, False)
fm.client_playback.testing = True
q = Queue.Queue()
assert not fm.state.flow_count()
fm.tick(q, 0)
assert fm.state.flow_count()
f.error = Error("error")
fm.handle_error(f)
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.flow.client_conn)
f = fm.handle_request(req)
assert s.flow_count() == 1
resp = tutils.tresp(req)
fm.handle_response(resp)
assert s.flow_count() == 1
rx = tutils.tresp()
rx.flow = None
assert not fm.handle_response(rx)
fm.handle_clientdisconnect(req.flow.client_conn)
f.error = Error("msg")
f.error.reply = controller.DummyReply()
fm.handle_error(f.error)
fm.load_script(tutils.test_data.path("scripts/a.py"))
fm.shutdown()
def test_script(self):
s = flow.State()
fm = flow.FlowMaster(None, s)
assert not fm.load_script(tutils.test_data.path("scripts/all.py"))
req = tutils.treq()
fm.handle_clientconnect(req.flow.client_conn)
assert fm.scripts[0].ns["log"][-1] == "clientconnect"
sc = ServerConnection((req.get_host(), req.get_port()), None)
sc.reply = controller.DummyReply()
fm.handle_serverconnection(sc)
assert fm.scripts[0].ns["log"][-1] == "serverconnect"
f = fm.handle_request(req)
assert fm.scripts[0].ns["log"][-1] == "request"
resp = tutils.tresp(req)
fm.handle_response(resp)
assert fm.scripts[0].ns["log"][-1] == "response"
#load second script
assert not fm.load_script(tutils.test_data.path("scripts/all.py"))
assert len(fm.scripts) == 2
fm.handle_clientdisconnect(sc)
assert fm.scripts[0].ns["log"][-1] == "clientdisconnect"
assert fm.scripts[1].ns["log"][-1] == "clientdisconnect"
#unload first script
fm.unload_scripts()
assert len(fm.scripts) == 0
assert not fm.load_script(tutils.test_data.path("scripts/all.py"))
err = tutils.terr()
err.reply = controller.DummyReply()
fm.handle_error(err)
assert fm.scripts[0].ns["log"][-1] == "error"
def test_flow(self):
"""
normal flow:
connect -> request -> response
"""
c = flow.State()
f = tutils.tflow()
c.add_flow(f)
assert f
assert c.flow_count() == 1
assert c.active_flow_count() == 1
newf = tutils.tflow()
assert c.add_flow(newf)
assert c.active_flow_count() == 2
f.response = HTTPResponse.wrap(netlib.tutils.tresp())
assert c.update_flow(f)
assert c.flow_count() == 2
assert c.active_flow_count() == 1
_ = HTTPResponse.wrap(netlib.tutils.tresp())
assert not c.update_flow(None)
assert c.active_flow_count() == 1
newf.response = HTTPResponse.wrap(netlib.tutils.tresp())
assert c.update_flow(newf)
assert c.active_flow_count() == 0
def test_concurrent2(self):
s = flow.State()
fm = flow.FlowMaster(None, s)
s = script.Script(
tutils.test_data.path("scripts/concurrent_decorator.py"), fm)
s.load()
m = mock.Mock()
class Dummy:
def __init__(self):
self.response = self
self.error = self
self.reply = m
t_start = time.time()
for hook in ("clientconnect", "serverconnect", "response", "error",
"clientconnect"):
d = Dummy()
assert s.run(hook, d)[0]
d.reply()
while (time.time() - t_start) < 5 and m.call_count <= 5:
if m.call_count == 5:
return
time.sleep(0.001)
assert False
def test_set_limit(self):
c = flow.State()
f = tutils.tflow()
assert len(c.view) == 0
c.add_flow(f)
assert len(c.view) == 1
c.set_limit("~s")
assert c.limit_txt == "~s"
assert len(c.view) == 0
f.response = HTTPResponse.wrap(netlib.tutils.tresp())
c.update_flow(f)
assert len(c.view) == 1
c.set_limit(None)
assert len(c.view) == 1
f = tutils.tflow()
c.add_flow(f)
assert len(c.view) == 2
c.set_limit("~q")
assert len(c.view) == 1
c.set_limit("~s")
assert len(c.view) == 1
assert "Invalid" in c.set_limit("~")
def __init__(self, config):
s = ProxyServer(config, 0)
state = flow.State()
flow.FlowMaster.__init__(self, s, state)
self.apps.add(testapp, "testapp", 80)
self.apps.add(errapp, "errapp", 80)
self.clear_log()
def test_flow(self):
"""
normal flow:
connect -> request -> response
"""
c = flow.State()
f = tutils.tflow()
c.add_request(f)
assert f
assert c.flow_count() == 1
assert c.active_flow_count() == 1
newf = tutils.tflow()
assert c.add_request(newf)
assert c.active_flow_count() == 2
f.response = tutils.tresp()
assert c.add_response(f)
assert c.flow_count() == 2
assert c.active_flow_count() == 1
_ = tutils.tresp()
assert not c.add_response(None)
assert c.active_flow_count() == 1
newf.response = tutils.tresp()
assert c.add_response(newf)
assert c.active_flow_count() == 0
def test_backup(self):
c = flow.State()
req = tutils.treq()
f = c.add_request(req)
f.backup()
c.revert(f)
def test_set_limit(self):
c = flow.State()
f = tutils.tflow()
assert len(c.view) == 0
c.add_request(f)
assert len(c.view) == 1
c.set_limit("~s")
assert c.limit_txt == "~s"
assert len(c.view) == 0
f.response = tutils.tresp()
c.add_response(f)
assert len(c.view) == 1
c.set_limit(None)
assert len(c.view) == 1
f = tutils.tflow()
c.add_request(f)
assert len(c.view) == 2
c.set_limit("~q")
assert len(c.view) == 1
c.set_limit("~s")
assert len(c.view) == 1
assert "Invalid" in c.set_limit("~")
def test_flow(self):
"""
normal flow:
connect -> request -> response
"""
bc = flow.ClientConnect(("address", 22))
c = flow.State()
req = tutils.treq(bc)
f = c.add_request(req)
assert f
assert c.flow_count() == 1
assert c._flow_map.get(req)
assert c.active_flow_count() == 1
newreq = tutils.treq()
assert c.add_request(newreq)
assert c._flow_map.get(newreq)
assert c.active_flow_count() == 2
resp = tutils.tresp(req)
assert c.add_response(resp)
assert c.flow_count() == 2
assert c._flow_map.get(resp.request)
assert c.active_flow_count() == 1
unseen_resp = tutils.tresp()
assert not c.add_response(unseen_resp)
assert not c._flow_map.get(unseen_resp.request)
assert c.active_flow_count() == 1
resp = tutils.tresp(newreq)
assert c.add_response(resp)
assert c.active_flow_count() == 0
def test_clear(self):
c = flow.State()
f = self._add_request(c)
f.intercepting = True
c.clear()
assert c.flow_count() == 0
def test_set_limit(self):
c = flow.State()
req = tutils.treq()
assert len(c.view) == 0
c.add_request(req)
assert len(c.view) == 1
c.set_limit("~s")
assert c.limit_txt == "~s"
assert len(c.view) == 0
resp = tutils.tresp(req)
c.add_response(resp)
assert len(c.view) == 1
c.set_limit(None)
assert len(c.view) == 1
req = tutils.treq()
c.add_request(req)
assert len(c.view) == 2
c.set_limit("~q")
assert len(c.view) == 1
c.set_limit("~s")
assert len(c.view) == 1
assert "Invalid" in c.set_limit("~")
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"))
f = tutils.tflow()
fm.handle_clientconnect(f.client_conn)
assert fm.handle_request(f)
def test_script_reqerr(self):
s = flow.State()
fm = flow.FlowMaster(None, s)
assert not fm.load_script("scripts/reqerr.py")
req = tutils.treq()
fm.handle_clientconnect(req.client_conn)
assert fm.handle_request(req)
def test_server_playback(self):
s = flow.State()
f = tutils.tflow()
f.response = tutils.tresp(f.request)
pb = [f]
fm = flow.FlowMaster(None, s)
fm.refresh_server_playback = True
assert not fm.do_server_playback(tutils.tflow())
fm.start_server_playback(pb, False, [], False, False)
assert fm.do_server_playback(tutils.tflow())
fm.start_server_playback(pb, False, [], True, False)
r = tutils.tflow()
r.request.content = "gibble"
assert not fm.do_server_playback(r)
assert fm.do_server_playback(tutils.tflow())
fm.start_server_playback(pb, False, [], True, False)
q = Queue.Queue()
fm.tick(q, 0)
assert fm.should_exit.is_set()
fm.stop_server_playback()
assert not fm.server_playback
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)
assert s.flow_count() == 1
resp = tutils.tresp(req)
fm.handle_response(resp)
assert s.flow_count() == 1
rx = tutils.tresp()
assert not fm.handle_response(rx)
dc = flow.ClientDisconnect(req.client_conn)
req.client_conn.requestcount = 1
fm.handle_clientdisconnect(dc)
err = flow.Error(f.request, "msg")
fm.handle_error(err)
fm.load_script("scripts/a.py")
fm.shutdown()
def test_tick(self):
first = tutils.tflow()
s = flow.State()
fm = flow.FlowMaster(None, s)
fm.start_client_playback([first, tutils.tflow()], True)
c = fm.client_playback
assert not c.done()
assert not s.flow_count()
assert c.count() == 2
c.tick(fm, testing=True)
assert s.flow_count()
assert c.count() == 1
c.tick(fm, testing=True)
assert c.count() == 1
c.clear(c.current)
c.tick(fm, testing=True)
assert c.count() == 0
c.clear(c.current)
assert c.done()
q = Queue.Queue()
fm.state.clear()
fm.tick(q, timeout=0)
fm.stop_client_playback()
assert not fm.client_playback
def __init__(self, server, options, sessionFactory):
FlowMaster.__init__(self, server, flow.State())
self.sessionFactory = sessionFactory
self.o = options
self.flows = FlowCollection()
self.anticache = options.anticache
self.anticomp = options.anticomp
if options.stickycookie:
self.set_stickycookie(options.stickycookie)
if options.stickyauth:
self.set_stickyauth(options.stickyauth)
if options.wfile:
path = os.path.abspath(os.path.expanduser(options.wfile))
directory = os.path.split(path)[0]
if not os.path.exists(directory):
os.makedirs(directory)
try:
f = file(path, "wb")
self.fwriter = flow.FlowWriter(f)
except IOError, v:
raise Exception(v.strerror)