Merge pull request #6 from mxrch/onlynamed

Adding the --named-keychains arg

.gitignore

@@ -7,3 +7,4 @@ protodeep.egg-info/
 *.proto
 *.pdeep
 .venv
+*_pb2.py

protodeep/cli.py

@@ -16,6 +16,7 @@ def parse_and_run():
     parser.add_argument('-t', '--type', required=True, type=str, help="Either protobuf (raw protobuf content), or protodeep (a ProtoDeep file).")
     parser.add_argument('-d', '--definitions', type=str, help="The file containing the custom protobuf definitions.")
     parser.add_argument('-na', '--no-autodetect', action='store_true', default=False, help="Don't try to autodetect if it's a raw HTTP request.")
+    parser.add_argument('-nk', '--named-keychains', action='store_true', default=False, help="Show and extract only named keychains.")
     parser.add_argument('-s', '--stdin', action='store_true', default=False, help="Parse from stdin.")
     parser.add_argument('-b', '--base64', action='store_true', default=False, help="If this is a base64 input, so it automatically decodes it.")
     parser.add_argument('-hx', '--hex', action='store_true', default=False, help="If this is a hex input, so it automatically decodes it.")
@@ -59,5 +60,6 @@ def process_args(args: argparse.Namespace):
         data_type=args.type,
         compile=args.compile,
         no_print=args.no_print,
-        schema_name=args.name
+        schema_name=args.name,
+        named_keychains=args.named_keychains
     )

protodeep/internals.py

@@ -16,12 +16,12 @@ def find_proto_schema(body: bytes, bruteforce_index: int=20):
             pass
     return output, i
 
-def clean_schema(schema, parsed, definitions={}) -> tuple[dict[str], bool]:
+def clean_schema(schema, parsed, definitions={}, named_keychains=False) -> tuple[dict[str], bool]:
     global custom_types_defined
 
     custom_types_defined = False
 
-    def _clean_schema(message: dict, parsed: dict, key_chain: str="", defs: dict={}):
+    def _clean_schema(message: dict, parsed: dict, key_chain: str="", defs: dict={}, named_keychains=False):
         global custom_types_defined
 
         original_message = deepcopy(message)
@@ -50,16 +50,20 @@ def _clean_schema(message: dict, parsed: dict, key_chain: str="", defs: dict={})
                         sparsed = parsed[skey]
                         if isinstance(sparsed, list):
                             sub["seen_repeated"] = True # Repeated sub
-                        sub["message_typedef"] = _clean_schema(sub["message_typedef"], sparsed, new_key_chain, defs)
+                        sub["message_typedef"] = _clean_schema(sub["message_typedef"], sparsed, new_key_chain, defs, named_keychains)
                 elif isinstance(parsed, list):
                     for item in parsed:
                         if skey in item:
                             sparsed = item[skey]
                             if isinstance(sparsed, list):
                                 sub["seen_repeated"] = True # Repeated sub
-                            sub["message_typedef"] = _clean_schema(sub["message_typedef"], sparsed, new_key_chain, defs)
-            message[key] = sub
+                            sub["message_typedef"] = _clean_schema(sub["message_typedef"], sparsed, new_key_chain, defs, named_keychains)
+
+            if named_keychains and not sub.get("message_typedef") and not sub.get("name"):
+                del(message[key])
+            else:
+                message[key] = sub
         return message
 
-    _clean_schema(schema, parsed, defs=definitions)
+    schema = _clean_schema(schema, parsed, defs=definitions, named_keychains=named_keychains)
     return schema, custom_types_defined

protodeep/lib.py

@@ -158,7 +158,7 @@ def _pretty_print(schema: dict[str, dict], parsed: dict, key_chain: str="", pret
         cs = Console(highlight=False)
         _pretty_print(self.schema, self.values)
 
-def guess_schema(data: bytes, definitions: dict={}, bruteforce_index=20, no_autodetect=False) -> ProtoDeepSchema:
+def guess_schema(data: bytes, definitions: dict={}, bruteforce_index=20, no_autodetect=False, named_keychains=False) -> ProtoDeepSchema:
     if not no_autodetect:
         if data.strip().startswith(b"HTTP") and b"\r\n\r\n" in data:
             print("[!] Full request detected, extracting the body...")
@@ -168,12 +168,12 @@ def guess_schema(data: bytes, definitions: dict={}, bruteforce_index=20, no_auto
     parsed = schema[0]
     new_schema = schema[1]
 
-    new_schema, custom_types_defined = clean_schema(new_schema, parsed, definitions=definitions)
+    new_schema, custom_types_defined = clean_schema(new_schema, parsed, definitions=definitions, named_keychains=named_keychains)
     if custom_types_defined:
         schema = decode_message(data[data_index:], new_schema)
         parsed = schema[0]
         new_schema = schema[1]
-        new_schema, _ = clean_schema(new_schema, parsed, definitions=definitions)
+        new_schema, _ = clean_schema(new_schema, parsed, definitions=definitions, named_keychains=named_keychains)
 
     protodeep_schema = ProtoDeepSchema(schema=new_schema, values=parsed)
     return protodeep_schema

protodeep/modules/analyze.py

@@ -12,6 +12,7 @@ def main(data: bytes=b"",
         no_print: bool=False,
         data_type: str="",
         no_autodetect: bool=False,
+        named_keychains: bool=False,
         bruteforce_index: int=10,
         match_any: str="",
         match_keychain: str="",
@@ -70,10 +71,12 @@ def main(data: bytes=b"",
                     data=data,
                     definitions=definitions,
                     bruteforce_index=bruteforce_index,
-                    no_autodetect=no_autodetect
+                    no_autodetect=no_autodetect,
+                    named_keychains=named_keychains
                 )
-            except:
-                exit("[-] Can't decode the data. Please verify your type and bruteforce_index. Otherwise, RIP. 🥹")
+            except Exception as err:
+                print("Error :", err)
+                exit("\n[-] Can't decode the data. Please verify your type and bruteforce_index. Otherwise, RIP. 🥹")
         case "protodeep":
             import pickle
             from protodeep.lib import ProtoDeepSchema, clean_schema

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "protodeep"
-version = "1.0.3"
+version = "1.1.0"
 authors = [
     {name = "mxrch", email = "mxrch.dev@pm.me"},
 ]