solanaizer · jpmonette · Mar 10, 2024 · Mar 10, 2024 · Mar 10, 2024
diff --git a/ai_validator.py b/ai_validator.py
@@ -29,9 +29,9 @@ def analyze_vulnerability_with_gpt(api_key, file_content, filename: Path):
 
 If you find no errors, you should return an empty array.
 
-You are an Solana smart contract auditor. You are an expert at finding vulnerabilities that can be exploited by bad people.
+The filename key should contain the name of the module.
 
-Only output vulnerabilities which you are certain pose security risks.
+You are an Solana smart contract auditor. You are an expert at finding vulnerabilities that can be exploited by bad people.
 
 ```rs
 '{file_content}'
@@ -52,7 +52,7 @@ def analyze_vulnerability_with_gpt(api_key, file_content, filename: Path):
     if response.ok:
         response_json = response.json()
         response_content = response_json["choices"][0]["message"]["content"].replace("```json", "").replace("```", "")
-        if (response_content != ""):
+        if (response_content != "" or response_content != None or response_content != []):
             parsed = json.loads(response_content)
 
             for item in parsed:
@@ -72,4 +72,3 @@ def analyze_vulnerability_with_gpt(api_key, file_content, filename: Path):
             error_message = f"Failed to get a valid response from OpenAI: {response.status_code} - {response.text}"
         raise IOError(error_message)
 
-
diff --git a/requirements.txt b/requirements.txt
diff --git a/solanaizer.py b/solanaizer.py
@@ -13,15 +13,29 @@ def validate_file_content(file_path: Path):
     with open(file_path, 'r') as file:
         content = file.read()
 
-    return analyze_vulnerability_with_gpt(API_KEY, content, file_path)
+    try:
+        return analyze_vulnerability_with_gpt(API_KEY, content, file_path)
+    except:
+        return []
+
+
+
+def get_rust_files(directory):
+    rust_files = []
+    for root, _, files in os.walk(directory):
+        for file in files:
+            if file.endswith(".rs"):
+                rust_files.append(os.path.join(root, file))
+    return rust_files
 
 if __name__ == "__main__":
-    suffix = "src/lib.rs"
-    bug_free = "programs/bug-free-contract-1/"
-    non_bug_free = "programs/buggy-contract-1/"
+    dir_to_search = "programs/"
 
+    rust_files = get_rust_files(dir_to_search)
 
-    file_path_bug_free = Path(bug_free + suffix)
-    file_path_buggy = Path(non_bug_free + suffix)
+    json_dumps = []
+    for rust_file in rust_files:
+        rust_file_path = Path(rust_file)
+        json_dumps += validate_file_content(rust_file_path)
 
-    print(json.dumps(validate_file_content(file_path_bug_free) + validate_file_content(file_path_buggy)))
+    print(json.dumps(json_dumps))