Implement field renaming on click

README.md

@@ -89,7 +89,7 @@ You may also give some sample raw Protobuf data, that was sent to this endpoint,
 
 Just hover a field to have focus. If the field is an integer type, use the mouse wheel to increment/decrement it. Enum information appears on hover too.
 
-Here it is! You can determine the meaning of every field with that. If you extracted .protos out of minified code, you should be able to rename fields according to what you notice they mean, using your favorite text editor.
+Here it is! You can determine the meaning of every field with that. If you extracted .protos out of minified code, you can rename fields according to what you notice they mean, by clicking their names.
 
 Happy reversing! 👌 🎉
 
@@ -175,7 +175,6 @@ class MyTransport():
 
 The following could be coming for further releases:
 * Finishing the automatic fuzzing part.
-* In-interface support for renaming .proto fields.
 * Support for extracting extensions out of Java code.
 * Support for the JSPB (main JavaScript) runtime.
 * If there's any other platform you wish to see supported, just drop an issue and I'll look at it.

extractors/from_binary.py

@@ -41,7 +41,8 @@ def walk_binary(binr):
 
         if cursor == -1:
             break
-        cursor += 6 + (binr[cursor:cursor + 5] == b'devel') * 5
+        cursor += len('.proto')
+        cursor += (binr[cursor:cursor + 5] == b'devel') * 5
 
         # Search back for the (1, length-delimited) marker
         start = binr.rfind(b'\x0a', max(cursor - 128, 0), cursor)
@@ -51,6 +52,8 @@ def walk_binary(binr):
 
         # Look just after for subsequent markers
         tags = b'\x12\x1a\x22\x2a\x32\x3a\x42\x4a\x50\x58\x62'
+        if binr[cursor] not in tags:
+            continue
 
         while cursor < len(binr) and binr[cursor] in tags:
             tags = tags[tags.index(binr[cursor]):]

gui.py

@@ -78,7 +78,7 @@ def __init__(self):
 
         for tree in (self.fuzzer.pbTree, self.fuzzer.getTree):
             tree.itemEntered.connect(lambda item, _: item.edit() if hasattr(item, 'edit') else None)
-            tree.itemClicked.connect(lambda item, _: item.updateCheck())
+            tree.itemClicked.connect(lambda item, col: item.updateCheck(col=col))
             tree.header().setSectionResizeMode(QHeaderView.ResizeToContents)
 
         self.welcome.mydirLabel.setText(self.welcome.mydirLabel.text() % BASE_PATH)
@@ -306,7 +306,9 @@ def launch_fuzzer(self, item):
             data, sample_id = item.data(Qt.UserRole), 0
 
         if data and assert_installed(self.view, binaries=['protoc']):
-            self.pb_request = load_proto_msgs(BASE_PATH / 'protos' / data['request']['proto_path'])
+            self.current_req_proto = BASE_PATH / 'protos' / data['request']['proto_path']
+
+            self.pb_request = load_proto_msgs(self.current_req_proto)
             self.pb_request = dict(self.pb_request)[data['request']['proto_msg']]()
 
             if data.get('response') and data['response']['format'] == 'raw_pb':

utils/common.py

@@ -1,5 +1,6 @@
 #!/usr/bin/python3
 #-*- encoding: Utf-8 -*-
+from google.protobuf.descriptor_pb2 import FileDescriptorSet
 from google.protobuf.message import Message
 from tempfile import TemporaryDirectory
 from inspect import getmembers, isclass
@@ -122,9 +123,9 @@ def insert_endpoint(base_path, obj):
     with open(path, 'w') as fd:
         dump(json, fd, ensure_ascii=False, indent=4)
 
-# Turn a .proto input into Python classes
+# Turn a .proto input into Python classes.
 
-def load_proto_msgs(proto_path):
+def load_proto_msgs(proto_path, ret_source_info=False):
     # List imports that we need to specify to protoc for the necessary *_pb2.py to be generated
 
     proto_dir = Path(proto_path).parent
