import Queue from cStringIO import StringIO from libmproxy import console, proxy, filt, flow, controller import tutils import libpry class uStickyCookieState(libpry.AutoTree): 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 test_handle_response(self): c = "SSID=mooo, FOO=bar; Domain=.google.com; Path=/; "\ "Expires=Wed, 13-Jan-2021 22:23:01 GMT; Secure; " s, f = self._response(c, "host") assert not s.jar.keys() s, f = self._response(c, "www.google.com") assert s.jar.keys() s, f = self._response("SSID=mooo", "www.google.com") assert s.jar.keys()[0] == ('www.google.com', 80, '/') def test_handle_request(self): s, f = self._response("SSID=mooo", "www.google.com") assert "cookie" not in f.request.headers s.handle_request(f) assert "cookie" in f.request.headers class uStickyAuthState(libpry.AutoTree): 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"] class uClientPlaybackState(libpry.AutoTree): 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) fm.stop_client_playback() assert not fm.client_playback class uServerPlaybackState(libpry.AutoTree): def test_hash(self): s = flow.ServerPlaybackState(None, [], False) r = tutils.tflow() r2 = tutils.tflow() assert s._hash(r) assert s._hash(r) == s._hash(r2) r.request.headers["foo"] = ["bar"] assert s._hash(r) == s._hash(r2) r.request.path = "voing" assert s._hash(r) != s._hash(r2) def test_headers(self): s = flow.ServerPlaybackState(["foo"], [], 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_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) class uFlow(libpry.AutoTree): def test_run_script(self): f = tutils.tflow() f.response = tutils.tresp() f.request = f.response.request se = f.run_script("scripts/a") assert "DEBUG" == se.strip() assert f.request.host == "TESTOK" def test_match(self): f = tutils.tflow() f.response = tutils.tresp() f.request = f.response.request assert not f.match(filt.parse("~b test")) assert f.match(None) def test_backup(self): f = tutils.tflow() f.response = tutils.tresp() f.request = f.response.request f.request.content = "foo" assert not f.modified() f.backup() f.request.content = "bar" assert f.modified() f.revert() assert f.request.content == "foo" def test_getset_state(self): f = tutils.tflow() f.response = tutils.tresp(f.request) state = f.get_state() assert f.get_state() == flow.Flow.from_state(state).get_state() f.response = None f.error = proxy.Error(f.request, "error") state = f.get_state() assert f.get_state() == flow.Flow.from_state(state).get_state() f2 = tutils.tflow() f2.error = proxy.Error(f.request, "e2") assert not f == f2 f.load_state(f2.get_state()) assert f.get_state() == f2.get_state() def test_kill(self): s = flow.State() fm = flow.FlowMaster(None, s) f = tutils.tflow() f.request = tutils.treq() f.intercept() assert not f.request.acked f.kill(fm) assert f.request.acked f.intercept() f.response = tutils.tresp() f.request = f.response.request f.request.ack() assert not f.response.acked f.kill(fm) assert f.response.acked def test_killall(self): s = flow.State() fm = flow.FlowMaster(None, s) r = tutils.treq() fm.handle_request(r) r = tutils.treq() fm.handle_request(r) for i in s.view: assert not i.request.acked s.killall(fm) for i in s.view: assert i.request.acked def test_accept_intercept(self): f = tutils.tflow() f.request = tutils.treq() f.intercept() assert not f.request.acked f.accept_intercept() assert f.request.acked f.response = tutils.tresp() f.request = f.response.request f.intercept() f.request.ack() assert not f.response.acked f.accept_intercept() assert f.response.acked def test_serialization(self): f = flow.Flow(None) f.request = tutils.treq() 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" f = tutils.tflow_err() f.replace("error", "bar") assert f.error.msg == "bar" class uState(libpry.AutoTree): def test_backup(self): bc = proxy.ClientConnect(("address", 22)) c = flow.State() req = tutils.treq() f = c.add_request(req) f.backup() c.revert(f) def test_flow(self): """ normal flow: connect -> request -> response """ bc = proxy.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 dc = proxy.ClientDisconnect(bc) def test_err(self): bc = proxy.ClientConnect(("address", 22)) c = flow.State() req = tutils.treq() f = c.add_request(req) e = proxy.Error(f.request, "message") assert c.add_error(e) e = proxy.Error(tutils.tflow().request, "message") assert not c.add_error(e) def test_set_limit(self): c = flow.State() req = tutils.treq() assert len(c.view) == 0 f = 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_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 _add_request(self, state): req = tutils.treq() f = state.add_request(req) return f def _add_response(self, state): req = tutils.treq() f = state.add_request(req) resp = tutils.tresp(req) state.add_response(resp) def _add_error(self, state): req = tutils.treq() f = state.add_request(req) f.error = proxy.Error(f.request, "msg") def test_clear(self): c = flow.State() f = self._add_request(c) f.intercepting = True c.clear() assert c.flow_count() == 0 def test_dump_flows(self): c = flow.State() self._add_request(c) self._add_response(c) self._add_request(c) self._add_response(c) self._add_request(c) self._add_response(c) self._add_error(c) flows = c.view[:] c.clear() c.load_flows(flows) assert isinstance(c._flow_list[0], flow.Flow) def test_accept_all(self): c = flow.State() self._add_request(c) self._add_response(c) self._add_request(c) c.accept_all() class uSerialize(libpry.AutoTree): 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_roundtrip(self): sio = StringIO() f = tutils.tflow() f.request.content = "".join(chr(i) for i in range(255)) w = flow.FlowWriter(sio) w.add(f) sio.seek(0) r = flow.FlowReader(sio) l = list(r.stream()) assert len(l) == 1 f2 = l[0] assert f2.get_state() == f.get_state() assert f2.request.assemble() == f.request.assemble() def test_load_flows(self): r = self._treader() s = flow.State() fm = flow.FlowMaster(None, s) fm.load_flows(r) assert len(s._flow_list) == 6 def test_error(self): sio = StringIO() sio.write("bogus") sio.seek(0) r = flow.FlowReader(sio) libpry.raises(flow.FlowReadError, list, r.stream()) f = flow.FlowReadError("foo") assert f.strerror == "foo" 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) 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 = proxy.ClientDisconnect(req.client_conn) fm.handle_clientdisconnect(dc) err = proxy.Error(f.request, "msg") fm.handle_error(err) 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(proxy.Error(f.request, "error")) 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) assert fm.do_server_playback(tutils.tflow()) fm.start_server_playback(pb, False, [], True) r = tutils.tflow() r.request.content = "gibble" assert not fm.do_server_playback(r) assert fm.do_server_playback(tutils.tflow()) q = Queue.Queue() fm.tick(q) assert controller.should_exit fm.stop_server_playback() assert not fm.server_playback 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) f = 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_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"] tests = [ uStickyCookieState(), uStickyAuthState(), uServerPlaybackState(), uClientPlaybackState(), uFlow(), uState(), uSerialize(), uFlowMaster(), ]