danger/LLM_patch
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

init commit

Browse Source
This commit is contained in:
danger
2024-10-15 10:49:25 +08:00
parent 96bd52bfcc
commit 8dd7505095
15 changed files with 1105 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

BIN
__pycache__/chat.cpython-312.pyc Normal file
View File

Binary file not shown.

106
chat.py Normal file
View File

@@ -0,0 +1,106 @@
"""
Interfaces to interact with various LLMs
"""
import json
import os
import atexit
import configparser
from typing import Dict, Optional, List
from openai import OpenAI
CUR_DIR = os.path.dirname(os.path.abspath(__file__))
CONFIG = os.path.join(CUR_DIR, 'config.ini')
def load_config(field: str, value: str) -> str:
config = configparser.ConfigParser()
config.read(CONFIG)
return config[field][value]
def llm_configured() -> bool:
model = load_config('LLM', 'model')
api_key = load_config('LLM', 'api_key')
api_base = load_config('LLM', 'api_base')
return bool(len(model) and len(api_key) and len(api_base))
class QueryChatGPT():
""" Interface for interacting with ChatGPT
"""
def __init__(self) -> None:
self.chat_context: List[Dict[str, str]] = []
self.chat_history: List[Dict[str, str]] = []
self.temperature:float = 0.2
self.use_history = False
self.system_prompt: Optional[str] = None
atexit.register(self.log_history)
def clear(self):
self.chat_context = []
def set_history(self, open: bool) -> None:
self.use_history = open
def insert_system_prompt(self, system_prompt: str) -> None:
""" add system_prompt in self.chat_context """
if self.chat_context and self.chat_context[0]["role"] == "system":
self.chat_context[0]['content'] = system_prompt
else:
self.chat_context.insert(0, {
"role": "system",
"content": system_prompt
})
def log_history(self, log_file: str = 'chat_log.json'):
if not os.path.exists(log_file):
with open(log_file, 'w') as w:
json.dump([], w, indent=4)
with open(log_file, 'r') as r:
log = json.load(r)
assert (isinstance(log, list))
log.append(self.chat_history)
with open(log_file, 'w') as w:
json.dump(log, w, indent=4)
def __query(self, prompt: str, model: str) -> Optional[str]:
self.chat_context.append({"role": "user", "content": prompt})
self.chat_history.append({"role": "user", "content": prompt})
client = OpenAI(api_key=load_config('LLM', 'api_key'), base_url=load_config('LLM', 'api_base'))
response = client.chat.completions.create(
messages=self.chat_context, # type: ignore
model=model,
temperature=self.temperature,
)
response_content = str(response.choices[0].message.content)
self.chat_context.append({
"role": "assistant",
"content": response_content
})
self.chat_history.append({
"role": "assistant",
"content": response_content
})
return response_content
def query(self,
prompt: str,
*,
model: str = load_config('LLM', 'model')) -> Optional[str]:
response = self.__query(prompt, model)
if not self.use_history:
self.clear()
return response

184
chat_log.json Normal file
View File

File diff suppressed because one or more lines are too long

7
config.ini Normal file
View File

@@ -0,0 +1,7 @@
[LLM]
model = gpt-3.5-turbo
api_key = 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
api_base = https://api.openai.com/v1/
; For openai online LLMs, the default api_base is https://api.openai.com/v1/
; For local LLMs, the api_base is usually like http://ip:port/v1/

14
danger_func.json Normal file
View File

@@ -0,0 +1,14 @@
[
{
"type": "buffer_overflow",
"name": "strcpy"
},
{
"type": "buffer_overflow",
"name": "recv"
},
{
"type": "format_string",
"name": "dprintf"
}
]

81
input/dprintf_extract.c Normal file
View File

