ChatGptStudy/staVul2proComp.py at master · AntChainOpenLabs/ChatGptStudy · GitHub

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
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
import json
#

# 创建一个空列表来存储所有字典
vulnerabilities = []
# 标准漏洞格式数据文件
readfilename = 'standardVulnerability.json'
# prompt&completion文件
wirtefilename = "prompt_completion.jsonl"

# prompt优化字符
prompt_end = "\n\n###\n\n"
completion_start = " "
completion_end = "###"

# 将标准漏洞格式文件加载到内存
def read_standard_vulnerability(filename):
    # 打开JSON文件并加载数据
    with open(filename, encoding='utf-8') as f:
        data = json.load(f)

    # 遍历JSON中的每个对象
    for item in data:
        # 创建一个字典来存储此对象的信息
        vulnerability = {}
        vulnerability['Code'] = item['Code']

        # 创建一个空列表来存储漏洞描述信息
        desc_list = []
        for desc in item['VulnerabilityDesc']:
            # 将每个漏洞描述信息存储到字典中
            desc_dict = {}
            desc_dict['Name'] = desc['Name']
            desc_dict['Location'] = desc['Location']
            desc_dict['Type'] = desc['Type']
            desc_dict['Description'] = desc['Description']
            desc_dict['Repair'] = desc['Repair']
            # 将此漏洞描述信息字典添加到漏洞描述列表中
            desc_list.append(desc_dict)

        # 将漏洞描述列表添加到漏洞字典中
        vulnerability['VulnerabilityDesc'] = desc_list

        # 将此漏洞字典添加到漏洞列表中
        vulnerabilities.append(vulnerability)

# 根据标准漏洞格式和多种提示生成规则，生成prompt和completion对。并写入文件
def generate_prompt_completion(filename):
    with open(filename,"w") as f :
        for vul in vulnerabilities:
            # prompt_completion_1(f, vul)
            # prompt_completion_2(f, vul)
            # prompt_completion_3(f, vul)
            # prompt_completion_4(f, vul)
            # prompt_completion_5(f, vul)
            prompt_completion_6(f, vul)
            prompt_completion_7(f, vul)
            prompt_completion_8(f, vul)
            prompt_completion_9(f, vul)


def prompt_optimization(prompt, completion):
    prompt_completion = {}
    prompt_completion['prompt'] = prompt + prompt_end
    prompt_completion['completion'] = completion_start + completion + completion_end
    return prompt_completion

def write_file(prompt_completion,f):
    json.dump(prompt_completion,f)
    f.write("\n")


def prompt_completion_1(f, vul):
    prompt = "Does this  smart contract  have any vulnerabilities?  " + vul['Code']
    completion = "here are some potential vulnerabilities in the given smart contract:  " + str(vul['VulnerabilityDesc'])

    prompt_completion = prompt_optimization(prompt, completion)
    write_file(prompt_completion,f)

def prompt_completion_2(f, vul):

    for Desc in vul['VulnerabilityDesc']:
        if Desc['Type'] == "":
            continue

        prompt = "Does this  smart contract  have any vulnerabilities? What is its vulnerability type?  " + vul['Code']
        completion = "here are some potential vulnerabilities in the given smart contract:  " + Desc['Description'] + "The vulnerability type is:  " + Desc['Type']

        prompt_completion = prompt_optimization(prompt, completion)
        write_file(prompt_completion,f)

def prompt_completion_3(f, vul):

    for Desc in vul['VulnerabilityDesc']:
        if Desc['Repair'] == "":
            continue

        prompt = "Does this  smart contract  have any vulnerabilities? How to fix it?  " + vul['Code']
        completion = "here are some potential vulnerabilities in the given smart contract:  " + Desc['Description'] + "To repair this Vulnerability:  " + Desc['Repair']

        prompt_completion = prompt_optimization(prompt, completion)
        write_file(prompt_completion,f)

def prompt_completion_4(f, vul):

    for Desc in vul['VulnerabilityDesc']:
        if Desc['Location'] == "":
            continue

        prompt = "Does this  smart contract  have any vulnerabilities? Where is the vulnerability?  " + vul['Code']
        completion = "here are some potential vulnerabilities in the given smart contract:  " + Desc['Description'] + "The vulnerability locates at:  " + Desc['Location']

        prompt_completion = prompt_optimization(prompt, completion)
        write_file(prompt_completion,f)

def prompt_completion_5(f, vul):

    for Desc in vul['VulnerabilityDesc']:
        if Desc['Type'] == "":
            continue

        prompt = "Does this  smart contract  have any vulnerabilities of " + Desc['Type'] +  "? "+ vul['Code']
        completion = "here are some potential vulnerabilities in the given smart contract:  " + Desc['Description']

        prompt_completion = prompt_optimization(prompt, completion)
        write_file(prompt_completion,f)


def prompt_completion_6(f, vul):

    for Desc in vul['VulnerabilityDesc']:
        if Desc['Type'] == "":
            continue
        if Desc['Location'] == "":
            continue
        prompt = "Does this  smart contract  have any vulnerabilities of " + Desc['Type'] +  "? "+ vul['Code']
        completion = "Yes , the vulnerabilities happens at " + Desc['Location']

        prompt_completion = prompt_optimization(prompt, completion)
        write_file(prompt_completion,f)

def prompt_completion_7(f, vul):

    for Desc in vul['VulnerabilityDesc']:
        if Desc['Type'] == "":
            continue
        if Desc['Location'] == "":
            continue
        prompt = "Does this  smart contract  have any vulnerabilities at " + Desc['Location'] +  "? "+ vul['Code']
        completion = "Yes , the contract have " + Desc['Type']

        prompt_completion = prompt_optimization(prompt, completion)
        write_file(prompt_completion,f)

def prompt_completion_8(f, vul):

    for Desc in vul['VulnerabilityDesc']:
        if Desc['Type'] == "":
            continue
        if Desc['Location'] == "":
            continue
        prompt = "Does this  smart contract  have any vulnerabilities ? "+ vul['Code']
        completion = "Yes , the contract have " + Desc['Type'] + " at " + Desc['Location']

        prompt_completion = prompt_optimization(prompt, completion)
        write_file(prompt_completion,f)

# 修复好的良性样本
def prompt_completion_9(f,vul):
    for Desc in vul['VulnerabilityDesc']:
        if Desc['Type'] == "":
            continue
        if Desc['Location'] == "":
            prompt = "Does this  smart contract  have any vulnerabilities ? " + vul['Code']
            completion = "No  , the contract don't have " + Desc['Type'] + " vulnerability"

            prompt_completion = prompt_optimization(prompt, completion)
            write_file(prompt_completion, f)


if __name__=="__main__":
    read_list = [
        r"json/bug_return.json",
        r"json/Benign_sample(return).json",
        r"json/bug_precision.json",
        r"json/Benign_sample(precision).json",
        r"json/bug_fake deposit.json",
        r"json/fixed_fake deposit.json",
        r"json/bug_fake notification.json",
        r"json/fixed_fake notification.json",

        r"json/bug_variable length array out of gas.json",
        r"json/fixed_variable length array out of gas.json"

    ]
    for f_name in read_list:
        read_standard_vulnerability(f_name)


    generate_prompt_completion(wirtefilename)