From da3ad91b2027316a8132039da19f42961da51ab5 Mon Sep 17 00:00:00 2001
From: Anonymous Maarten <madebr@users.noreply.github.com>
Date: Wed, 21 Jun 2023 23:33:08 +0200
Subject: [PATCH] recomp.py: use argparse to parse arguments (#30)

* recomp.py: use argparse to parse arguments

* Address code revew comments

* reccomp.py: -h/--help for help -H/--htmp for html

* update CI to use new arg

* slight string updates

---------

Co-authored-by: itsmattkc <34096995+itsmattkc@users.noreply.github.com>
---
 .github/workflows/build.yml |  4 +--
 tools/reccomp/reccomp.py    | 69 ++++++++++++++-----------------------
 2 files changed, 27 insertions(+), 46 deletions(-)

diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index 335b9699..e865a1c0 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -60,8 +60,8 @@ jobs:
         C:\msys64\usr\bin\wget.exe https://legoisland.org/download/ISLE.EXE
         C:\msys64\usr\bin\wget.exe https://legoisland.org/download/LEGO1.DLL
         pip install capstone
-        python3 tools/reccomp/reccomp.py -h ISLEPROGRESS.HTML ISLE.EXE Release/ISLE.EXE Release/ISLE.PDB ISLE
-        python3 tools/reccomp/reccomp.py -h LEGO1PROGRESS.HTML LEGO1.DLL Release/LEGO1.DLL Release/LEGO1.PDB LEGO1
+        python3 tools/reccomp/reccomp.py -H ISLEPROGRESS.HTML ISLE.EXE Release/ISLE.EXE Release/ISLE.PDB ISLE
+        python3 tools/reccomp/reccomp.py -H LEGO1PROGRESS.HTML LEGO1.DLL Release/LEGO1.DLL Release/LEGO1.PDB LEGO1
       
     - name: Upload Artifact
       uses: actions/upload-artifact@master
diff --git a/tools/reccomp/reccomp.py b/tools/reccomp/reccomp.py
index f1f378ae..4e234a61 100755
--- a/tools/reccomp/reccomp.py
+++ b/tools/reccomp/reccomp.py
@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 
+import argparse
 from capstone import *
 import difflib
 import struct
@@ -7,60 +8,40 @@
 import os
 import sys
 
-def print_usage():
-  print('Usage: %s [options] <original-binary> <recompiled-binary> <recompiled-pdb> <decomp-dir>\n' % sys.argv[0])
-  print('\t-v, --verbose <offset>\t\t\tPrint assembly diff for specific function (original file\'s offset)')
-  print('\t-h, --html <output-file>\t\t\tGenerate searchable HTML summary of status and diffs')
-  sys.exit(1)
+parser = argparse.ArgumentParser(allow_abbrev=False,
+  description='Recomp Compare: compare an original EXE with a recompiled EXE + PDB.')
+parser.add_argument('original', metavar='original-binary', help='The original binary')
+parser.add_argument('recompiled', metavar='recompiled-binary', help='The recompiled binary')
+parser.add_argument('pdb', metavar='recompiled-pdb', help='The PDB of the recompiled binary')
+parser.add_argument('decomp_dir', metavar='decomp-dir', help='The decompiled source tree')
+parser.add_argument('--verbose', '-v', metavar='offset', help='Print assembly diff for specific function (original file\'s offset)')
+parser.add_argument('--html', '-H', metavar='output-file', help='Generate searchable HTML summary of status and diffs')
+
+args = parser.parse_args()
 
-positional_args = []
 verbose = None
-skip = False
-html = None
+if args.verbose:
+  try:
+    verbose = int(args.verbose, 16)
+  except ValueError:
+    parser.error('invalid verbose argument')
+html = args.html
 
-for i, arg in enumerate(sys.argv):
-  if skip:
-    skip = False
-    continue
-
-  if arg.startswith('-'):
-    # A flag rather than a positional arg
-    flag = arg[1:]
-
-    if flag == 'v' or flag == '-verbose':
-      verbose = int(sys.argv[i + 1], 16)
-      skip = True
-    elif flag == 'h' or flag == '-html':
-      html = sys.argv[i + 1]
-      skip = True
-    else:
-      print('Unknown flag: %s' % arg)
-      print_usage()
-  else:
-    positional_args.append(arg)
-
-if len(positional_args) != 5:
-  print_usage()
-
-original = positional_args[1]
+original = args.original
 if not os.path.isfile(original):
-  print('Invalid input: Original binary does not exist')
-  sys.exit(1)
+  parser.error('Original binary does not exist')
 
-recomp = positional_args[2]
+recomp = args.recompiled
 if not os.path.isfile(recomp):
-  print('Invalid input: Recompiled binary does not exist')
-  sys.exit(1)
+  parser.error('Recompiled binary does not exist')
 
-syms = positional_args[3]
+syms = args.pdb
 if not os.path.isfile(syms):
-  print('Invalid input: Symbols PDB does not exist')
-  sys.exit(1)
+  parser.error('Symbols PDB does not exist')
 
-source = positional_args[4]
+source = args.decomp_dir
 if not os.path.isdir(source):
-  print('Invalid input: Source directory does not exist')
-  sys.exit(1)
+  parser.error('Source directory does not exist')
 
 # Declare a class that can automatically convert virtual executable addresses
 # to file addresses