@@ -0,0 +1,81 @@
//Function: echo_handler ->0x4199190 7 perm->5
// bad sp value at call has been detected, the output may be wrong!
int __cdecl echo_handler(int sock)
{
char buffer[256]; // [rsp+10h] [rbp-100h] BYREF
memset(buffer, 0, sizeof(buffer));
return recv(sock, &buffer[8], 0x100uLL, 0) > 0 && dprintf(sock, buffer) > 0;
}
//Function: main ->0x4199603 7 perm->5
// bad sp value at call has been detected, the output may be wrong!
int __fastcall main(int argc, const char **argv, const char **envp)
{
uint16_t v3; // ax
char client_addr_str[24]; // [rsp+0h] [rbp-40h] BYREF
int addrlen; // [rsp+18h] [rbp-28h] BYREF
int opt; // [rsp+1Ch] [rbp-24h] BYREF
sockaddr_in address; // [rsp+20h] [rbp-20h] BYREF
int new_socket; // [rsp+38h] [rbp-8h]
int server_fd; // [rsp+3Ch] [rbp-4h]
opt = 1;
addrlen = 16;
server_fd = socket(2, 1, 0);
if ( !server_fd )
{
perror("socket failed");
exit(1);
}
if ( setsockopt(server_fd, 1, 15, &opt, 4u) )
{
perror("setsockopt");
exit(1);
}
address.sin_family = 2;
address.sin_addr.s_addr = 0;
address.sin_port = htons(0x2B00u);
if ( bind(server_fd, (const struct sockaddr *)&address, 0x10u) < 0 )
{
perror("bind failed");
exit(1);
}
if ( listen(server_fd, 3) < 0 )
{
perror("listen");
exit(1);
}
printf("TCP server listening on port %d\n", 11008);
new_socket = accept(server_fd, (struct sockaddr *)&address, (socklen_t *)&addrlen);
if ( new_socket < 0 )
{
perror("accept");
exit(1);
}
inet_ntop(2, &address.sin_addr, client_addr_str, 0x10u);
v3 = ntohs(address.sin_port);
printf("Accept %s:%d\n", client_addr_str, v3);
while ( echo_handler(new_socket) )
;
close(new_socket);
return 0;
}
//Function: backdoor ->0x4200059 7 perm->5
// bad sp value at call has been detected, the output may be wrong!
int __cdecl backdoor()
{
char *new_envp[2]; // [rsp+0h] [rbp-20h] BYREF
char *new_argv[2]; // [rsp+10h] [rbp-10h] BYREF
dup2(4, 0);
dup2(4, 1);
dup2(4, 2);
execve("/bin/sh", new_argv, new_envp);
return 0;
}

191
input/edit_extract.c Normal file
View File

