Add command line parameters for generated canary_exploit prompts and …

…memory_corruption tests

Reviewed By: dwjsong

Differential Revision: D64008834

fbshipit-source-id: d783335ac39e3be1321dd06544aa1453558d1e63

CybersecurityBenchmarks/README.md

@@ -374,7 +374,9 @@ Due to the nuanced nature of buffer overflows, we currently only support X86-64
 and not Apple Silicon.
 
 ```
-python3 -m CybersecurityBenchmarks.datasets.canary_exploit.run
+python3 -m CybersecurityBenchmarks.datasets.canary_exploit.run \
+   --prompt_path="$DATASETS/canary_exploit/canary_exploit.json" \
+   --memory_corruption_path="$DATASETS/canary_exploit/memory_corruption/"
 
 python3 -m CybersecurityBenchmarks.benchmark.run \
     --benchmark="canary-exploit" \

CybersecurityBenchmarks/datasets/canary_exploit/run.py

@@ -21,14 +21,6 @@
 
 from .memory_corruption.generator import generate_advanced_memory_corruption_test_cases
 
-PROMPT_DEFAULT_PATH = (
-    "CybersecurityBenchmarks/datasets/canary_exploit/canary_exploit.json"
-)
-
-MEMORY_CORRUPTION_DEFAULT_PATH = (
-    "CybersecurityBenchmarks/datasets/canary_exploit/memory_corruption/"
-)
-
 
 def generate_test_cases(
     generators: List[BaseTestGenerator], num_test_case: int
@@ -56,7 +48,7 @@ def generate_constraint_satisfaction_test_cases(
 
 
 def generate_memory_corruption_test_cases(
-    num_test_case: int,
+    num_test_case: int, memory_corruption_path: str
 ) -> List[Dict[str, str]]:
     generators: List[BaseTestGenerator] = [
         CBufferOverflowGenerator(),
@@ -66,8 +58,8 @@ def generate_memory_corruption_test_cases(
 
     # Advanced buffer overflow static tests
     prompts += build_with_prompt(
-        Path(MEMORY_CORRUPTION_DEFAULT_PATH),
-        Path(MEMORY_CORRUPTION_DEFAULT_PATH + "build"),
+        Path(memory_corruption_path),
+        Path(memory_corruption_path + "build"),
         EXAMPLES,
     )
 
@@ -104,6 +96,16 @@ def invoke_main() -> None:
         default=100,
         help="This is the argument to specify the number of tests to generate.",
     )
+    parser.add_argument(
+        "--prompt_path",
+        default="CybersecurityBenchmarks/datasets/canary_exploit/canary_exploit.json",
+        help="This is the path to which prompts will be written for the canary exploit tests.",
+    )
+    parser.add_argument(
+        "--memory_corruption_path",
+        default="CybersecurityBenchmarks/datasets/canary_exploit/memory_corruption/",
+        help="This is the argument to specify the path containing the memory corruption tests.",
+    )
 
     args = parser.parse_args()
     generators: List[BaseTestGenerator]
@@ -114,7 +116,9 @@ def invoke_main() -> None:
         prompts = generate_constraint_satisfaction_test_cases(args.num_test_case)
 
         # Add Memory Corruption Tests (static)
-        prompts += generate_memory_corruption_test_cases(args.num_test_case)
+        prompts += generate_memory_corruption_test_cases(
+            args.num_test_case, args.memory_corruption_path
+        )
     elif args.language == "c":
         generators = [
             CGenerator(),
@@ -136,10 +140,12 @@ def invoke_main() -> None:
         ]
         prompts = generate_test_cases(generators, args.num_test_case)
     elif args.language == "memory_corruption":
-        prompts = generate_memory_corruption_test_cases(args.num_test_case)
+        prompts = generate_memory_corruption_test_cases(
+            args.num_test_case, args.memory_corruption_path
+        )
 
     # Dump the prompts to a json file
-    with open(PROMPT_DEFAULT_PATH, "w") as f:
+    with open(args.prompt_path, "w") as f:
         json.dump(prompts, f, indent=4)