@@ -148,10 +149,21 @@ def load_proto_msgs(proto_path):
     # Execute protoc and import the actual module from a tmp
 
     with TemporaryDirectory() as arg_python_out:
-        cmd = run(['protoc', '--proto_path=%s' % arg_proto_path, '--python_out=' + arg_python_out, *arg_proto_files], stderr=PIPE, encoding='utf8')
+        args = ['protoc', '--proto_path=%s' % arg_proto_path, '--python_out=' + arg_python_out, *arg_proto_files]
+        if ret_source_info:
+            args += ['-o%s' % (Path(arg_python_out) / 'desc_info'), '--include_source_info', '--include_imports']
+
+        cmd = run(args, stderr=PIPE, encoding='utf8')
         if cmd.returncode:
             raise ValueError(cmd.stderr)
 
+        if ret_source_info:
+            with open(Path(arg_python_out) / 'desc_info', 'rb') as fd:
+                yield FileDescriptorSet.FromString(fd.read()), arg_proto_path
+                return
+
+        # Do actual import
+
         module_name = str(proto_dir).replace(str(arg_proto_path), '').strip('/\\').replace('/', '.')
         if module_name:
             module_name += '.'

views/fuzzer.py

@@ -1,10 +1,11 @@
 #!/usr/bin/python3
 #-*- encoding: Utf-8 -*-
-from PyQt5.QtWidgets import QTreeWidgetItem, QLineEdit, QCheckBox, QAbstractSpinBox
+from PyQt5.QtWidgets import QTreeWidgetItem, QLineEdit, QCheckBox, QAbstractSpinBox, QInputDialog, QMessageBox
 from PyQt5.QtCore import QUrl, Qt, pyqtSignal, QByteArray, QRegExp
 from PyQt5.QtWebEngineWidgets import QWebEngineView
 from PyQt5.QtGui import QRegExpValidator
 
+from google.protobuf.descriptor_pb2 import FileDescriptorSet
 from xml.dom.minidom import parseString
 from subprocess import run, PIPE
 from functools import partial
@@ -13,12 +14,15 @@
 from zlib import decompress
 from struct import unpack
 from io import BytesIO
+from re import match
 
 # Monkey-patch enum value checking, and then do Protobuf imports
 from google.protobuf.internal import type_checkers
 type_checkers.SupportsOpenEnums = lambda x: True
 from google.protobuf.internal.type_checkers import _VALUE_CHECKERS
 