@@ -0,0 +1,191 @@
//Function: add ->0x4199478 7 perm->5
void __cdecl add(char *str)
{
Node *newNode; // [rsp+18h] [rbp-8h]
newNode = (Node *)malloc(0x108uLL);
strcpy(newNode->data, str);
newNode->next = head;
head = newNode;
}
//Function: delete ->0x4199559 7 perm->5
void __cdecl delete(char *str)
{
Node *entry; // [rsp+10h] [rbp-10h]
Node **current; // [rsp+18h] [rbp-8h]
for ( current = &head; *current; current = &entry->next )
{
entry = *current;
if ( !strcmp((*current)->data, str) )
{
*current = entry->next;
free(entry);
return;
}
}
}
//Function: edit ->0x4199684 7 perm->5
void __cdecl edit(char *oldStr, char *newStr)
{
Node *current; // [rsp+18h] [rbp-8h]
for ( current = head; current; current = current->next )
{
if ( !strcmp(current->data, oldStr) )
{
strcpy(current->data, newStr);
return;
}
}
}
//Function: show ->0x4199787 7 perm->5
void __cdecl show(int client_sock)
{
size_t v1; // rax
char buffer[1024]; // [rsp+10h] [rbp-410h] BYREF
Node *current; // [rsp+418h] [rbp-8h]
for ( current = head; current; current = current->next )
{
snprintf(buffer, 0x400uLL, "%s\n", current->data);
v1 = strlen(buffer);
send(client_sock, buffer, v1, 0);
}
}
//Function: main ->0x4199929 7 perm->5
// local variable allocation has failed, the output may be wrong!
// bad sp value at call has been detected, the output may be wrong!
int __fastcall __noreturn main(int argc, const char **argv, const char **envp)
{
int opt; // [rsp+Ch] [rbp-C44h] BYREF
char arg2[1024]; // [rsp+10h] [rbp-C40h] BYREF
char arg1[1035]; // [rsp+410h] [rbp-840h] BYREF
char command[5]; // [rsp+81Bh] [rbp-435h] BYREF
_BYTE buffer[1032]; // [rsp+820h] [rbp-430h] OVERLAPPED BYREF
int addrlen; // [rsp+C2Ch] [rbp-24h] BYREF
sockaddr_in address; // [rsp+C30h] [rbp-20h] BYREF
int new_socket; // [rsp+C48h] [rbp-8h]
int server_fd; // [rsp+C4Ch] [rbp-4h]
addrlen = 16;
*(_QWORD *)buffer = 0LL;
*(_QWORD *)&buffer[8] = 0LL;
memset(&buffer[24], 0, 0x3F0uLL);
opt = 1;
server_fd = socket(2, 1, 0);
if ( !server_fd )
{
perror("socket failed");
exit(1);
}
if ( setsockopt(server_fd, 1, 15, &opt, 4u) )
{
perror("setsockopt");
exit(1);
}
address.sin_family = 2;
address.sin_addr.s_addr = 0;
address.sin_port = htons(0x2B04u);
if ( bind(server_fd, (const struct sockaddr *)&address, 0x10u) < 0 )
{
perror("bind failed");
exit(1);
}
if ( listen(server_fd, 3) < 0 )
{
perror("listen");
exit(1);
}
printf("Server listening on port %d\n", 11012);
while ( 1 )
{
new_socket = accept(server_fd, (struct sockaddr *)&address, (socklen_t *)&addrlen);
if ( new_socket < 0 )
break;
read(new_socket, buffer, 0x400uLL);
memset(command, 0, sizeof(command));
memset(arg1, 0, 0x400uLL);
memset(arg2, 0, sizeof(arg2));
((void (*)(_BYTE *, const char *, ...))__isoc99_sscanf)(buffer, "%4s %1023s %1023s", command, arg1, arg2);
if ( !strcmp(command, "ADD") )
{
add(arg1);
}
else if ( !strcmp(command, "DEL") )
{
delete(arg1);
}
else if ( !strcmp(command, "EDIT") )
{
edit(arg1, arg2);
}
else if ( !strcmp(command, "SHOW") )
{
show(new_socket);
}
else
{
puts("Unknown command.");
}
close(new_socket);
}
perror("accept");
exit(1);
}
//Function: backdoor ->0x4200706 7 perm->5
// bad sp value at call has been detected, the output may be wrong!
int __cdecl backdoor()
{
char *new_envp[2]; // [rsp+0h] [rbp-40h] BYREF
char *new_argv[2]; // [rsp+10h] [rbp-30h] BYREF
sockaddr_in serv_addr; // [rsp+20h] [rbp-20h] BYREF
int sock; // [rsp+3Ch] [rbp-4h]
new_argv[0] = "/bin/sh";
new_argv[1] = 0LL;
new_envp[0] = "PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin";
new_envp[1] = 0LL;
sock = socket(2, 1, 0);
if ( sock >= 0 )
{
serv_addr.sin_family = 2;
serv_addr.sin_port = htons(0x2EECu);
if ( inet_pton(2, "127.0.0.1", &serv_addr.sin_addr) > 0 )
{
if ( connect(sock, (const struct sockaddr *)&serv_addr, 0x10u) >= 0 )
{
dup2(sock, 0);
dup2(sock, 1);
dup2(sock, 2);
execve(new_argv[0], new_argv, new_envp);
return 0;
}
else
{
puts("\nConnection Failed ");
return -1;
}
}
else
{
puts("\nInvalid address/ Address not supported ");
return -1;
}
}
else
{
puts("\n Socket creation error ");
return -1;
}
}

