Python taddons.context函数代码示例

def test_cut():
    c = cut.Cut()
    with taddons.context():
        tflows = [tflow.tflow(resp=True)]
        assert c.cut(tflows, ["request.method"]) == [["GET"]]
        assert c.cut(tflows, ["request.scheme"]) == [["http"]]
        assert c.cut(tflows, ["request.host"]) == [["address"]]
        assert c.cut(tflows, ["request.port"]) == [["22"]]
        assert c.cut(tflows, ["request.path"]) == [["/path"]]
        assert c.cut(tflows, ["request.url"]) == [["http://address:22/path"]]
        assert c.cut(tflows, ["request.content"]) == [[b"content"]]
        assert c.cut(tflows, ["request.header[header]"]) == [["qvalue"]]
        assert c.cut(tflows, ["request.header[unknown]"]) == [[""]]

        assert c.cut(tflows, ["response.status_code"]) == [["200"]]
        assert c.cut(tflows, ["response.reason"]) == [["OK"]]
        assert c.cut(tflows, ["response.content"]) == [[b"message"]]
        assert c.cut(tflows, ["response.header[header-response]"]) == [["svalue"]]
        assert c.cut(tflows, ["moo"]) == [[""]]
        with pytest.raises(exceptions.CommandError):
            assert c.cut(tflows, ["__dict__"]) == [[""]]

    with taddons.context():
        tflows = [tflow.tflow(resp=False)]
        assert c.cut(tflows, ["response.reason"]) == [[""]]
        assert c.cut(tflows, ["response.header[key]"]) == [[""]]

    c = cut.Cut()
    with taddons.context():
        tflows = [tflow.ttcpflow()]
        assert c.cut(tflows, ["request.method"]) == [[""]]
        assert c.cut(tflows, ["response.status"]) == [[""]]

def test_cut():
    v = view.View()
    c = cut.Cut()
    with taddons.context() as tctx:
        v.add([tflow.tflow(resp=True)])
        tctx.master.addons.add(v, c)
        assert c.cut("q.method|@all") == [["GET"]]
        assert c.cut("q.scheme|@all") == [["http"]]
        assert c.cut("q.host|@all") == [["address"]]
        assert c.cut("q.port|@all") == [["22"]]
        assert c.cut("q.path|@all") == [["/path"]]
        assert c.cut("q.url|@all") == [["http://address:22/path"]]
        assert c.cut("q.content|@all") == [[b"content"]]
        assert c.cut("q.header[header]|@all") == [["qvalue"]]
        assert c.cut("q.header[unknown]|@all") == [[""]]

        assert c.cut("s.status_code|@all") == [["200"]]
        assert c.cut("s.reason|@all") == [["OK"]]
        assert c.cut("s.content|@all") == [[b"message"]]
        assert c.cut("s.header[header-response]|@all") == [["svalue"]]
        assert c.cut("moo") == [[""]]
        with pytest.raises(exceptions.CommandError):
            assert c.cut("__dict__") == [[""]]

    v = view.View()
    c = cut.Cut()
    with taddons.context() as tctx:
        tctx.master.addons.add(v, c)
        v.add([tflow.ttcpflow()])
        assert c.cut("q.method|@all") == [[""]]
        assert c.cut("s.status|@all") == [[""]]

def test_decorator():
    with taddons.context() as tctx:
        c = command.CommandManager(tctx.master)
        a = TDec()
        c.collect_commands(a)
        assert "cmd1" in c.commands
        assert c.execute("cmd1 bar") == "ret bar"
        assert "empty" in c.commands
        assert c.execute("empty") is None

    with taddons.context() as tctx:
        tctx.master.addons.add(a)
        assert tctx.master.commands.execute("cmd1 bar") == "ret bar"

def test_movement():
    v = view.View()
    with taddons.context():
        v.go(0)
        v.add([
            tflow.tflow(),
            tflow.tflow(),
            tflow.tflow(),
            tflow.tflow(),
            tflow.tflow(),
        ])
        assert v.focus.index == 0
        v.go(-1)
        assert v.focus.index == 4
        v.go(0)
        assert v.focus.index == 0
        v.go(1)
        assert v.focus.index == 1
        v.go(999)
        assert v.focus.index == 4
        v.go(-999)
        assert v.focus.index == 0

        v.focus_next()
        assert v.focus.index == 1
        v.focus_prev()
        assert v.focus.index == 0

