From 36352c95397351f4e69d857fcd8110fe125388f7 Mon Sep 17 00:00:00 2001 From: Thomas Kriechbaumer Date: Sun, 19 Feb 2017 12:37:50 +0100 Subject: [PATCH] protobuf: coverage++ --- examples/simple/custom_contentview.py | 2 +- mitmproxy/contentviews/__init__.py | 5 +- mitmproxy/contentviews/protobuf.py | 21 ++++---- setup.cfg | 2 - setup.py | 3 -- test/mitmproxy/contentviews/test_protobuf.py | 52 +++++++++++++++++--- 6 files changed, 57 insertions(+), 28 deletions(-) diff --git a/examples/simple/custom_contentview.py b/examples/simple/custom_contentview.py index 352163970..1f3a38ecb 100644 --- a/examples/simple/custom_contentview.py +++ b/examples/simple/custom_contentview.py @@ -10,7 +10,7 @@ class ViewSwapCase(contentviews.View): # We don't have a good solution for the keyboard shortcut yet - # you manually need to find a free letter. Contributions welcome :) - prompt = ("swap case text", "p") + prompt = ("swap case text", "z") content_types = ["text/plain"] def __call__(self, data: bytes, **metadata): diff --git a/mitmproxy/contentviews/__init__.py b/mitmproxy/contentviews/__init__.py index 357172e3e..c7db66904 100644 --- a/mitmproxy/contentviews/__init__.py +++ b/mitmproxy/contentviews/__init__.py @@ -159,6 +159,7 @@ def get_content_view(viewmode: View, data: bytes, **metadata): return desc, safe_to_print(content), error +# The order in which ContentViews are added is important! add(auto.ViewAuto()) add(raw.ViewRaw()) add(hex.ViewHex()) @@ -172,9 +173,7 @@ add(urlencoded.ViewURLEncoded()) add(multipart.ViewMultipart()) add(image.ViewImage()) add(query.ViewQuery()) - -if protobuf.ViewProtobuf.is_available(): - add(protobuf.ViewProtobuf()) +add(protobuf.ViewProtobuf()) __all__ = [ "View", "VIEW_CUTOFF", "KEY_MAX", "format_text", "format_dict", diff --git a/mitmproxy/contentviews/protobuf.py b/mitmproxy/contentviews/protobuf.py index 620d9444d..4bbb15809 100644 --- a/mitmproxy/contentviews/protobuf.py +++ b/mitmproxy/contentviews/protobuf.py @@ -15,31 +15,28 @@ class ViewProtobuf(base.View): "application/x-protobuffer", ] - @staticmethod - def is_available(): + def is_available(self): try: p = subprocess.Popen( ["protoc", "--version"], stdout=subprocess.PIPE ) out, _ = p.communicate() - return out.startswith("libprotoc") + return out.startswith(b"libprotoc") except: return False - def decode_protobuf(self, content): + def __call__(self, data, **metadata): + if not self.is_available(): + raise NotImplementedError("protoc not found. Please make sure 'protoc' is available in $PATH.") + # if Popen raises OSError, it will be caught in # get_content_view and fall back to Raw p = subprocess.Popen(['protoc', '--decode_raw'], stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE) - out, err = p.communicate(input=content) - if out: - return out - else: - return err - - def __call__(self, data, **metadata): - decoded = self.decode_protobuf(data) + decoded, _ = p.communicate(input=data) + if not decoded: + raise ValueError("Failed to parse input.") return "Protobuf", base.format_text(decoded) diff --git a/setup.cfg b/setup.cfg index 4f2e86b33..52f51f260 100644 --- a/setup.cfg +++ b/setup.cfg @@ -22,7 +22,6 @@ exclude_lines = [tool:full_coverage] exclude = mitmproxy/contentviews/__init__.py - mitmproxy/contentviews/protobuf.py mitmproxy/contentviews/wbxml.py mitmproxy/contentviews/xml_html.py mitmproxy/net/tcp.py @@ -58,7 +57,6 @@ exclude = mitmproxy/certs.py mitmproxy/connections.py mitmproxy/contentviews/base.py - mitmproxy/contentviews/protobuf.py mitmproxy/contentviews/wbxml.py mitmproxy/contentviews/xml_html.py mitmproxy/controller.py diff --git a/setup.py b/setup.py index 1b3f08a4a..a942a9e30 100644 --- a/setup.py +++ b/setup.py @@ -110,9 +110,6 @@ setup( "sphinx_rtd_theme>=0.1.9, <0.2", ], 'contentviews': [ - "protobuf>=3.1.0, <3.3", - # TODO: Find Python 3 replacement - # "pyamf>=0.8.0, <0.9", ], 'examples': [ "beautifulsoup4>=4.4.1, <4.6", diff --git a/test/mitmproxy/contentviews/test_protobuf.py b/test/mitmproxy/contentviews/test_protobuf.py index 1224b8db2..31e382ecb 100644 --- a/test/mitmproxy/contentviews/test_protobuf.py +++ b/test/mitmproxy/contentviews/test_protobuf.py @@ -1,12 +1,50 @@ +from unittest import mock +import pytest + from mitmproxy.contentviews import protobuf from mitmproxy.test import tutils from . import full_eval -if protobuf.ViewProtobuf.is_available(): - def test_view_protobuf_request(): - v = full_eval(protobuf.ViewProtobuf()) - p = tutils.test_data.path("mitmproxy/data/protobuf01") - content_type, output = v(open(p, "rb").read()) - assert content_type == "Protobuf" - assert output.next()[0][1] == '1: "3bbc333c-e61c-433b-819a-0b9a8cc103b8"' +def test_view_protobuf_request(): + v = full_eval(protobuf.ViewProtobuf()) + p = tutils.test_data.path("mitmproxy/data/protobuf01") + + with mock.patch('mitmproxy.contentviews.protobuf.ViewProtobuf.is_available'): + with mock.patch('subprocess.Popen') as n: + m = mock.Mock() + attrs = {'communicate.return_value': (b'1: "3bbc333c-e61c-433b-819a-0b9a8cc103b8"', True)} + m.configure_mock(**attrs) + n.return_value = m + + content_type, output = v(open(p, "rb").read()) + assert content_type == "Protobuf" + assert output[0] == [('text', b'1: "3bbc333c-e61c-433b-819a-0b9a8cc103b8"')] + + m.communicate = mock.MagicMock() + m.communicate.return_value = (None, None) + with pytest.raises(ValueError, matches="Failed to parse input."): + v(b'foobar') + + +def test_view_protobuf_availability(): + with mock.patch('subprocess.Popen') as n: + m = mock.Mock() + attrs = {'communicate.return_value': (b'libprotoc fake version', True)} + m.configure_mock(**attrs) + n.return_value = m + assert protobuf.ViewProtobuf().is_available() + + m = mock.Mock() + attrs = {'communicate.return_value': (b'command not found', True)} + m.configure_mock(**attrs) + n.return_value = m + assert not protobuf.ViewProtobuf().is_available() + + +def test_view_protobuf_fallback(): + with mock.patch('subprocess.Popen.communicate') as m: + m.side_effect = OSError() + v = full_eval(protobuf.ViewProtobuf()) + with pytest.raises(NotImplementedError, matches='protoc not found'): + v(b'foobar')