From 7d45d7f15e04ffd5dc1e593d2d3f3c4c9bfd1193 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Wed, 20 Dec 2017 11:53:44 +1300 Subject: cuts: don't crash when retrieving header of a non-existent response --- mitmproxy/addons/cut.py | 2 ++ test/mitmproxy/addons/test_cut.py | 5 +++++ 2 files changed, 7 insertions(+) diff --git a/mitmproxy/addons/cut.py b/mitmproxy/addons/cut.py index b90df549..f4b560e8 100644 --- a/mitmproxy/addons/cut.py +++ b/mitmproxy/addons/cut.py @@ -36,6 +36,8 @@ def extract(cut: str, f: flow.Flow) -> typing.Union[str, bytes]: if spec == "host" and is_addr(current): return str(current[0]) elif spec.startswith("header["): + if not current: + return "" return current.headers.get(headername(spec), "") elif isinstance(part, bytes): return part diff --git a/test/mitmproxy/addons/test_cut.py b/test/mitmproxy/addons/test_cut.py index 0a523fff..71e699db 100644 --- a/test/mitmproxy/addons/test_cut.py +++ b/test/mitmproxy/addons/test_cut.py @@ -135,6 +135,11 @@ def test_cut(): 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()] -- cgit v1.2.3