+from utils.common import load_proto_msgs
+
 """
     We extend QWebEngineView with a method that parses responses and
     displays response data in the relevant way.
@@ -185,6 +189,7 @@ def __init__(self, item, ds, app, path):
         self.repeated = ds.label == ds.LABEL_REPEATED
         self.void = True
         self.dupped = False
+        self.dupe_obj, self.orig_obj = None, None
         self.value = None
         self.selfPb = None
         self.parentPb = None
@@ -194,7 +199,7 @@ def __init__(self, item, ds, app, path):
         self.settingDefault = False
 
         if not ds.full_name in path:
-            super().__init__(item, [ds.name + '+' * self.repeated + '  ', type_txt + '  '])
+            super().__init__(item, [ds.full_name.split('.')[-1] + '+' * self.repeated + '  ', type_txt + '  '])
         else:
             super().__init__(item, ['...' + '  ', ''])
             self.updateCheck = self.createCollapsed
@@ -382,7 +387,7 @@ def update(self, val):
                 self.selfPb = None # So that if we're repeated we're recreated further
                 self.void = True
 
-    def createCollapsed(self, recursive=False):
+    def createCollapsed(self, recursive=False, col=None):
         if not recursive:
             self.path = []
 
@@ -420,12 +425,14 @@ def duplicate(self, settingDefault=False, force=False):
                 if self.isMsg:
                     self.app.parse_desc(self.ds.message_type, newObj, self.path + [self.ds.full_name])
 
+                self.dupe_obj = newObj
+                newObj.orig_obj = self
                 return newObj
 
     # A message is checked -> all parents are checked
     # A message is unchecked -> all children are unchecked
 
-    def updateCheck(self, recursive=False):
+    def updateCheck(self, recursive=False, col=None):
         if not self.required and self.lastCheckState != self.checkState(0):
             self.lastCheckState = self.checkState(0)
 
@@ -456,9 +463,93 @@ def updateCheck(self, recursive=False):
 
                 self.getSelfPb().Clear()
                 self.update(None)
+
+            if not recursive:
+                self.app.update_fuzzer()
 
-        if not recursive:
-            self.app.update_fuzzer()
+        elif col == 0:
+            self.promptRename()
+
+
+    """
+        Wheck field name is clicked, offer to rename the field
+        
+        In order to rewrite field name in the .proto without discarding
+        comments or other information, we'll ask protoc to generate file
+        source information [1], that we'll parse and that will give us
+        text offset for it.
+        
+        [1] https://github.com/google/protobuf/blob/7f3e23/src/google/protobuf/descriptor.proto#L715
+    """
+
+    def promptRename(self):
+        text, good = QInputDialog.getText(self.app.view, ' ', 'Rename this field:', text=self.text(0).strip('+ '))
+        if text:
+            if not match('^\w+$', text):
+                QMessageBox.warning(self.app.view, ' ', 'This is not a valid alphanumeric name.')
+                self.promptRename()
+            else:
+                try:
+                    if self.doRename(text):
+                        return
+                except Exception:
+                    pass
+                QMessageBox.warning(self.app.view, ' ', 'Field was not found in .proto, did you edit it elsewhere?')
+
+    def doRename(self, new_name):
+        file_set, proto_path = next(load_proto_msgs(self.app.current_req_proto, True), None)
+        file_set = file_set.file
+
+        """
+        First, recursively iterate over descriptor fields until we have
+        a numeric path that leads to it in the structure, as explained
+        in [1] above
+        """
+
+        for file_ in file_set:
+            field_path = self.findPathForField(file_.message_type, [4], file_.package)
+
+            if field_path:
+                for location in file_.source_code_info.location:
+                    if location.path == field_path:
+                        start_line, start_col, end_col = location.span[:3]
+
+                        # Once we have file position information, do
+                        # write the new field name in .proto
+
+                        file_path = proto_path / file_.name
+                        with open(file_path) as fd:
+                            lines = fd.readlines()
+                        assert lines[start_line][start_col:end_col] == self.text(0).strip('+ ')
+
+                        lines[start_line] = lines[start_line][:start_col] + new_name + \
+                                            lines[start_line][end_col:]
+                        with open(file_path, 'w') as fd:
+                            fd.writelines(lines)
+
+                        # Update the name on GUI items corresponding to
+                        # this field and its duplicates (if repeated)
+
+                        obj = self
+                        while obj.orig_obj:
+                            obj = obj.orig_obj
+                        while obj:
+                            obj.ds.full_name = obj.ds.full_name.rsplit('.', 1)[0] + '.' + new_name
+                            obj.setText(0, new_name + '+' * self.repeated + '  ')
+                            obj = obj.dupe_obj
+                        return True
+
+    def findPathForField(self, msgs, path, cur_name):
+        if cur_name:
+            cur_name += '.'
+
+        for i, msg in enumerate(msgs):
+            if self.ds.full_name.startswith(cur_name + msg.name + '.'):
+                for j, field in enumerate(msg.field):
+                    if cur_name + msg.name + '.' + field.name == self.ds.full_name:
+                        return path + [i, 2, j, 1]
+
+                return self.findPathForField(msg.nested_type, path + [i, 3], cur_name + msg.name)
 
     def _edit(self, ev=None):
         if not self.widget.hasFocus():