Kaitai parser for protobuf

1 files changed, 60 insertions, 22 deletions

diff --git a/mitmproxy/contentviews/protobuf.py b/mitmproxy/contentviews/protobuf.py
index 4bbb1580..abd3985a 100644
--- a/mitmproxy/contentviews/protobuf.py
+++ b/mitmproxy/contentviews/protobuf.py
@@ -1,6 +1,63 @@
-import subprocess
+import io
 
+from kaitaistruct import KaitaiStream
 from . import base
+from mitmproxy.contrib.kaitaistruct import google_protobuf
+
+
+def write_buf(out, field_tag, body, indent_level):
+    if body is not None:
+        out.write("{: <{level}}{}: {}\n".format('', field_tag, body if isinstance(body, int) else str(body, 'utf-8'),
+                                                level=indent_level))
+    elif field_tag is not None:
+        out.write(' ' * indent_level + str(field_tag) + " {\n")
+    else:
+        out.write(' ' * indent_level + "}\n")
+
+
+def format_pbuf(raw):
+    out = io.StringIO()
+    stack = []
+
+    try:
+        buf = google_protobuf.GoogleProtobuf(KaitaiStream(io.BytesIO(raw)))
+    except:
+        return False
+    stack.extend([(pair, 0) for pair in buf.pairs[::-1]])
+
+    while len(stack):
+        pair, indent_level = stack.pop()
+
+        if pair.wire_type == pair.WireTypes.group_start:
+            body = None
+        elif pair.wire_type == pair.WireTypes.group_end:
+            body = None
+            pair._m_field_tag = None
+        elif pair.wire_type == pair.WireTypes.len_delimited:
+            body = pair.value.body
+        elif pair.wire_type == pair.WireTypes.varint:
+            body = pair.value.value
+        else:
+            body = pair.value
+
+        try:
+            next_buf = google_protobuf.GoogleProtobuf(KaitaiStream(io.BytesIO(body)))
+            stack.extend([(pair, indent_level + 2) for pair in next_buf.pairs[::-1]])
+            write_buf(out, pair.field_tag, None, indent_level)
+        except:
+            write_buf(out, pair.field_tag, body, indent_level)
+
+        if stack:
+            prev_level = stack[-1][1]
+        else:
+            prev_level = 0
+
+        if prev_level < indent_level:
+            levels = int((indent_level - prev_level) / 2)
+            for i in range(1, levels + 1):
+                write_buf(out, None, None, indent_level - i * 2)
+
+    return out.getvalue()
 
 
 class ViewProtobuf(base.View):
@@ -15,28 +72,9 @@ class ViewProtobuf(base.View):
         "application/x-protobuffer",
     ]
 
-    def is_available(self):
-        try:
-            p = subprocess.Popen(
-                ["protoc", "--version"],
-                stdout=subprocess.PIPE
-            )
-            out, _ = p.communicate()
-            return out.startswith(b"libprotoc")
-        except:
-            return False
-
     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)
-        decoded, _ = p.communicate(input=data)
+        decoded = format_pbuf(data)
         if not decoded:
             raise ValueError("Failed to parse input.")
+
         return "Protobuf", base.format_text(decoded)