def test_ignore_content():
    s = serverplayback.ServerPlayback()
    with taddons.context(s) as tctx:
        tctx.configure(s, server_replay_ignore_content=False)

        r = tflow.tflow(resp=True)
        r2 = tflow.tflow(resp=True)

        r.request.content = b"foo"
        r2.request.content = b"foo"
        assert s._hash(r) == s._hash(r2)
        r2.request.content = b"bar"
        assert not s._hash(r) == s._hash(r2)

        tctx.configure(s, server_replay_ignore_content=True)
        r = tflow.tflow(resp=True)
        r2 = tflow.tflow(resp=True)
        r.request.content = b"foo"
        r2.request.content = b"foo"
        assert s._hash(r) == s._hash(r2)
        r2.request.content = b"bar"
        assert s._hash(r) == s._hash(r2)
        r2.request.content = b""
        assert s._hash(r) == s._hash(r2)
        r2.request.content = None
        assert s._hash(r) == s._hash(r2)

def test_load():
    s = serverplayback.ServerPlayback()
    with taddons.context(s) as tctx:
        tctx.configure(s)

        r = tflow.tflow(resp=True)
        r.request.headers["key"] = "one"

        r2 = tflow.tflow(resp=True)
        r2.request.headers["key"] = "two"

        s.load_flows([r, r2])

        assert s.count() == 2

        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 not s.flowmap
        assert s.count() == 0

        assert not s.next_flow(r)

def test_encoding():
    sa = core.Core()
    with taddons.context():
        f = tflow.tflow()
        assert sa.encode_options()
        sa.encode([f], "request", "deflate")
        assert f.request.headers["content-encoding"] == "deflate"

        sa.encode([f], "request", "br")
        assert f.request.headers["content-encoding"] == "deflate"

        sa.decode([f], "request")
        assert "content-encoding" not in f.request.headers

        sa.encode([f], "request", "br")
        assert f.request.headers["content-encoding"] == "br"

        sa.encode_toggle([f], "request")
        assert "content-encoding" not in f.request.headers
        sa.encode_toggle([f], "request")
        assert f.request.headers["content-encoding"] == "deflate"
        sa.encode_toggle([f], "request")
        assert "content-encoding" not in f.request.headers

        with pytest.raises(exceptions.CommandError):
            sa.encode([f], "request", "invalid")

 def test_normal(self, load_flow, tmpdir, data):
     rf = readfile.ReadFileStdin()
     with taddons.context():
         tfile = tmpdir.join("tfile")
         tfile.write(data.getvalue())
         rf.load_flows_from_path(str(tfile))
         assert load_flow.called

    def test_authenticate(self):
        up = proxyauth.ProxyAuth()
        with taddons.context(up, loadcore=False) as ctx:
            ctx.configure(up, proxyauth="any", mode="regular")

            f = tflow.tflow()
            assert not f.response
            up.authenticate(f)
            assert f.response.status_code == 407

            f = tflow.tflow()
            f.request.headers["Proxy-Authorization"] = proxyauth.mkauth(
                "test", "test"
            )
            up.authenticate(f)
            assert not f.response
            assert not f.request.headers.get("Proxy-Authorization")

            f = tflow.tflow()
            ctx.configure(up, mode="reverse")
            assert not f.response
            up.authenticate(f)
            assert f.response.status_code == 401

            f = tflow.tflow()
            f.request.headers["Authorization"] = proxyauth.mkauth(
                "test", "test"
            )
            up.authenticate(f)
            assert not f.response
            assert not f.request.headers.get("Authorization")

def test_bind():
    with taddons.context() as tctx:
        km = keymap.Keymap(tctx.master)
        km.executor = mock.Mock()

        with pytest.raises(ValueError):
            km.add("foo", "bar", ["unsupported"])

        km.add("key", "str", ["options", "commands"])
        assert km.get("options", "key")
        assert km.get("commands", "key")
        assert not km.get("flowlist", "key")
        assert len((km.list("commands"))) == 1

        km.handle("unknown", "unknown")
        assert not km.executor.called

        km.handle("options", "key")
        assert km.executor.called

        km.add("glob", "str", ["global"])
        km.executor = mock.Mock()
        km.handle("options", "glob")
        assert km.executor.called

        assert len((km.list("global"))) == 1

