From 0b394c39966bd451918db26aed30758a332c2719 Mon Sep 17 00:00:00 2001 From: psrok1 Date: Wed, 31 Jul 2024 16:51:25 +0200 Subject: [PATCH] Fix: add proper argument parser for CLI tool --- drakpdb/main.py | 37 ++++++++++++++++++++++++++++++++----- 1 file changed, 32 insertions(+), 5 deletions(-) diff --git a/drakpdb/main.py b/drakpdb/main.py index 1f7cb20..7ba3fe9 100644 --- a/drakpdb/main.py +++ b/drakpdb/main.py @@ -7,12 +7,39 @@ def main(): parser = argparse.ArgumentParser(description="drakpdb") - parser.add_argument("action", type=str, help="one of: fetch_pdb, parse_pdb") - parser.add_argument( - "pdb_name", type=str, help="name of pdb file without extension, e.g. ntkrnlmp" + + pdbname_subparser = argparse.ArgumentParser(add_help=False) + pdbname_subparser.add_argument( + "pdb_name", type=str, help="name of the pdb file" + ) + + dllname_subparser = argparse.ArgumentParser(add_help=False) + dllname_subparser.add_argument( + "dll_name", type=str, help="path to the dll file" + ) + + guidage_subparser = argparse.ArgumentParser(add_help=False) + guidage_subparser.add_argument( + "guid_age", type=str, help="guid/age of the pdb file" ) - parser.add_argument("guid_age", nargs="?", help="guid/age of the pdb file") + action = parser.add_subparsers(help="Action commands", dest="action") + action.required = True + action.add_parser( + "parse_pdb", + parents=[pdbname_subparser], + help="Parse PDB file into Rekall profile", + ) + action.add_parser( + "fetch_pdb", + parents=[pdbname_subparser, guidage_subparser], + help="Fetch PDB file matching PDB name and GUID/Age", + ) + action.add_parser( + "pe_codeview_data", + parents=[dllname_subparser], + help="Get PDB name and GUID/Age for DLL path" + ) args = parser.parse_args() if args.action == "parse_pdb": @@ -21,6 +48,6 @@ def main(): elif args.action == "fetch_pdb": fetch_pdb(args.pdb_name, args.guid_age) elif args.action == "pe_codeview_data": - print(pe_codeview_data(args.file)) + print(pe_codeview_data(args.dll_name)) else: raise RuntimeError("Unknown action")