【A9】GPT代码漏洞挖掘小实践

【A9】GPT代码漏洞挖掘小实践
2023-4-11 21:45:41 Author: mp.weixin.qq.com(查看原文) 阅读量:42 收藏

—

简介

在GPT出来后，发现能做的事情实在太多了。尤其是技术方面，当有思路就可以找GPT锻炼(tou lan)一下它的技术能力，让它帮忙写代码啥的，这感觉是真的舒服，但是GPT目前弊端也很明显，无法上传文件，只能以文本输入问题，而输出也有局限性，3.5版本输出只有4096个tokens，4的版本虽然有32K的tokens，回答也比3.5准确多了，但是GPT4每天有使用限制，而且API接口也没有全开放，只能申请。

—

背景

近期，朋友告诉我可以用llama_index库对GPT训练，就能做代码审计，不需要将源代码上传到GPT上就能实现，虽然GPT也没有上传按钮。

听到这个消息，让我这个不会代码审计的菜鸡欲欲跃试，要是实现了，是不是对外称“我也会代码审计”哈哈哈哈。于是利用周末时间好好学了一下，翻阅了一下llama_index库的简单使用方法，便有了这篇文章。

—
什么是llama_index

llama_index主要通过对大规模语言模型的训练数据进行索引，使模型能够更快地找到与查询相关的信息。基本原理是将数据分成多个块，并将每个块存储在多个节点上。每个节点负责处理其本地块的查询请求，并将结果广播给其他节点。这种分布式存储和查询方式可以有效地减少数据传输和处理时间。

如向GPT-3提问时，GPT可能会在其知识库中找到多个相关片段，可能给出的答案有误。但是如果通过用模型的训练的数据建立索引，就可以快速查找到相关信息，从而提高回答问题的准确性和效率。

—

如何使用llama_index对GPT训练

1、准备一个GPT和训练数据集。训练数据集包含了大量文本信息；

2、使用llama_index对训练数据集进行索引。这个过程就像给一本书建立目录一样，目的是帮助模型更快地查找到相关的信息；

3、将训练数据集的文本向量化；

4、使用llama_index查找与问题相关的文本。查找会从训练数据集的索引中找到；

5、将找到的相关文本片段作为提示提供给GPT模型；

6、GPT模型根据提示生成回答，并反馈回答。

主要核心是llama_index提供的索引信息，让GPT模型更快地找到与问题相关的知识片段，从而提高回答问题的准确性和效率。

代码如下：

# -*- coding: utf-8 -*-# 导入所需库import openaiimport osfrom langchain.chat_models import ChatOpenAI
# 设置openai的API密钥openai.api_key = "your api key"os.environ["OPENAI_API_KEY"] = openai.api_key
# 导入llama_index库中的相关组件from llama_index import (    GPTSimpleVectorIndex,    SimpleDirectoryReader,    LLMPredictor,    PromptHelper,    ServiceContext)
# 设置最大输入大小max_input_size = 4096# 设置输出令牌数量num_output = 1000# 设置最大块重叠max_chunk_overlap = 20
# 设置块大小限制chunk_size_limit = 1000
# 初始化PromptHelperprompt_helper = PromptHelper(max_input_size, num_output, max_chunk_overlap, chunk_size_limit=chunk_size_limit)
# 初始化LLMPredictorllm_predictor = LLMPredictor(llm=ChatOpenAI(temperature=0, model_name="gpt-3.5-turbo", max_tokens=num_output))
# 初始化ServiceContextservice_context = ServiceContext.from_defaults(llm_predictor=llm_predictor, prompt_helper=prompt_helper, chunk_size_limit=1000)
# 从文档中读取数据documents = SimpleDirectoryReader("./taocms", recursive=True, errors='ignore').load_data()
# 构建索引index = GPTSimpleVectorIndex.from_documents(documents, service_context=service_context)
#保存索引index1.jsonindex.save_to_disk('index1.json')
# 对输入问题进行查询response = index.query(    "请审计这套源码，请将审计出的源码构造出http包,并用中文描述。",    service_context=service_context,    similarity_top_k=3,    mode="default")
# 输出查询结果print(response)

效果测试：

拿靶场测试了一下，有时候会构造HTTP包，但也有时候不会构造。存在有不稳定性，而且也因为GPT3输出tokens大小只有4096，反馈的时候也会停顿，或者只能审计出一个漏洞后就没下文了。

—

总结

由于llama_index目前还在实验阶段，实操的时候踩了好多坑，翻阅官方资料和操作demo，目前也将踩坑记录下。

踩坑一：

源码中可能因为图片的原因导致报错，解决方案去掉源码里的图片即可。

踩坑二：

分块的时候出现错误，表示无解，换套源码试试玩吧...

踩坑三：

GPT3容易胡说八道，回答的答案不是所需的，多运行几次就可以，但是也会消耗tokens。简直就是拿钱在消耗，幸好免费送了18美刀。

总的来说，目前阶段只能玩一玩，简单的靶机源码还是可以审计出来，如果是大型一点的源代码，就无法审计出或者误报。

由于时间关系和对于llama_index模块了解比较浅，做的实验存在的局限性很大，如审计只能审计一个漏洞、json文件每次训练都会将训练过的被覆盖等问题，当前也没时间去解决这些问题。

如果再深入了解AI模型训练的话，创建个专注代码审计的知识库，随着不断训练，知识库的数据不断增加，再加上GPT4来审计，未来是否真的可以让GPT做代码审计呢？

看来目前还是不能对外称“我会代码审计”。不过，有一个入安全圈不久的小同学，拿ChatGPT挖了两个高危以上CVE。有三年代码审计的同学，被刚入圈不久的同学，拿ChatGPT反杀了，以后怎么办？

文章来源: https://mp.weixin.qq.com/s?__biz=MzUzODU3ODA0MA==&mid=2247488819&idx=1&sn=32f76f80e4e9574357f55333e47915c2&chksm=fad4c824cda34132effefe3e5f2d72ea720622847aacad7ab8bb739fd10dc9b283bc946f5cb3&scene=58&subscene=0#rd
如有侵权请联系:admin#unsafe.sh