87
input/recv_extract.c Normal file
View File

@@ -0,0 +1,87 @@
//Function: echo_handler ->0x4199222 7 perm->5
// bad sp value at call has been detected, the output may be wrong!
int __cdecl echo_handler(int sock)
{
char buffer[256]; // [rsp+10h] [rbp-100h] BYREF
memset(buffer, 0, sizeof(buffer));
if ( recv(sock, &buffer[8], 0x400uLL, 0) <= 0 )
return 0;
printf("Message from client: %s\n", buffer);
if ( send(sock, "Hello from server\n", 0x12uLL, 0) <= 0 )
return 0;
puts("Hello message sent");
return 1;
}
//Function: main ->0x4199683 7 perm->5
// bad sp value at call has been detected, the output may be wrong!
int __fastcall main(int argc, const char **argv, const char **envp)
{
uint16_t v3; // ax
char client_addr_str[24]; // [rsp+0h] [rbp-40h] BYREF
int addrlen; // [rsp+18h] [rbp-28h] BYREF
int opt; // [rsp+1Ch] [rbp-24h] BYREF
sockaddr_in address; // [rsp+20h] [rbp-20h] BYREF
int new_socket; // [rsp+38h] [rbp-8h]
int server_fd; // [rsp+3Ch] [rbp-4h]
opt = 1;
addrlen = 16;
server_fd = socket(2, 1, 0);
if ( !server_fd )
{
perror("socket failed");
exit(1);
}
if ( setsockopt(server_fd, 1, 15, &opt, 4u) )
{
perror("setsockopt");
exit(1);
}
address.sin_family = 2;
address.sin_addr.s_addr = 0;
address.sin_port = htons(0x2AFFu);
if ( bind(server_fd, (const struct sockaddr *)&address, 0x10u) < 0 )
{
perror("bind failed");
exit(1);
}
if ( listen(server_fd, 3) < 0 )
{
perror("listen");
exit(1);
}
printf("TCP server listening on port %d\n", 11007);
new_socket = accept(server_fd, (struct sockaddr *)&address, (socklen_t *)&addrlen);
if ( new_socket < 0 )
{
perror("accept");
exit(1);
}
inet_ntop(2, &address.sin_addr, client_addr_str, 0x10u);
v3 = ntohs(address.sin_port);
printf("Accept %s:%d\n", client_addr_str, v3);
while ( echo_handler(new_socket) )
;
close(new_socket);
return 0;
}
//Function: backdoor ->0x4200139 7 perm->5
// bad sp value at call has been detected, the output may be wrong!
int __cdecl backdoor()
{
char *new_envp[2]; // [rsp+0h] [rbp-20h] BYREF
char *new_argv[2]; // [rsp+10h] [rbp-10h] BYREF
dup2(4, 0);
dup2(4, 1);
dup2(4, 2);
execve("/bin/sh", new_argv, new_envp);
return 0;
}

10
input/test.c Normal file
View File

@@ -0,0 +1,10 @@
//Function: main ->0x4199603 7 perm->5
// bad sp value at call has been detected, the output may be wrong!
int __fastcall main(int argc, const char **argv, const char **envp)
{
return 0;
}

26
output/dprintf_extract/patch_dprintf.json Normal file
View File

