From 1dc433445f5ffbb6f7b8ed286e23c2abd562b924 Mon Sep 17 00:00:00 2001
From: Jean-Philippe Monette <contact@jpmonette.net>
Date: Sun, 10 Mar 2024 10:19:06 +0000
Subject: [PATCH] up

---
 ai_validator.py  |  7 +++----
 requirements.txt |  2 --
 solanaizer.py    | 21 +++++++++++++++------
 3 files changed, 18 insertions(+), 12 deletions(-)
 delete mode 100644 requirements.txt

diff --git a/ai_validator.py b/ai_validator.py
index 05ca347..fc26c69 100644
--- 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
deleted file mode 100644
index 1c62f1d..0000000
--- a/requirements.txt
+++ /dev/null
@@ -1,2 +0,0 @@
-requests
-pathlib
\ No newline at end of file
diff --git a/solanaizer.py b/solanaizer.py
index a155d54..7f887e9 100644
--- a/solanaizer.py
+++ b/solanaizer.py
@@ -15,13 +15,22 @@ def validate_file_content(file_path: Path):
 
     return analyze_vulnerability_with_gpt(API_KEY, content, file_path)
 
+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))