Kaitai parser for protobuf

6 files changed, 92 insertions, 43 deletions

diff --git a/test/mitmproxy/contentviews/test_protobuf.py b/test/mitmproxy/contentviews/test_protobuf.py
index 71e51576..6c6e37f2 100644
--- a/test/mitmproxy/contentviews/test_protobuf.py
+++ b/test/mitmproxy/contentviews/test_protobuf.py
@@ -1,52 +1,31 @@
-from unittest import mock
 import pytest
 
 from mitmproxy.contentviews import protobuf
 from mitmproxy.test import tutils
 from . import full_eval
 
+data = tutils.test_data.push("mitmproxy/contentviews/test_protobuf_data/")
+
 
 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
-
-            with open(p, "rb") as f:
-                data = f.read()
-            content_type, output = v(data)
-            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')
+    p = data.path("protobuf01")
+
+    with open(p, "rb") as f:
+        raw = f.read()
+    content_type, output = v(raw)
+    assert content_type == "Protobuf"
+    assert output == [[('text', '1: 3bbc333c-e61c-433b-819a-0b9a8cc103b8')]]
+    with pytest.raises(ValueError, matches="Failed to parse input."):
+        v(b'foobar')
+
+
+@pytest.mark.parametrize("filename", ["protobuf02", "protobuf03"])
+def test_format_pbuf(filename):
+    path = data.path(filename)
+    with open(path, "rb") as f:
+        input = f.read()
+    with open(path + "-decoded") as f:
+        expected = f.read()
+
+    assert protobuf.format_pbuf(input) == expected
diff --git a/test/mitmproxy/data/protobuf01 b/test/mitmproxy/contentviews/test_protobuf_data/protobuf01
index fbfdbff3..fbfdbff3 100644
--- a/test/mitmproxy/data/protobuf01
+++ b/test/mitmproxy/contentviews/test_protobuf_data/protobuf01
diff --git a/test/mitmproxy/contentviews/test_protobuf_data/protobuf02 b/test/mitmproxy/contentviews/test_protobuf_data/protobuf02
new file mode 100644
index 00000000..a47c45d5
--- /dev/null
+++ b/test/mitmproxy/contentviews/test_protobuf_data/protobuf02
diff --git a/test/mitmproxy/contentviews/test_protobuf_data/protobuf02-decoded b/test/mitmproxy/contentviews/test_protobuf_data/protobuf02-decoded
new file mode 100644
index 00000000..9be61e28
--- /dev/null
+++ b/test/mitmproxy/contentviews/test_protobuf_data/protobuf02-decoded
@@ -0,0 +1,65 @@
+1 {
+  1: tpbuf
+  4 {
+    1: Person
+    2 {
+      1: name
+      3: 1
+      4: 2
+      5: 9
+    }
+    2 {
+      1: id
+      3: 2
+      4: 2
+      5: 5
+    }
+    2 {
+      1 {
+        12: 1818845549
+      }
+      3: 3
+      4: 1
+      5: 9
+    }
+    2 {
+      1: phone
+      3: 4
+      4: 3
+      5: 11
+      6: .Person.PhoneNumber
+    }
+    3 {
+      1: PhoneNumber
+      2 {
+        1: number
+        3: 1
+        4: 2
+        5: 9
+      }
+      2 {
+        1: type
+        3: 2
+        4: 1
+        5: 14
+        6: .Person.PhoneType
+        7: HOME
+      }
+    }
+    4 {
+      1: PhoneType
+      2 {
+        1: MOBILE
+        2: 0
+      }
+      2 {
+        1: HOME
+        2: 1
+      }
+      2 {
+        1: WORK
+        2: 2
+      }
+    }
+  }
+}
diff --git a/test/mitmproxy/contentviews/test_protobuf_data/protobuf03 b/test/mitmproxy/contentviews/test_protobuf_data/protobuf03
new file mode 100644
index 00000000..9fb230b3
--- /dev/null
+++ b/test/mitmproxy/contentviews/test_protobuf_data/protobuf03
@@ -0,0 +1 @@
+€ ð
\ No newline at end of file
diff --git a/test/mitmproxy/contentviews/test_protobuf_data/protobuf03-decoded b/test/mitmproxy/contentviews/test_protobuf_data/protobuf03-decoded
new file mode 100644
index 00000000..3d3392e1
--- /dev/null
+++ b/test/mitmproxy/contentviews/test_protobuf_data/protobuf03-decoded
@@ -0,0 +1,4 @@
+2 {
+3: 3840
+4: 2160
+}