def test_flow_set():
    sa = core.Core()
    with taddons.context(loadcore=False):
        f = tflow.tflow(resp=True)
        assert sa.flow_set_options()

        assert f.request.method != "post"
        sa.flow_set([f], "method", "post")
        assert f.request.method == "POST"

        assert f.request.host != "testhost"
        sa.flow_set([f], "host", "testhost")
        assert f.request.host == "testhost"

        assert f.request.path != "/test/path"
        sa.flow_set([f], "path", "/test/path")
        assert f.request.path == "/test/path"

        assert f.request.url != "http://foo.com/bar"
        sa.flow_set([f], "url", "http://foo.com/bar")
        assert f.request.url == "http://foo.com/bar"
        with pytest.raises(exceptions.CommandError):
            sa.flow_set([f], "url", "oink")

        assert f.response.status_code != 404
        sa.flow_set([f], "status_code", "404")
        assert f.response.status_code == 404
        assert f.response.reason == "Not Found"
        with pytest.raises(exceptions.CommandError):
            sa.flow_set([f], "status_code", "oink")

        assert f.response.reason != "foo"
        sa.flow_set([f], "reason", "foo")
        assert f.response.reason == "foo"

def test_options(tmpdir):
    p = str(tmpdir.join("path"))
    sa = core.Core()
    with taddons.context() as tctx:
        tctx.options.listen_host = "foo"
        assert tctx.options.listen_host == "foo"
        sa.options_reset_one("listen_host")
        assert tctx.options.listen_host != "foo"

        with pytest.raises(exceptions.CommandError):
            sa.options_reset_one("unknown")

        tctx.options.listen_host = "foo"
        sa.options_save(p)
        with pytest.raises(exceptions.CommandError):
            sa.options_save("/")

        sa.options_reset()
        assert tctx.options.listen_host == ""
        sa.options_load(p)
        assert tctx.options.listen_host == "foo"

        sa.options_load("/nonexistent")

        with open(p, 'a') as f:
            f.write("'''")
        with pytest.raises(exceptions.CommandError):
            sa.options_load(p)

    def test_playback(self):
        cp = clientplayback.ClientPlayback()
        with taddons.context() as tctx:
            assert cp.count() == 0
            f = tflow.tflow(resp=True)
            cp.start_replay([f])
            assert cp.count() == 1
            RP = "mitmproxy.proxy.protocol.http_replay.RequestReplayThread"
            with mock.patch(RP) as rp:
                assert not cp.current_thread
                cp.tick()
                assert rp.called
                assert cp.current_thread

            cp.flows = []
            cp.current_thread.is_alive.return_value = False
            assert cp.count() == 1
            cp.tick()
            assert cp.count() == 0
            assert tctx.master.has_event("update")
            assert tctx.master.has_event("processing_complete")

            cp.current_thread = MockThread()
            cp.tick()
            assert cp.current_thread is None

            cp.start_replay([f])
            cp.stop_replay()
            assert not cp.flows

 def fill_history(self, commands):
     with taddons.context() as tctx:
         history = commander.CommandHistory(tctx.master, size=3)
         for c in commands:
             cbuf = commander.CommandBuffer(tctx.master, c)
             history.add_command(cbuf)
     return history, tctx.master

def test_simple():
    r = intercept.Intercept()
    with taddons.context(options=Options()) as tctx:
        assert not r.filt
        tctx.configure(r, intercept="~q")
        assert r.filt
        with pytest.raises(exceptions.OptionsError):
            tctx.configure(r, intercept="~~")
        tctx.configure(r, intercept=None)
        assert not r.filt

        tctx.configure(r, intercept="~s")

        f = tflow.tflow(resp=True)
        tctx.cycle(r, f)
        assert f.intercepted

        f = tflow.tflow(resp=False)
        tctx.cycle(r, f)
        assert not f.intercepted

        f = tflow.tflow(resp=True)
        f.reply._state = "handled"
        r.response(f)
        assert f.intercepted