diff options
| author | yifei cheng <[email protected]> | 2024-09-06 06:31:11 +0000 |
|---|---|---|
| committer | yifei cheng <[email protected]> | 2024-09-06 06:31:11 +0000 |
| commit | d4df6fa16bfc8031e13433137ad7db0ad8a7af7b (patch) | |
| tree | 9347af8f42c29c0946c23ddc1c205f3540daf161 | |
| parent | bc78f6cf67ac5470cd572e62ada6a114fdf1a789 (diff) | |
Upload New File
| -rw-r--r-- | strategy_adjust_code/kimi_main.py | 105 |
1 files changed, 105 insertions, 0 deletions
diff --git a/strategy_adjust_code/kimi_main.py b/strategy_adjust_code/kimi_main.py new file mode 100644 index 0000000..87d0282 --- /dev/null +++ b/strategy_adjust_code/kimi_main.py @@ -0,0 +1,105 @@ +from openai import OpenAI
+
+client = OpenAI(
+ api_key="sk-wFEcyFj8sS2mbncCEqYK5uxYa6VyQyJrQ8reKPFYzxQcwt3i",
+ base_url="https://api.moonshot.cn/v1",
+)
+
+# 加载背景介绍信息
+with open('./content_role_system.txt', 'r', encoding='utf-8') as f:
+ content_system = f.read()
+
+output_constraints = "请按下面规定的格式输出策略调整结果，从下面的集合中选择一种决策输出（不需要其他任何文字反馈，我需要处理格式化输出的内容）：{保持当前策略执行、调整脚本参数：参数代称 原始的数值->新的数值、更换攻击脚本：脚本名称}"
+
+mode_mapping = {"保持当前策略执行": "keep",
+ "调整脚本参数": "parameter",
+ "更换攻击脚本": "script"}
+
+history = [
+ {"role": "system",
+ "content": content_system}
+]
+
+
+def chat(query, history):
+ history.append({
+ "role": "user",
+ "content": query
+ })
+ completion = client.chat.completions.create(
+ model="moonshot-v1-32k",
+ messages=history,
+ temperature=0.3,
+ )
+ result = completion.choices[0].message.content
+ history.append({
+ "role": "assistant",
+ "content": result
+ })
+ return result
+
+
+def initial(content):
+ global history
+ history = [
+ {"role": "system",
+ "content": content_system}
+ ]
+
+ try:
+ clxz = content["clxz"] # 策略选择
+ mbgz = content["mbgz"] # 目标感知
+ ztgz = content["ztgz"] # 状态感知
+ script = content["script"] # 脚本及参数
+ except Exception as e:
+ raise Exception('post上传的数据key解析失败', e)
+
+ # 构建待上传ai的策略初始化信息
+ assert clxz in ["拒绝服务", "数据欺骗"], "策略选择只能是'拒绝服务'或者'数据欺骗'"
+ clxz_context = f"1. 测试的最终目的是：{clxz}"
+
+ mbgz_context = f"2. 当前目标节点的18维不变特性是：{mbgz}"
+
+ script = script.replace('；', ';').split(';')
+ script_context = f"3. 当前执行的脚本名称为{script[0]}；" + ("脚本参数选择为" + "；".join(script[1:]) if len(script) > 1 else "")
+
+ ztgz_context = f"4. 当前（未执行脚本前）目标节点的状态参数为：{ztgz}"
+
+ context = '\n'.join([clxz_context, mbgz_context, script_context, ztgz_context])
+
+ # 初始化交互不需要输出响应
+ chat(context, history)
+ return 1
+
+
+def adjusting(content):
+ try:
+ ztgz = content["ztgz"] # 状态感知
+ except Exception as e:
+ raise Exception("post上传的数据key解析失败", e)
+
+ context = f"当前目标节点的状态感知结果为：{ztgz}\n{output_constraints}"
+
+ response = chat(context, history)
+
+ mode = response.replace(":", "：").split("：")[0]
+ assert mode in ["保持当前策略执行", "调整脚本参数", "更换攻击脚本"], f"响应出现非标情况，当前模式响应为{mode}"
+ mode_context = mode_mapping[mode]
+
+ try:
+ if mode_context == "keep":
+ adjust_context = "none"
+ elif mode_context == "parameter":
+ # 例如：模型响应：调整脚本参数：doh_ddos_n 4->8 输出adjust参数为doh_ddos_n-8
+ para = response.replace(":", "：").split("：")[1]
+ adjust_context = "-".join([para.split(" ")[0], para.split(">")[-1]])
+ else:
+ # 例如：模型响应：更换攻击脚本：DoT_DDoS 输出adjust参数为DoT_DDoS
+ adjust_context = response.replace(":", "：").split("：")[1]
+ except Exception as e:
+ raise Exception("处理模型规范化输出错误", e)
+
+ # # debug
+ # print(history)
+
+ return {"mode": mode_context, "adjust": adjust_context}
\ No newline at end of file |