def test_stream(self): with tutils.tmpdir() as tdir: p = os.path.join(tdir, "foo") def r(): r = flow.FlowReader(open(p)) return list(r.stream()) s = flow.State() fm = flow.FlowMaster(None, s) tf = tutils.tflow_full() fm.start_stream(file(p, "ab"), None) fm.handle_request(tf.request) fm.handle_response(tf.response) fm.stop_stream() assert r()[0].response tf = tutils.tflow_full() fm.start_stream(file(p, "ab"), None) fm.handle_request(tf.request) fm.shutdown() assert not r()[1].response
def test_handle_response(self): s = flow.StickyAuthState(filt.parse(".*")) f = tutils.tflow_full() f.request.headers["authorization"] = ["foo"] s.handle_request(f) assert "host" in s.hosts f = tutils.tflow_full() s.handle_request(f) assert f.request.headers["authorization"] == ["foo"]
def test_load_with_nopop(self): r = tutils.tflow_full() r.request.headers["key"] = ["one"] r2 = tutils.tflow_full() r2.request.headers["key"] = ["two"] s = flow.ServerPlaybackState(None, [r, r2], False, True) assert s.count() == 2 s.next_flow(r) assert s.count() == 2
def test_headers(self): s = flow.ServerPlaybackState(["foo"], [], False, False) r = tutils.tflow_full() r.request.headers["foo"] = ["bar"] r2 = tutils.tflow_full() assert not s._hash(r) == s._hash(r2) r2.request.headers["foo"] = ["bar"] assert s._hash(r) == s._hash(r2) r2.request.headers["oink"] = ["bar"] assert s._hash(r) == s._hash(r2) r = tutils.tflow_full() r2 = tutils.tflow_full() assert s._hash(r) == s._hash(r2)
def test_client_playback(self): s = flow.State() f = tutils.tflow_full() pb = [tutils.tflow_full(), f] fm = flow.FlowMaster(None, s) assert not fm.start_server_playback(pb, False, [], False) assert not fm.start_client_playback(pb, False) q = Queue.Queue() assert not fm.state.flow_count() fm.tick(q) assert fm.state.flow_count() fm.handle_error(flow.Error(f.request, "error"))
def test_client_playback(self): s = flow.State() f = tutils.tflow_full() pb = [tutils.tflow_full(), f] fm = flow.FlowMaster(None, s) assert not fm.start_server_playback(pb, False, [], False, False) assert not fm.start_client_playback(pb, False) q = Queue.Queue() assert not fm.state.flow_count() fm.tick(q) assert fm.state.flow_count() fm.handle_error(flow.Error(f.request, "error"))
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) tf = tutils.tflow_full() fm.start_stream(file(p, "ab"), None) fm.handle_request(tf.request) fm.handle_response(tf.response) fm.stop_stream() assert r()[0].response tf = tutils.tflow() fm.start_stream(file(p, "ab"), None) fm.handle_request(tf.request) fm.shutdown() assert not r()[1].response
def _response(self, cookie, host): s = flow.StickyCookieState(filt.parse(".*")) f = tutils.tflow_full() f.server_conn.address = tcp.Address((host, 80)) f.response.headers["Set-Cookie"] = [cookie] s.handle_response(f) return s, f
def test_replay(self): cs = StringIO() o = dump.Options(server_replay="nonexistent", kill=True) libpry.raises(dump.DumpError, dump.DumpMaster, None, o, None, outfile=cs) t = self.tmpdir() p = os.path.join(t, "rep") f = open(p, "w") fw = flow.FlowWriter(f) t = tutils.tflow_full() t.response = tutils.tresp(t.request) fw.add(t) f.close() o = dump.Options(server_replay=p, kill=True) m = dump.DumpMaster(None, o, None, outfile=cs) self._cycle(m, "content") self._cycle(m, "content") o = dump.Options(server_replay=p, kill=False) m = dump.DumpMaster(None, o, None, outfile=cs) self._cycle(m, "nonexistent") o = dump.Options(client_replay=p, kill=False) m = dump.DumpMaster(None, o, None, outfile=cs)
def _response(self, cookie, host): s = flow.StickyCookieState(filt.parse(".*")) f = tutils.tflow_full() f.request.host = host f.response.headers["Set-Cookie"] = [cookie] s.handle_response(f) return s, f
def _flowfile(self, path): f = open(path, "w") fw = flow.FlowWriter(f) t = tutils.tflow_full() t.response = tutils.tresp(t.request) fw.add(t) f.close()
def test_copy(self): f = tutils.tflow_full() f2 = f.copy() assert not f is f2 assert not f.request is f2.request assert f.request.headers == f2.request.headers assert not f.request.headers is f2.request.headers
def test_filter(self): sio = StringIO() fl = filt.parse("~c 200") w = flow.FilteredFlowWriter(sio, fl) f = tutils.tflow_full() f.response.code = 200 w.add(f) f = tutils.tflow_full() f.response.code = 201 w.add(f) sio.seek(0) r = flow.FlowReader(sio) assert len(list(r.stream()))
def _flowfile(self, path): f = open(path, "wb") fw = flow.FlowWriter(f) t = tutils.tflow_full() t.response = tutils.tresp(t.request) fw.add(t) f.close()
def test_client_playback(self): s = flow.State() f = tutils.tflow_full() pb = [tutils.tflow_full(), f] fm = flow.FlowMaster(None, s) assert not fm.start_server_playback(pb, False, [], False, False) assert not fm.start_client_playback(pb, False) q = Queue.Queue() assert not fm.state.flow_count() fm.tick(q) assert fm.state.flow_count() f.error = Error("error") f.error.reply = controller.DummyReply() fm.handle_error(f.error)
def test_replay(self): s = flow.State() fm = flow.FlowMaster(None, s) f = tutils.tflow_full() f.request.content = flow.CONTENT_MISSING assert "missing" in fm.replay_request(f) f.intercepting = True assert "intercepting" in fm.replay_request(f)
def test_duplicate_flow(self): s = flow.State() fm = flow.FlowMaster(None, s) f = tutils.tflow_full() fm.load_flow(f) assert s.flow_count() == 1 f2 = fm.duplicate_flow(f) assert f2.response assert s.flow_count() == 2 assert s.index(f2)
def test_stickyauth(self): s = flow.State() fm = flow.FlowMaster(None, s) assert "Invalid" in fm.set_stickyauth("~h") fm.set_stickyauth(".*") assert fm.stickyauth_state fm.set_stickyauth(None) assert not fm.stickyauth_state fm.set_stickyauth(".*") tf = tutils.tflow_full() tf.request.headers["authorization"] = ["foo"] fm.handle_request(tf.request) f = tutils.tflow_full() assert fm.stickyauth_state.hosts assert not "authorization" in f.request.headers fm.handle_request(f.request) assert f.request.headers["authorization"] == ["foo"]
def test_match(self): f = tutils.tflow_full() assert not f.match("~b test") assert f.match(None) assert not f.match("~b test") f = tutils.tflow_err() assert f.match("~e") tutils.raises(ValueError, f.match, "~")
def test_load(self): r = tutils.tflow_full() r.request.headers["key"] = ["one"] r2 = tutils.tflow_full() r2.request.headers["key"] = ["two"] s = flow.ServerPlaybackState(None, [r, r2], False) assert s.count() == 2 assert len(s.fmap.keys()) == 1 n = s.next_flow(r) assert n.request.headers["key"] == ["one"] assert s.count() == 1 n = s.next_flow(r) assert n.request.headers["key"] == ["two"] assert s.count() == 0 assert not s.next_flow(r)
def test_load(self): r = tutils.tflow_full() r.request.headers["key"] = ["one"] r2 = tutils.tflow_full() r2.request.headers["key"] = ["two"] s = flow.ServerPlaybackState(None, [r, r2], False, False) assert s.count() == 2 assert len(s.fmap.keys()) == 1 n = s.next_flow(r) assert n.request.headers["key"] == ["one"] assert s.count() == 1 n = s.next_flow(r) assert n.request.headers["key"] == ["two"] assert s.count() == 0 assert not s.next_flow(r)
def _treader(self): sio = StringIO() w = flow.FlowWriter(sio) for i in range(3): f = tutils.tflow_full() w.add(f) for i in range(3): f = tutils.tflow_err() w.add(f) sio.seek(0) return flow.FlowReader(sio)
def test_setheaders(): h = flow.SetHeaders() h.add("~q", "foo", "bar") assert h.lst h.set( [ (".*", "one", "two"), (".*", "three", "four"), ] ) assert h.count() == 2 h.clear() assert not h.lst h.add("~q", "foo", "bar") h.add("~s", "foo", "bar") v = h.get_specs() assert v == [('~q', 'foo', 'bar'), ('~s', 'foo', 'bar')] assert h.count() == 2 h.clear() assert h.count() == 0 f = tutils.tflow() f.request.content = "foo" h.add("~s", "foo", "bar") h.run(f) assert f.request.content == "foo" h.clear() h.add("~s", "one", "two") h.add("~s", "one", "three") f = tutils.tflow_full() f.request.headers["one"] = ["xxx"] f.response.headers["one"] = ["xxx"] h.run(f) assert f.request.headers["one"] == ["xxx"] assert f.response.headers["one"] == ["two", "three"] h.clear() h.add("~q", "one", "two") h.add("~q", "one", "three") f = tutils.tflow() f.request.headers["one"] = ["xxx"] h.run(f) assert f.request.headers["one"] == ["two", "three"] assert not h.add("~", "foo", "bar")
def test_replace(self): f = tutils.tflow_full() f.request.headers["foo"] = ["foo"] f.request.content = "afoob" f.response.headers["foo"] = ["foo"] f.response.content = "afoob" assert f.replace("foo", "bar") == 6 assert f.request.headers["bar"] == ["bar"] assert f.request.content == "abarb" assert f.response.headers["bar"] == ["bar"] assert f.response.content == "abarb"
def test_concurrent2(self): ctx = TScriptContext() s = script.Script([tutils.test_data.path("scripts/concurrent_decorator.py")], ctx) s.load() f = tutils.tflow_full() f.error = tutils.terr(f.request) f.reply = f.request.reply print s.run("response", f) print s.run("error", f) print s.run("clientconnect", f) print s.run("clientdisconnect", f) print s.run("serverconnect", f) time.sleep(0.1) assert ctx.count == 5
def test_simple(self): f = tutils.tflow_full() resp = f.response assert resp._assemble() assert resp.size() == len(resp._assemble()) resp2 = resp.copy() assert resp2 == resp resp.content = None assert resp._assemble() assert resp.size() == len(resp._assemble()) resp.content = CONTENT_MISSING tutils.raises("Cannot assemble flow with CONTENT_MISSING", resp._assemble)
def test_replace_encoded(self): f = tutils.tflow_full() f.request.content = "afoob" f.request.encode("gzip") f.response.content = "afoob" f.response.encode("gzip") f.replace("foo", "bar") assert f.request.content != "abarb" f.request.decode() assert f.request.content == "abarb" assert f.response.content != "abarb" f.response.decode() assert f.response.content == "abarb"
def test_concurrent2(self): ctx = TScriptContext() s = script.Script( [tutils.test_data.path("scripts/concurrent_decorator.py")], ctx) s.load() f = tutils.tflow_full() f.error = tutils.terr(f.request) f.reply = f.request.reply print s.run("response", f) print s.run("error", f) print s.run("clientconnect", f) print s.run("clientdisconnect", f) print s.run("serverconnect", f) time.sleep(0.1) assert ctx.count == 5
def test_replacehooks(): h = flow.ReplaceHooks() h.add("~q", "foo", "bar") assert h.lst h.set( [ (".*", "one", "two"), (".*", "three", "four"), ] ) assert h.count() == 2 h.clear() assert not h.lst h.add("~q", "foo", "bar") h.add("~s", "foo", "bar") v = h.get_specs() assert v == [('~q', 'foo', 'bar'), ('~s', 'foo', 'bar')] assert h.count() == 2 h.clear() assert h.count() == 0 f = tutils.tflow() f.request.content = "foo" h.add("~s", "foo", "bar") h.run(f) assert f.request.content == "foo" f = tutils.tflow_full() f.request.content = "foo" f.response.content = "foo" h.run(f) assert f.response.content == "bar" assert f.request.content == "foo" f = tutils.tflow() h.clear() h.add("~q", "foo", "bar") f.request.content = "foo" h.run(f) assert f.request.content == "bar" assert not h.add("~", "foo", "bar") assert not h.add("foo", "*", "bar")
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() f = tutils.tflow_full() f.error = tutils.terr(f.request) f.reply = f.request.reply with mock.patch("libmproxy.controller.DummyReply.__call__") as m: s.run("clientconnect", f) s.run("serverconnect", f) s.run("response", f) s.run("error", f) s.run("clientdisconnect", f) time.sleep(0.1) assert m.call_count == 5
def test_getset_state(self): f = tutils.tflow_full() state = f._get_state() assert f._get_state() == protocol.http.HTTPFlow._from_state(state)._get_state() f.response = None f.error = Error("error") state = f._get_state() assert f._get_state() == protocol.http.HTTPFlow._from_state(state)._get_state() f2 = f.copy() assert f._get_state() == f2._get_state() assert not f == f2 f2.error = Error("e2") assert not f == f2 f._load_state(f2._get_state()) assert f._get_state() == f2._get_state()
def test_copy(self): f = tutils.tflow_full() f2 = f.copy() assert not f is f2 assert not f.request is f2.request assert f.request.headers == f2.request.headers assert not f.request.headers is f2.request.headers assert f.response == f2.response assert not f.response is f2.response f = tutils.tflow_err() f2 = f.copy() assert not f is f2 assert not f.request is f2.request assert f.request.headers == f2.request.headers assert not f.request.headers is f2.request.headers assert f.error == f2.error assert not f.error is f2.error
def test_stickycookie(self): s = flow.State() fm = flow.FlowMaster(None, s) assert "Invalid" in fm.set_stickycookie("~h") fm.set_stickycookie(".*") assert fm.stickycookie_state fm.set_stickycookie(None) assert not fm.stickycookie_state fm.set_stickycookie(".*") tf = tutils.tflow_full() tf.response.headers["set-cookie"] = ["foo=bar"] fm.handle_request(tf.request) fm.handle_response(tf.response) assert fm.stickycookie_state.jar assert not "cookie" in tf.request.headers fm.handle_request(tf.request) assert tf.request.headers["cookie"] == ["foo=bar"]
def test_getset_state(self): f = tutils.tflow_full() state = f._get_state() assert f._get_state() == protocol.http.HTTPFlow._from_state( state)._get_state() f.response = None f.error = Error("error") state = f._get_state() assert f._get_state() == protocol.http.HTTPFlow._from_state( state)._get_state() f2 = f.copy() assert f._get_state() == f2._get_state() assert not f == f2 f2.error = Error("e2") assert not f == f2 f._load_state(f2._get_state()) assert f._get_state() == f2._get_state()
def test_add_remove(self): h = flow.ReplaceHooks() h.add("~q", "foo", "bar") assert h.lst h.remove("~q", "foo", "bar") assert not h.lst h.add("~q", "foo", "bar") h.add("~s", "foo", "bar") v = h.get_specs() assert v == [('~q', 'foo', 'bar'), ('~s', 'foo', 'bar')] assert h.count() == 2 h.remove("~q", "foo", "bar") assert h.count() == 1 h.remove("~q", "foo", "bar") assert h.count() == 1 h.clear() assert h.count() == 0 f = tutils.tflow() f.request.content = "foo" h.add("~s", "foo", "bar") h.run(f) assert f.request.content == "foo" f = tutils.tflow_full() f.request.content = "foo" f.response.content = "foo" h.run(f) assert f.response.content == "bar" assert f.request.content == "foo" f = tutils.tflow() h.clear() h.add("~q", "foo", "bar") f.request.content = "foo" h.run(f) assert f.request.content == "bar" assert not h.add("~", "foo", "bar") assert not h.add("foo", "*", "bar")
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() f = tutils.tflow_full() f.error = tutils.terr(f.request) f.reply = f.request.reply with mock.patch("libmproxy.controller.DummyReply.__call__") as m: t_start = time.time() s.run("clientconnect", f) s.run("serverconnect", f) s.run("response", f) s.run("error", f) s.run("clientdisconnect", f) while (time.time() - t_start) < 1 and m.call_count <= 5: if m.call_count == 5: return time.sleep(0.001) assert False
def test_replace_unicode(self): f = tutils.tflow_full() f.response.content = "\xc2foo" f.replace("foo", u"bar")
def test_format_flow(): f = tutils.tflow_full() assert common.format_flow(f, True) assert common.format_flow(f, True, hostheader=True) assert common.format_flow(f, True, extended=True)
def test_backup_idempotence(self): f = tutils.tflow_full() f.backup() f.revert() f.backup() f.revert()