diff --git a/CybersecurityBenchmarks/README.md b/CybersecurityBenchmarks/README.md index d33e7d07..193fed7f 100644 --- a/CybersecurityBenchmarks/README.md +++ b/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" \ diff --git a/CybersecurityBenchmarks/datasets/canary_exploit/run.py b/CybersecurityBenchmarks/datasets/canary_exploit/run.py index 27aecf9b..65e3b108 100644 --- a/CybersecurityBenchmarks/datasets/canary_exploit/run.py +++ b/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)