Repository: trailofbits/protofuzz

Issue:
Allow import of generated `_pb2.py`
This may be a misunderstanding of how protobufs work.

If I have a compiled protobuf module, I can try to import the python file as a module

error as expected:
```
message_fuzzers = protofuzz.from_file("/Users/Spellchaser/PycharmProjects/test/module_pb2.py")
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/Spellchaser/PycharmProjects/test/venv/lib/python3.7/site-packages/protofuzz/protofuzz.py", line 251, in from_file
    module = pbimport.from_file(protobuf_file)
  File "/Users/Spellchaser/PycharmProjects/test/venv/lib/python3.7/site-packages/protofuzz/pbimport.py", line 114, in from_file
    raise BadProtobuf()
protofuzz.pbimport.BadProtobuf
```

Discussion:
woodruffw: One small suggestion, otherwise LGTM!
Spellchaser: got some tests to write first

On Wed, Aug 18, 2021, 3:18 PM William Woodruff ***@***.***>
wrote:

> One small suggestion, otherwise LGTM!
>
> —
> You are receiving this because you authored the thread.
> Reply to this email directly, view it on GitHub
> <https://github.com/trailofbits/protofuzz/pull/63#issuecomment-901366732>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AASX7C7LOHTBHIA5HUYUB53T5QBP7ANCNFSM5CMRWQPA>
> .
>

woodruffw: JFYI: I'm currently working on generally fixing the tests in #59. So don't worry about unbreaking any of the old tests; I'll take care of those.

Patch:
diff --git a/protofuzz/pbimport.py b/protofuzz/pbimport.py
index 9de657c..9961af8 100644
--- a/protofuzz/pbimport.py
+++ b/protofuzz/pbimport.py
@@ -92,11 +92,16 @@ def _compile_proto(full_path, dest):
 
 
 def from_file(proto_file):
-    """Take a filename |protoc_file|, compile it via the Protobuf compiler, and import the module.
+    """Takes either a |protoc_file| or a generated |module_file|
+    If given a `_pb2.py` file, this will try to just import the module. This should be the output of the Protobuf compiler; users should not attempt to import arbitrary Python files.
+    If given a `.proto` file, this will compile it via the Protobuf compiler, and import the module.
 
     Return the module if successfully compiled, otherwise raise either a ProtocNotFound or BadProtobuf exception.
 
     """
+    if proto_file.endswith('_pb2.py'):
+        return _load_module(proto_file)
+
     if not proto_file.endswith('.proto'):
         raise BadProtobuf()
 
diff --git a/protofuzz/protofuzz.py b/protofuzz/protofuzz.py
index 4c8af74..c666945 100644
--- a/protofuzz/protofuzz.py
+++ b/protofuzz/protofuzz.py
@@ -221,16 +221,20 @@ def _module_to_generators(pb_module):
 
 
 def from_file(protobuf_file):
-    """Return dict of generators from a path to a .proto file.
-
+    """Return dict of generators from a path to a .proto file or pre-generated _pb2.py file.
+    _pb2.py file should be the output of the Protobuf compiler; users should not attempt to import arbitrary Python files.
+    
     Args:
-        protobuf_file(str) -- The path to the .proto file.
+        protobuf_file(str) -- The path to the .proto file or pre-generated _pb2.py file.
 
     Returns:
         A dict indexed by message name of ProtobufGenerator objects.
         These can be used to create inter-field dependencies or to generate messages.
 
     Raises:
+        AttributeError: If the _pb2.py file is not a valid generated file
+        FileNotFoundError: If the _pb2.py file is not found
+        ModuleNotFoundError: If there is a nested protobuf import, see issue #11
         BadProtobuf: If the .proto file is incorrectly formatted or named.
         ProtocNotFound: If the protoc compiler was not found on $PATH.
 


Test patch:
diff --git a/protofuzz/tests/test_pbimport.py b/protofuzz/tests/test_pbimport.py
index 17fd9ea..2edcfe3 100644
--- a/protofuzz/tests/test_pbimport.py
+++ b/protofuzz/tests/test_pbimport.py
@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 
 import os
+import re
 import secrets
 import stat
 import shutil
@@ -14,6 +15,7 @@ class TestPbimport(unittest.TestCase):
 
     def setUp(self):
         self.tempdir = tempfile.mkdtemp()
+        self.valid_filename = os.path.join(self.tempdir, "test.proto")
 
     def tearDown(self):
         shutil.rmtree(self.tempdir)
@@ -41,18 +43,58 @@ class TestPbimport(unittest.TestCase):
         module = pbimport.from_string(contents)
         self.assertTrue(hasattr(module, name))
 
-    def test_generate_and_import(self):
+    def test_from_generated(self):
         """
-        Test generation and loading of protobuf artifacts
+        Get a protobuf module from generated file
         """
         name = f'Msg{secrets.token_hex(16)}'
         contents = 'message {} {{ required int32 var = 1; }}\n'.format(name)
 
-        filename = os.path.join(self.tempdir, "test.proto")
+        with open(self.valid_filename, 'w') as f:
+            f.write(contents)
+    
+        dest = self.tempdir
+        full_path = os.path.abspath(self.valid_filename)
+        pbimport._compile_proto(full_path, dest)
+        target = os.path.join(dest, 'test_pb2.py')
+
+        module = pbimport.from_file(target)
+        self.assertTrue(hasattr(module, name))
+
+    def test_failing_import(self):
+        """
+        Test the failing of malformed generated protobuf file
+        """
+        contents = 'print("malformed generated code")'
+
+        filename = os.path.join(self.tempdir, "test_pb2.py")    
         with open(filename, 'w') as f:
             f.write(contents)
 
         module = pbimport.from_file(filename)
+        with self.assertRaises(AttributeError):
+            pbimport.types_from_module(module)
+
+    def test_failing_import_not_found(self):
+        """
+        Get a protobuf module from generated file
+        """
+        filename = os.path.join(self.tempdir, "test_pb2.py")
+    
+        with self.assertRaises(FileNotFoundError):
+            module = pbimport.from_file(filename)
+
+    def test_generate_and_import(self):
+        """
+        Test generation and loading of protobuf artifacts
+        """
+        name = f'Msg{secrets.token_hex(16)}'
+        contents = 'message {} {{ required int32 var = 1; }}\n'.format(name)
+
+        with open(self.valid_filename, 'w') as f:
+            f.write(contents)
+
+        module = pbimport.from_file(self.valid_filename)
         self.assertTrue(hasattr(module, name))
 
     def test_failing_generate_and_import(self):
@@ -61,9 +103,8 @@ class TestPbimport(unittest.TestCase):
         """
         contents = 'malformed protoc'
 
-        filename = os.path.join(self.tempdir, "test.proto")
-        with open(filename, 'w') as f:
+        with open(self.valid_filename, 'w') as f:
             f.write(contents)
 
         with self.assertRaises(pbimport.BadProtobuf):
-            module = pbimport.from_file(filename)
+            module = pbimport.from_file(self.valid_filename)