@@ -0,0 +1,26 @@
[
{
"file_path": "/root/LLM_prompt/output//dprintf_extract/patch_dprintf.json",
"vul_info": " there is a dprintf format string vulnerability in the echo_handler function. The vulnerability lies in the line `dprintf(sock, buffer)`, where the buffer is being passed as the format string to the dprintf function without proper formatting. This can potentially lead to a format string vulnerability."
},
{
"file_path": "dprintf_extract.c",
"vul_info": " there is a dprintf format string vulnerability in the echo_handler function. The vulnerability lies in the line `dprintf(sock, buffer)`, where the buffer is being passed as the format string to dprintf without proper formatting. This can potentially lead to a format string vulnerability."
},
{
"file_path": "dprintf_extract.c",
"vul_info": " there is a dprintf format string vulnerability in the echo_handler function. The vulnerability lies in the line `dprintf(sock, buffer)`, where the buffer is passed as the format string to the dprintf function without proper formatting. This can potentially lead to a format string vulnerability."
},
{
"file_path": "dprintf_extract.c",
"vul_info": " there is a dprintf format string vulnerability in the echo_handler function. The vulnerability lies in the line `dprintf(sock, buffer)`, where the buffer is being passed as the format string to dprintf without proper formatting. This can potentially lead to a format string vulnerability."
},
{
"file_path": "dprintf_extract.c",
"vul_info": " there is a dprintf format string vulnerability in the echo_handler function. The vulnerability lies in the line `dprintf(sock, buffer)`, where the buffer is directly passed as the format string to the dprintf function without proper formatting. This can potentially lead to a format string vulnerability."
},
{
"file_path": "dprintf_extract.c",
"vul_info": " there is a dprintf format string vulnerability in the echo_handler function. The vulnerability lies in the line `dprintf(sock, buffer)`, where the buffer is passed as the format string to the dprintf function without proper formatting. This can lead to a format string vulnerability, allowing an attacker to potentially read or write arbitrary memory locations."
}
]

14
output/dprintf_extract/patch_recv.json Normal file
View File

@@ -0,0 +1,14 @@
[
{
"file_name": "dprintf_extract.c",
"vul_info": " reason: The recv function in the echo_handler function does not check the size of the buffer before writing data to it. The buffer size is 256 bytes, but the recv function is writing up to 0x100 bytes (256 bytes) to the buffer, which can lead to a buffer overflow vulnerability.",
"fix_size": 248,
"patch_info": "In the echo_handler function, the buffer size is 256 bytes, but only 248 bytes are being used for the recv function (starting from buffer[8]). This leaves 8 bytes unused, which could potentially lead to a buffer overflow if the recv function receives more than 248 bytes of data. To ensure the program works safely, the recv function should be prepared to receive a maximum of 248 bytes of data."
},
{
"file_name": "dprintf_extract.c",
"vul_info": " reason: The recv function in the echo_handler function does not check the size of the buffer before writing data to it. This can lead to a buffer overflow vulnerability if the data received is larger than the size of the buffer (256 bytes). An attacker could potentially exploit this vulnerability to overwrite memory beyond the buffer and execute arbitrary code.",
"fix_size": 248,
"patch_info": "In the echo_handler function, the buffer size is 256 bytes. However, only 248 bytes are being used for the recv function (buffer[8] to buffer[255]). This leaves 8 bytes at the beginning of the buffer unused, which could potentially lead to a buffer overflow if the recv function receives more than 248 bytes of data."
}
]

26
output/recv_extract/patch_recv.json Normal file
View File

