智能体工具

#coding=utf-8
import autogen
import ktoolkits
from autogen import register_function

ktoolkits.api_key="XXXXXXXXXX"

llm_config = {#大模型配置#}

"""
创建用户代理
"""
user_proxy = autogen.UserProxyAgent(
    name="用户代理",
    #human_input_mode="NEVER",
    is_termination_msg=lambda x: x.get("content", "").rstrip().endswith("TERMINATE"),
    code_execution_config={
        "last_n_messages": 10,
        "work_dir": "code",
        "use_docker": False,
    }
)

"""
创建AI渗透测试专家
"""
engineer = autogen.AssistantAgent(
    name="渗透测试专家",
    system_message=ktoolkits.Agent.role("渗透测试专家"),
    llm_config=llm_config
)

#注册渗透测试工具
pentest_tool_list = [
    "dig",
    "nmap",
    "nuclei",
    "k_pentest_report",
]

agent_uuid = ktoolkits.Agent.create_agent(agent_name="渗透测试专家")
def k_register_tool(tool_name:str):
    tool = ktoolkits.FunctionCall.create_tool(tool_name=tool_name,agent_uuid=agent_uuid)
    register_function(
        f=tool.get("tool_func"),
        description=tool.get("tool_desc"),
        caller=engineer,
        executor=user_proxy,
    )

for tool_name in pentest_tool_list:
    try:
        k_register_tool(tool_name)
    except Exception as e:
        print(f"{tool_name}工具注册失败")


tasks=[
    """请对www.baidu.com进行渗透测试""",
]
user_proxy.initiate_chat(engineer,message=tasks[0])