proto: add type stubs for generate schema

.github/mypy/mypy.ini

@@ -1,7 +1,4 @@
 [mypy]
-exclude = (?x)(
-    _pb2\.py$    # protobuf generated files
-  )
 
 [mypy-halo.*]
 ignore_missing_imports = True

.github/workflows/tests.yml

@@ -38,7 +38,7 @@ jobs:
     - name: Lint with black
       run: black -l 120 --extend-exclude ".*_pb2.py" --check .
     - name: Lint with pycodestyle
-      run: pycodestyle --show-source capa/ scripts/ tests/
+      run: pycodestyle --exclude="*_pb2.py" --show-source capa/ scripts/ tests/
     - name: Check types with mypy
       run: mypy --config-file .github/mypy/mypy.ini --check-untyped-defs capa/ scripts/ tests/
 

capa/render/proto/__init__.py

@@ -318,7 +318,6 @@ def emit_proto_message(out: StringIO, deferred_types: Dict, message):
             out.writeln(f"  oneof {name} {{")
 
             for j, of in enumerate(prop["anyOf"]):
-
                 if is_ref(of):
                     ptype = get_ref_type_name(of)
                     out.writeln(f"    {ptype} v{j} = {i};")

setup.py

@@ -28,6 +28,7 @@ requirements = [
     "dnfile==0.13.0",
     "dncil==1.0.2",
     "pydantic==1.10.4",
+    "protobuf==4.21.12",
 ]
 
 # this sets __version__
@@ -79,6 +80,7 @@ setuptools.setup(
             "psutil==5.9.2",
             "stix2==3.0.1",
             "requests==2.28.0",
+            "mypy-protobuf==3.4.0",
             # type stubs for mypy
             "types-backports==0.1.3",
             "types-colorama==0.4.15",
@@ -87,6 +89,7 @@ setuptools.setup(
             "types-termcolor==1.1.4",
             "types-psutil==5.8.23",
             "types_requests==2.28.1",
+            "types-protobuf 4.21.0.5",
         ],
         "build": [
             "pyinstaller==5.7.0",

tests/fixtures.py

@@ -1101,3 +1101,13 @@ def _0953c_dotnetfile_extractor():
 @pytest.fixture
 def _039a6_dotnetfile_extractor():
     return get_dnfile_extractor(get_data_path_by_name("_039a6"))
+
+
+@pytest.fixture
+def pma0101_rd():
+    path = os.path.join(CD, "data", "Practical Malware Analysis Lab 01-01.dll.json")
+    with open(path, "rb") as f:
+        buf = f.read()
+
+    src = buf.decode("utf-8")
+    return capa.render.result_document.ResultDocument.parse_raw(src)

tests/test_proto.py

@@ -8,12 +8,16 @@
 import pathlib
 import subprocess
 
+from fixtures import *
+
 import capa.render
 import capa.render.proto
 import capa.render.utils
 import capa.features.freeze
+import capa.render.proto.capa_pb2
 import capa.render.result_document
 import capa.features.freeze.features
+from capa.render.result_document import ResultDocument
 
 
 def test_generate_proto(tmp_path: pathlib.Path):
@@ -30,8 +34,25 @@ def test_generate_proto(tmp_path: pathlib.Path):
     print("=====================================")
     proto_path.write_text(proto)
 
-    subprocess.run(["protoc", "-I=" + str(tmp_path), "--python_out=" + str(tmp_path), str(proto_path)], check=True)
+    subprocess.run(
+        [
+            "protoc",
+            "-I=" + str(tmp_path),
+            "--python_out=" + str(tmp_path),
+            "--mypy_out=" + str(tmp_path),
+            str(proto_path),
+        ],
+        check=True,
+    )
 
     pb = tmp_path / "capa_pb2.py"
     print(pb.read_text())
     print("=====================================")
+
+
+def test_translate_to_proto(pma0101_rd: ResultDocument):
+    src = pma0101_rd
+
+    dst = capa.render.proto.capa_pb2.ResultDocument()
+
+    assert True