-
Notifications
You must be signed in to change notification settings - Fork 1
/
scan_eosio.py
93 lines (69 loc) · 3.5 KB
/
scan_eosio.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
import os
import openai
import fnmatch
import spacy
def find_files(path, extensions):
found_files = []
for root, _, filenames in os.walk(path):
for ext in extensions:
for filename in fnmatch.filter(filenames, f'*.{ext}'):
found_files.append(os.path.join(root, filename))
return found_files
nlp = spacy.load("en_core_web_sm")
def split_code_into_chunks(code, max_tokens, overlap_tokens):
doc = nlp(code)
tokens = [token.text for token in doc]
chunks = []
start_token = 0
current_line = 1
line_counts = []
while start_token < len(tokens):
end_token = start_token + max_tokens
while end_token < len(tokens) and tokens[end_token] != '\n': # Look for newline character
end_token -= 1
chunk_text = "".join(tokens[start_token:end_token + 1])
chunks.append(chunk_text)
next_line = current_line + chunk_text.count('\n')
line_counts.append((current_line, next_line - 1))
current_line = next_line
start_token = end_token + 1 - overlap_tokens
return chunks, line_counts
def scan_eosio_file(file_path, api_key, engine, max_tokens=1000, overlap_tokens=10):
# Connect to the OpenAI API
openai.api_key = api_key
# Read the Solidity file content
with open(file_path) as f:
solidity_code = f.read()
# Split the code into smaller chunks
code_chunks, line_counts = split_code_into_chunks(solidity_code, max_tokens, overlap_tokens)
vulnerabilities = []
for i, (chunk, (line_start, line_end)) in enumerate(zip(code_chunks, line_counts)):
total_chunks = len(code_chunks)
if i == 0:
prompt = (f"Analyze vulnerabilities in the following EOSIO smart contract code. "
f"The code is divided into chunks. This is chunk {i+1}/{total_chunks} "
f"(lines {line_start}-{line_end}):\n{chunk}\n\n"
f"Please wait for all the chunks to complete before analyzing and "
f"providing your answer with the line numbers of where the vulnerabilities exist.")
else:
prompt = f"Here is chunk {i+1}/{total_chunks} (lines {line_start}-{line_end}):\n{chunk}\n\nPlease wait for all the chunks to complete before analyzing and providing your answer with the line numbers of where the vulnerabilities exist."
#Make a request to the OpenAI API to analyze the Solidity file
#result = openai.Completion.create(engine=engine, prompt=prompt, max_tokens=1000, n=1, stop=None, temperature=0.5) #/v1/completions
result = openai.ChatCompletion.create(model=engine, messages=[{"role": "system", "content": prompt}], max_tokens=1000, n=1, stop=None, temperature=0.5)
#vulnerabilities.append(result.choices[0].text.strip()) #/v1/completions
vulnerabilities.append(result['choices'][0]['message']['content'].strip())
return "\n".join(vulnerabilities)
def main():
api_key = os.environ['INPUT_OPENAI_API_KEY']
engine = os.environ['INPUT_ENGINE']
repo_path = os.getcwd()
file_extensions = ['cpp', 'hpp']
target_files = find_files(repo_path, file_extensions)
for file in target_files:
print(f"Scanning {file}...")
vulnerabilities = scan_eosio_file(file, api_key, engine)
with open(f"{file}_vulnerabilities.txt", "w") as f:
f.write("Service provided to you by www.sentnl.io.\n")
f.write(f"Vulnerabilities found in {file}:\n{vulnerabilities}\n")
if __name__ == "__main__":
main()