Team Ai
Apppublic

puneethshree/security-code-reviewer

sourceHugging Faceupdated 5mo agoView on Hugging Face
0likes
create_test_data.py85 linesDownload Raw Back to scripts
1"""Generate test ZIP files and diffs from eval examples"""2 3import json4import zipfile5import os6from pathlib import Path7 8BASE_DIR = Path(__file__).parent.parent9 10def create_test_case_1():11    """Test case 1: Mixed CWE-89 and CWE-78"""12    13    # Load examples14    with open(str(BASE_DIR) + '/eval/ground_truth/cwe_89_sql_injection.json') as f:15        cwe_89 = json.load(f)16    with open(str(BASE_DIR) + '/eval/ground_truth/cwe_78_command_injection.json') as f:17        cwe_78 = json.load(f)18    19    # Pick examples20    examples = {21        'src/admin_routes.py': cwe_89[0],  # cwe_89_fastapi_00122        'src/network_utils.py': cwe_78[1],  # cwe_78_flask_00223        'src/product_routes.py': cwe_89[1], # cwe_89_flask_00224    }25    26    # Create temp directory27    test_dir = Path('test_data/test_case_1')28    test_dir.mkdir(parents=True, exist_ok=True)29    30    # Create source directory31    src_dir = test_dir / 'src'32    src_dir.mkdir(exist_ok=True)33    34    # Write files35    for filepath, example in examples.items():36        file_path = test_dir / filepath37        file_path.parent.mkdir(parents=True, exist_ok=True)38        file_path.write_text(example['vulnerable_code'])39        print(f"Created: {file_path}")40    41    # Create ZIP42    zip_path = Path('test_data/test_case_1.zip')43    with zipfile.ZipFile(zip_path, 'w') as zipf:44        for filepath in examples.keys():45            file_path = test_dir / filepath46            zipf.write(file_path, filepath)47    48    print(f"\nZIP created: {zip_path}")49    50    # Generate diff51    diff_lines = []52    for filepath, example in examples.items():53        vuln_line = example['vulnerability_line']54        code_lines = example['vulnerable_code'].split('\n')55        56        # Create minimal diff showing the vulnerable line57        diff_lines.append(f"diff --git a/{filepath} b/{filepath}")58        diff_lines.append(f"index 0000000..1111111 100644")59        diff_lines.append(f"--- a/{filepath}")60        diff_lines.append(f"+++ b/{filepath}")61        diff_lines.append(f"@@ -{vuln_line},1 +{vuln_line},1 @@")62        diff_lines.append(f"+{code_lines[vuln_line - 1]}")  # -1 for 0-indexing63        diff_lines.append("")64    65    diff_text = '\n'.join(diff_lines)66    diff_path = Path('test_data/test_case_1_diff.txt')67    diff_path.write_text(diff_text)68    69    print(f"Diff created: {diff_path}")70    print(f"\nExpected findings:")71    for filepath, example in examples.items():72        print(f"  - {filepath}: {example['cwe']} (line {example['vulnerability_line']})")73    74    return zip_path, diff_path75 76if __name__ == '__main__':77    print("Creating test case 1...")78    zip_path, diff_path = create_test_case_1()79    80    print(f"\n{'='*60}")81    print("Test data ready!")82    print(f"{'='*60}")83    print(f"\nUpload to FastAPI:")84    print(f"  ZIP file: {zip_path}")85    print(f"  Diff text: Copy from {diff_path}")