@@ -0,0 +1,26 @@
[
{
"file_name": "recv_extract.c",
"vul_info": " reason: The recv function in the echo_handler function does not check the size of the buffer before writing data to it, which can lead to a buffer overflow vulnerability. The buffer size is 256 bytes, but the recv function is writing data starting from buffer[8], potentially allowing for more data to be written than the buffer can hold.",
"fix_size": 256,
"patch_info": "In the echo_handler function, the buffer size is set to 256 bytes. Therefore, when using recv to receive data from the socket, the program should be prepared to handle up to 256 bytes of data to prevent buffer overflow."
},
{
"file_name": "recv_extract.c",
"vul_info": " reason: The recv function in the echo_handler function does not check the size of the buffer being passed to it. It receives data into the buffer starting from index 8, which could potentially lead to a buffer overflow if the received data is larger than the allocated buffer size.",
"fix_size": 256,
"patch_info": "In the echo_handler function, the buffer size is set to 256 bytes. Therefore, when using recv to receive data from the socket, the program should be prepared to handle up to 256 bytes of data to prevent buffer overflow."
},
{
"file_name": "recv_extract.c",
"vul_info": " reason: The recv function in the echo_handler function does not check the size of the buffer before copying data into it, which can lead to a buffer overflow vulnerability. The buffer size is 256 bytes, but the recv function is copying data into buffer starting from index 8, potentially allowing more data to be copied than the buffer can hold.",
"fix_size": 256,
"patch_info": "In the echo_handler function, the buffer size is set to 256 bytes. Therefore, when using recv to receive data from the socket, the program should be prepared to handle up to 256 bytes of data to avoid buffer overflow."
},
{
"file_name": "recv_extract.c",
"vul_info": " reason: The recv function in the echo_handler function does not check the size of the buffer before copying data into it. This can lead to a buffer overflow vulnerability if the data received is larger than the size of the buffer (256 bytes). An attacker could potentially exploit this vulnerability to overwrite memory beyond the buffer and execute malicious code.",
"fix_size": 248,
"patch_info": "The buffer size should be reduced by 8 to account for the offset used in recv to prevent buffer overflow."
}
]

42
prompt.json Normal file
View File

@@ -0,0 +1,42 @@
[
{
"type": "attack",
"name": "dprintf",
"prompt": {
"role": "user",
"content": "Please analyze whether the following program contains a dprintf format string vulnerability. If you find dprintf format string vulnerability please answer with 'yes, reason here', or else just output 'no dprintf vulnerability'.\n{code}"
}
},
{
"type": "attack",
"name": "strcpy",
"prompt": {
"role": "user",
"content": "Please analyze the following code to determine if the strcpy used in the program poses a risk of buffer overflow. If it actually exists then answer with 'yes, reason here', or else output 'no strcpy vulnerability'.\n{code}"
}
},
{
"type": "attack",
"name": "recv",
"prompt": {
"role": "user",
"content": "Please analyze whether the following program contains the risk of buffer overflow and only check recv function. If it actually exists then answer with 'yes, reason here', or else output 'no recv vulnerability'.\n{code}"
}
},
{
"type": "patch",
"name": "strcpy",
"prompt": {
"role": "user",
"content": "There is a risk of buffer overflow when using strcpy in the program, and I want to replace it with the strncpy function. Please tell me explicitly the size shoud be prepared for strncpy to make sure the program work safely.\n{code}"
}
},
{
"type": "patch",
"name": "recv",
"prompt": {
"role": "user",
"content": "There is a risk of buffer overflow when using recv in the program. Please tell me the specific size be prepared for recv to ensure the program work safely. Answer with 'size=value, ', and then give me the reason.\n{code}"
}
}
]

69
prompt.py Normal file
View File

@@ -0,0 +1,69 @@
"""
This script is designed for editing the bare query
bare query - questions or information that inspire
ChatGPT to understand the program deeply instead of
caring about the response of ChatGPT
{
'type': 'bare_query',
'name': 'XXX',
'prompt': {
'role': 'user',
'content': 'You should provide programming advice.',
}
}
"""
import os
import json
BPATH = os.path.dirname(os.path.abspath(__file__))
BPATH = os.path.join(BPATH, 'prompt.json')
print(BPATH)
assert(os.path.exists(BPATH))
with open(BPATH, 'r') as r:
lst = json.load(r)
while True:
choice = input('1. list 2. add new 3. keyword-based search 4. delete 5. exit: ')
if choice == '1':
print('\n-----------------')
for _id, _query in enumerate(lst):
_type = _query['type']
_name = _type
if 'name' in _query.keys():
_name = _query['name']
print(_id, f': ({_type}) ({_name})', repr(_query['prompt']['content']))
elif choice == '2':
_type = input('Input your query type:')
_name = input('Input your query name (use type as default):')
if not _name:
_name = _typeprompt
_query = input('Input your query content:')
new_entry = {
'type': _type,
'name': _name,
'prompt': {
'role': 'user',
'content': _query,
}
}
lst.append(new_entry)
elif choice == '3':
k = input('Input your keyword:')
for _id, _query in enumerate(lst):
_c = _query['prompt']['content']
if k in _c:
print(_id, ':', _c)
elif choice == '4':
i = input('Input the id of deleted enetry:')
lst.pop(int(i))
elif choice == '5':
with open(BPATH, 'w') as w:
json.dump(lst, w, indent=4)
exit(0)
else:
pass

