scripts: allow run_ida_script to forward command line arguments to th…

…e script. Mimic -S switch when using IDA. idapython: add ARGV to idc
williballenthin · Aug 20, 2020 · 17392e9 · 17392e9
1 parent 5b7a300
commit 17392e9
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 3 deletions.
diff --git a/idb/idapython.py b/idb/idapython.py
@@ -769,6 +769,10 @@ def __init__(self, db, api):
         else:
             raise RuntimeError("unexpected wordsize")
 
+        # Command line arguments passed to idapython scripts. Args passed via -S
+        # switch in IDA
+        self.ARGV = []
+
         ## Mantain API compatibility for API < 7
         self.GetMnem = self.print_insn_mnem
         self.GetOpnd = self.print_operand

diff --git a/scripts/run_ida_script.py b/scripts/run_ida_script.py
@@ -5,6 +5,11 @@
 author: Willi Ballenthin
 email: willi.ballenthin@gmail.com
 """
+import re
+import sys
+import os.path
+import logging
+
 import argparse
 import logging
 import os.path
@@ -16,6 +21,17 @@
 logger = logging.getLogger(__name__)
 
 
+def script_path_to_args(script_path):
+    # Split based on spaces but preserve words inside double and single quotes
+    regex_extract_quoted = "[^\s\"']+|\"([^\"]*)\"|'([^']*)'"
+    matches = re.finditer(regex_extract_quoted, script_path)
+    # Strip quotes away using groups
+    return [
+        m.group(2) if m.group(2) else m.group(1) if m.group(1) else m.group(0)
+        for m in matches
+    ]
+
+
 def main(argv=None):
     # TODO: do version check for 3.x
 
@@ -25,7 +41,13 @@ def main(argv=None):
     parser = argparse.ArgumentParser(
         description="Dump an IDB B-tree to a textual representation."
     )
-    parser.add_argument("script_path", type=str, help="Path to script file")
+    parser.add_argument(
+        "script_path",
+        type=str,
+        help='Path to script file. \
+                              Command line arguments can be passed using quotes: \
+                              "myscrypt.py arg1 arg2 \\"arg3 arg3\\""',
+    )
     parser.add_argument("idbpath", type=str, help="Path to input idb file")
     parser.add_argument(
         "-v", "--verbose", action="store_true", help="Enable debug logging"
@@ -61,12 +83,15 @@ def main(argv=None):
 
         hooks = idb.shim.install(db, ScreenEA=screenea)
 
+        script_args = script_path_to_args(args.script_path)
         # update sys.path to point to directory containing script.
         # so scripts can import .py files in the same directory.
-        script_dir = os.path.dirname(args.script_path)
+        script_dir = os.path.dirname(script_args[0])
         sys.path.insert(0, script_dir)
+        # update idc.ARGV
+        hooks["idc"].ARGV = script_args
 
-        with open(args.script_path, "rb") as f:
+        with open(script_args[0], "rb") as f:
             g = {
                 "__name__": "__main__",
             }