基于openai通过prompt实现清洗nuclei脚本

前言：基于prompt实现nuclei过滤脚本，拿到的nuclei脚本中带有脏数据，这边基于openai通过prompt来实现清洗脏数据

from concurrent.futures import ThreadPoolExecutor, as_completed

import requests
import json
import threading
from datetime import datetime

from common.log_manage import log
from common.sql_manage import DatabaseHandler

# mock group_id
group_id = 51112445814284

# url
target_url = 'https://wx.zsxq.com/dweb2/index/group/51112445814284/{topic_id}'

# 获取当前日期和时间
current_datetime = datetime.now()

# 将日期格式化为 "YYYY-MM-DD"
formatted_date = current_datetime.strftime("%Y-%m-%d")

# 创建一个全局锁
gLock = threading.Lock()

# remote url
remote_url = "http://api.remote/"

# nuclei整理 prompt
prompt_data = '''
你是一名专业的基于nuclei的漏洞编写者。

你的任务是从给定的带有脏数据的nuclei模板脚本内容中提取出完整且可运行的nuclei格式的yaml文件。

请根据以下要求返回结果：

1. 优先匹配id字段：确保nuclei模板内容中存在id字段。

2. 返回的JSON格式：{"status": "", "title": "", "content": ""}

- status：布尔类型，如果成功识别nuclei模板，则为true；如果识别失败，则为false；
- title：字符串类型，提取的漏洞名称。
- content：字符串类型，识别出的nuclei内容。

注意：

1. 仅返回JSON格式的数据，确保没有其他任何形式的数据出现。
'''

# url
url = "http://10.12.1.94:8001/v1/chat/completions"

# 你可能需要的API密钥或其他认证信息
headers = {"Content-Type": "application/json"}


def get_prompt_data(data):
    # 要发送的数据，这是一个JSON字典，包含你需要向API传递的信息
    tongyi_data = {
        "model": "qwen-72b",
        "messages": [
            {"role": "system", "content": prompt_data},
            {"role": "user", "content": data}
        ]
    }
    return tongyi_data


def send_data(text):
    tongyi_prompt_data = get_prompt_data(text)

    # 将数据转换为JSON格式
    data_json = json.dumps(tongyi_prompt_data)

    try:
        # 发送POST请求
        response = requests.post(url=url, headers=headers, data=data_json)

        # 检查响应状态码以确保请求成功
        if response.status_code == 200:
            json_data = response.json()
            return 'success', json_data['choices'][0]['message']['content']

        return 'fail', response.text

    except Exception as e:
        return 'error', e.__str__()


def read_file():
    # 读取 JSON 文件
    with open('./downloads/2024-06-13/马赛克安全/topics.json', 'r', encoding='utf-8') as file:
        data = json.load(file)

    return data if data else None


def process_page(item):
    db_handler = DatabaseHandler()

    # 遍历items数据，查找特定 user_id 并输出对应的 text
    if not isinstance(item, dict):
        return

    item_keys = item.keys()
    if 'talk' not in item_keys or 'topic_id' not in item_keys:
        return

    if not isinstance(item["talk"], dict):
        return

    talk_keys = item["talk"].keys()
    if 'owner' not in talk_keys or 'text' not in talk_keys:
        return

    if not isinstance(item["talk"]["owner"], dict):
        return

    owner_keys = item["talk"]["owner"].keys()
    if 'user_id' not in owner_keys:
        return

    owner_id = item["talk"]["owner"]["user_id"]

    if owner_id != 812511422885242:
        return

    result_dic = {}
    result_dic["source_url"] = target_url.format(topic_id=item['topic_id'])
    try:
        print("正在爬取数据：" + result_dic["source_url"])

        # 爬取地址之前，先检查一次
        if db_handler.check_vul_info_exist_by_sourceurl(result_dic["source_url"]):
            return f'文章地址：{result_dic["source_url"]} 已存在'

        response_state, content = send_data(item['talk']['text'])
        # print(content)
        if response_state == 'success':
            json_content = json.loads(content)
            if json_content['status']:
                result_dic["title"] = json_content['title']
                result_dic["content"] = str(json_content['content'])
                result_dic["content_type"] = 'nuclei'
                result_dic["source_name"] = '知识星球-马赛克安全'
                result_dic["content_sha1"] = None
                result_dic["media_press_url"] = ''
                result_dic["created_at"] = datetime.now()
                result_dic["synced_at"] = datetime.now()

                # 插入数据库
                db_handler.insert_vul_info_by_dict(result_dic)

                # 发送remote 
                # send(page)

                print(result_dic)

                return f"处理完成：{result_dic['source_url']}"
            else:
                return f"处理失败：{result_dic['source_url']}, 失败原因：不符合nuclei脚本格式"
        else:
            return f"处理失败：{result_dic['source_url']}, 失败原因：{content}"
    except Exception as e:
        return f"处理失败：{result_dic['source_url']}, 失败原因：{e.__str__()}"
    finally:
        if db_handler.connection is not None:
            db_handler.connection.close()


if __name__ == '__main__':

    items = read_file()
    all_count = len(items)

    # 设置线程池的最大线程数为5
    with ThreadPoolExecutor(max_workers=5) as executor:
        # 提交所有页面处理任务
        future_to_page = {executor.submit(process_page, item): item for item in items}
        # 初始化进度计数器
        count = 0
        # 等待每个任务完成并获取结果
        for future in as_completed(future_to_page):
            # 获取返回结果
            result = future.result()
            if result is not None:
                # 结果输出
                print(result)
                # 日志记录
                log('知识星球_马赛克安全', result)
                # 进度条
            count += 1
            # 进度输出
            print(f"当前进度：{count}/{all_count}")

    print("所有页面处理完成。")