248
test.py Normal file
View File

@@ -0,0 +1,248 @@
import os
import re
import sys
import json
import warnings
from chat import QueryChatGPT, llm_configured, load_config
from typing import Optional, Dict, List, Tuple
DIR = os.path.dirname(os.path.abspath(__file__))
PROMPT_PATH = os.path.join(DIR, 'prompt.json')
DANGER_FUNC = os.path.join(DIR, 'danger_func.json')
IN_DIR = os.path.join(DIR, 'input')
OUT_DIR = os.path.join(DIR, 'output')
def get_prompt(name: str, _type: str, prompt_path: str = PROMPT_PATH) -> Optional[Dict[str, str]]:
"""
Access the prompt
Args:
name: the name of the prompt
_type: the type of the prompt
prompt_path: the path of the prompt file
Returns:
a dict containing two keys: 'role' and 'content'
"""
prompts = None
with open(prompt_path, 'r') as f:
prompts = json.load(f)
assert(prompts)
for _p in prompts:
if _p['name'] == name and _p['type'] == _type:
return _p['prompt']
return None
def read_decompile_code(file_path: str) -> Optional[str]:
if not os.path.exists(file_path):
warnings.warn("Fail to find {file_path}!".format(file_path, file_path))
sys.exit(1)
with open(file_path, 'r') as r:
code_data = r.read()
return code_data
data = {
"name": "Henry",
"age": 30,
"fuzzing": True,
"languages": ["C", "Python", "Assembly"]
}
def handle_dprintf(file_name: str, code: str, patch_dprintf_file: str = 'patch_dprintf.json'):
if 'dprintf' not in code:
return
print("begin to test dprintf in {file_name}.".format(file_name = file_name))
# generate output file
output_dir = OUT_DIR + '/' + file_name[:-2]
os.makedirs(output_dir, exist_ok=True)
output_file = output_dir + '/' + patch_dprintf_file
print("The dprintf info store into the " + output_file)
# get prompt from prompt.json
prompt = get_prompt('dprintf', 'attack')
assert (prompt)
q = QueryChatGPT()
response = q.query(prompt['content'].format(code = code))
# judge whether the program exists dprintf format string vulnerability
print("response info : " + response[:4].lower())
if 'yes' not in response[:4].lower():
return
# store the relevant information into INPUT_DIR + "patch_dprintf.json"
data = {
"file_path": file_name,
"vul_info": response[4:] # set response[4:] to skip string "yes, "
}
if not os.path.exists(output_file):
with open(output_file, 'w') as w:
json.dump([], w, indent=4)
with open(output_file, 'r') as r:
log = json.load(r)
assert (isinstance(log, list))
# insert relevant vulnerability info into json file
log.append(data)
with open(output_file, 'w') as w:
json.dump(log, w, indent=4)
# invoke patch api here
# patch_dprintf(file_path)
def handle_recv(file_name: str, code: str, patch_recv_file: str = 'patch_recv.json'):
if 'recv' not in code:
return
print("begin to test recv in {file_name}.".format(file_name = file_name))
# generate output file
output_dir = OUT_DIR + '/' + file_name[:-2]
os.makedirs(output_dir, exist_ok=True)
output_file = output_dir + '/' + patch_recv_file
print("The recv info store into the " + output_file)
# get prompt from prompt.json
prompt = get_prompt('recv', 'attack')
assert (prompt)
q = QueryChatGPT()
response = q.query(prompt['content'].format(code = code))
print(response)
# judge whether the program exists buffer overflow vulnerability due to recv func
print("response info : " + response[:4].lower())
if 'yes' not in response[:4].lower():
return
# store the relevant information into INPUT_DIR + "patch_recv.json"
data = {
"file_name": file_name,
"vul_info": response[4:] # set response[4:] to skip string "yes, "
}
# determine the specific size to fix recv func
# get patch prompt for recv
prompt = get_prompt('recv', 'patch')
assert (prompt)
q = QueryChatGPT()
response = q.query(prompt['content'].format(code = code))
print(response)
# record modified size
match = re.search(r'size=(\d+)', response)
data['fix_size'] = int(match.group(1))
match = re.search(r',\s*(.*)', response)
data['patch_info'] = match.group(1)
if not os.path.exists(output_file):
with open(output_file, 'w') as w:
json.dump([], w, indent=4)
with open(output_file, 'r') as r:
log = json.load(r)
assert (isinstance(log, list))
# insert relevant vulnerability info into json file
log.append(data)
with open(output_file, 'w') as w:
json.dump(log, w, indent=4)
# invoke patch api here
# patch_recv(file_path)
def handle_strcpy(file_name: str, code: str, patch_strcpy_file: str = 'patch_strcpy.json'):
if 'strcpy' not in code:
return
print("begin to test strcpy in {file_name}.".format(file_name = file_name))
# generate output file
output_dir = OUT_DIR + '/' + file_name[:-2]
os.makedirs(output_dir, exist_ok=True)
output_file = output_dir + '/' + patch_strcpy_file
print("The recv info store into the " + output_file)
# get prompt from prompt.json
prompt = get_prompt('strcpy', 'attack')
assert (prompt)
q = QueryChatGPT()
response = q.query(prompt['content'].format(code = code))
print(response)
# judge whether the program exists buffer overflow vulnerability due to recv func
print("response info : " + response[:4].lower())
if 'yes' not in response[:4].lower():
return
# store the relevant information into INPUT_DIR + "patch_recv.json"
data = {
"file_name": file_name,
"vul_info": response[4:] # set response[4:] to skip string "yes, "
}
# determine the specific size to fix recv func
# get patch prompt for recv
prompt = get_prompt('strcpy', 'patch')
assert (prompt)
q = QueryChatGPT()
response = q.query(prompt['content'].format(code = code))
print(response)
# record modified size
# match = re.search(r'size=(\d+)', response)
# data['fix_size'] = int(match.group(1))
# match = re.search(r',\s*(.*)', response)
# data['patch_info'] = match.group(1)
if not os.path.exists(output_file):
with open(output_file, 'w') as w:
json.dump([], w, indent=4)
with open(output_file, 'r') as r:
log = json.load(r)
assert (isinstance(log, list))
# insert relevant vulnerability info into json file
log.append(data)
with open(output_file, 'w') as w:
json.dump(log, w, indent=4)
# invoke patch api here
# patch_recv(file_path)
def exp():
print("Trying to test function normally")
test_file = "edit_extract.c"
# test_file = "recv_extract.c"
# test_file = "dprintf_extract.c"
# read code which will be used to analyze by LLM
code = read_decompile_code(IN_DIR + "/" + test_file)
# read danger_func.json to determine the scope of the func checked
with open(DANGER_FUNC, 'r') as f:
d_func = json.load(f)
assert(d_func)
for _d in d_func:
if _d['name'] == 'dprintf':
handle_dprintf(test_file, code)
elif _d['name'] == 'recv':
handle_recv(test_file, code)
elif _d['name'] == 'strcpy':
handle_strcpy(test_file, code)
if __name__ == '__main__':
if not llm_configured():
print('please complete llm access setup first...')
exit()
exp()