一、项目介绍1.1 仓库地址GIT URL: https://github.com/Joe-zsc/April-AE.git1.2 项目说明用强化学习动作嵌入SBERT训练“自治渗透测试”智能体在大规模动作空间中学习高效的扫描/利用策略对应论文 APRIL。我们应该看到的“结果”训练过程曲线与指标奖励、成功率、步数、覆盖率等TensorBoard 中查看。学到的渗透策略给定环境智能体能按更少步骤、更高成功率达成目标成功利用/拿到 flag。最优动作序列渗透路径每回合的具体操作序列可用于复现与报告。训练后的结果“在哪里”训练日志与曲线runs/TensorBoard 可视化。动作向量缓存actions/Action-/Embedding-.npy只是加速用的缓存不是训练产物。预训练文本嵌入模型SBERTNLP_Module/Embedding_models/…onnx 是预训练模型导出不是你的RL训练结果。强化学习策略权重需要你显式保存April-AE-actor.pt / critic_.pt 以及 statenorm.pt保存到你指定目录例如 log/checkpoints/。二、Docker镜像构建2.1 准备工作requirements.txt调整原来requirements中torch的版本是。annoy1.17.3# numpy1.19.5 这个版本要升级numpy1.26.4#pandas2.2.3 这个版本要升级pandas2.3.1scikit_learn1.4.1.post1sentence_transformers2.5.1torch2.5.1tqdm4.66.2Linux安装软件包(windows需要单独安装git客户端)sudoapt-getupdatesudoapt-getinstall-ygitgit-lfsgitlfsinstall下载模型git拉取(linux和windows)模型地址在https://huggingface.co/sentence-transformers/all-MiniLM-L12-v2/tree/main此处但由于网络不稳定只能下载到个人服务器中专一次。如果网络允许,可以直接git拉取:gitclone https://huggingface.co/sentence-transformers/all-MiniLM-L12-v2\./NLP_Module/Embedding_models/all-MiniLM-L12-v2linuxgit不能拉取的情况cdNLP_Module/Embedding_modelswgethttp://118.25.230.32:9000/all-MiniLM-L12-v2.tar.gztar-zxvfall-MiniLM-L12-v2.tar.gz Windowsgit不能拉取的情况Windows直接通过浏览器下载http://118.25.230.32:9000/all-MiniLM-L12-v2.tar.gz该模型文件到April-AE/ 目录解压即可。修改配置文件config.ini修改NLP_Module\Embedding_models为NLP_Module/Embedding_models。[common] log_path log actions_file Action-1000 project_name April-AE [Embedding] embedding_modelsNLP_Module/Embedding_models sbert_model all-MiniLM-L12-v22.2 制作Dockerfile# 基础镜像 FROM python:3.11-slim # 防止 Python 缓冲标准输出/错误 ENV PYTHONUNBUFFERED1 \ PIP_NO_CACHE_DIR1 \ PYTHONDONTWRITEBYTECODE1 # 切换 APT 源为国内镜像并安装系统依赖兼容新旧 Debian带重试与 fix-missing ARG DEBIAN_MIRRORhttps://mirrors.tuna.tsinghua.edu.cn/debian ARG DEBIAN_SECURITY_MIRRORhttps://mirrors.tuna.tsinghua.edu.cn/debian-security RUN set -eux; \ if [ -f /etc/apt/sources.list ]; then \ sed -i s|http://deb.debian.org/debian|${DEBIAN_MIRROR}|g; s|http://security.debian.org/debian-security|${DEBIAN_SECURITY_MIRROR}|g /etc/apt/sources.list; \ elif [ -f /etc/apt/sources.list.d/debian.sources ]; then \ sed -i s|URIs: http://deb.debian.org/debian|URIs: ${DEBIAN_MIRROR}|g /etc/apt/sources.list.d/debian.sources || true; \ sed -i s|URIs: http://security.debian.org/debian-security|URIs: ${DEBIAN_SECURITY_MIRROR}|g /etc/apt/sources.list.d/debian.sources || true; \ else \ printf deb %s stable main contrib non-free non-free-firmware\n ${DEBIAN_MIRROR} /etc/apt/sources.list; \ printf deb %s stable-updates main contrib non-free non-free-firmware\n ${DEBIAN_MIRROR} /etc/apt/sources.list; \ printf deb %s stable-security main contrib non-free non-free-firmware\n ${DEBIAN_SECURITY_MIRROR} /etc/apt/sources.list; \ fi; \ apt-get -o Acquire::Retries5 update; \ apt-get install -y --no-install-recommends \ build-essential \ git \ curl \ ca-certificates \ || (apt-get -o Acquire::Retries5 update --fix-missing apt-get install -y --no-install-recommends build-essential git curl ca-certificates); \ rm -rf /var/lib/apt/lists/* # 工作目录 WORKDIR /app # 仅先拷贝 requirements 以利用分层缓存 COPY requirements.txt /app/requirements.txt # 使用官方 pip 命令安装 PyTorchCPU 版本 # 注意torchaudio 2.5.1 可能在某些索引中不可用使用更宽松的版本约束 #RUN python -m pip install --upgrade pip \ # pip install --no-cache-dir torch2.5.1 torchvision0.20.1 torchaudio2.5.1 #--index-url https://download.pytorch.org/whl/cpu # 使用官方 pip 命令安装 PyTorchGPU 版本 RUN python -m pip install --upgrade pip \ pip install --no-cache-dir torch2.5.1 torchvision0.20.1 torchaudio2.5.1 --index-url https://download.pytorch.org/whl/cu118 # 安装其他 Python 依赖排除 torch 相关包 RUN awk !/^torch/ requirements.txt /tmp/requirements.no_torch.txt \ pip install --no-cache-dir -r /tmp/requirements.no_torch.txt \ pip install --no-cache-dir tensorboard # 拷贝项目其余文件 COPY . /app # 为模型、日志与运行结果声明挂载点 VOLUME [/app/NLP_Module/Embedding_models, /app/log, /app/runs] # 暴露 README 中的 TensorBoard 端口 EXPOSE 6666 # 默认可配置参数 ARG ENV_FILEsingle/env-CVE-2018-11776.json ARG AGENTApril-AE ARG GPU-1 # 支持在运行时通过环境变量覆盖 ENV ENV_FILE${ENV_FILE} \ AGENT${AGENT} \ GPU${GPU} # 启动命令按 README 运行训练在 Linux 中使用 POSIX 路径 CMD [/bin/sh, -lc, python April.py --env_file \$ENV_FILE\ --agent \$AGENT\ --gpu \$GPU\]2.3 构建镜像注意需要把/home/kali换位实际的地址。linux版本cd/home/kali/April-AEdockerbuild-tapril-ae:latest.# 或者# docker build -t april-ae:latest /home/kali/April-AEWindows版本cd D:/workspace/April-AE docker build -t april-ae:latest . # 或者 # docker build -t april-ae:latest D:/workspace/April-AE2.3 运行模型Linux版本需要把/home/kali/改为实际的地址。dockerrun--rm-it\-v/home/kali/April-AE/NLP_Module/Embedding_models:/app/NLP_Module/Embedding_models\-v/home/kali/April-AE/log:/app/log\-v/home/kali/April-AE/runs:/app/runs\-eAGENTApril-AE\april-ae:latestWindows版本docker run --rm -it -v D:/workspace/April-AE/NLP_Module/Embedding_models:/app/NLP_Module/Embedding_models -v D:/workspace/April-AE/log:/app/log -v D:/workspace/April-AE/runs:/app/runs -e AGENTApril-AE april-ae:latest其他命令:指定不同环境文件或 agent保持 POSIX 路径dockerrun--rm-it\-vD:/workspace/April-AE/NLP_Module/Embedding_models:/app/NLP_Module/Embedding_models\-vD:/workspace/April-AE/log:/app/log\-vD:/workspace/April-AE/runs:/app/runs\-eENV_FILEsingle/env-CVE-2019-9193.json\-eAGENTApril-AE\-eGPU-1\--nameapril-ae\april-ae:latest启动 tensorboard映射 6666 端口实时可视化单独起一个 TensorBoard 容器训练容器退出也能看历史访问一定要改为http://127.0.0.1:6006。dockerrun--rm-it\-v/mnt/d/workspace/April-AE/runs:/app/runs\-p6006:6666\--entrypointbash\april-ae:latest-lctensorboard --logdir runs --host 0.0.0.0 --port 6666dockerrun--rm-it-vd:\workspace\April-AE\runs:/app/runs-p6006:6666--entrypointbashapril-ae:latest-lctensorboard --logdir runs --host 0.0.0.0 --port 6666GPU 训练推荐分离训练与可视化在wsl中执行下面命令使用GPU加速。dockerrun--rm-it--gpusall\-v/mnt/d/workspace/April-AE/NLP_Module/Embedding_models:/app/NLP_Module/Embedding_models\-v/mnt/d/workspace/April-AE/log:/app/log\-v/mnt/d/workspace/April-AE/runs:/app/runs\-eAGENTApril-AE\-eGPU0\--nameapril-ae\april-ae:latest如果希望训练完容器不立即退出不推荐一般没必要前台常驻在容器里追加一个阻塞命令docker run --gpus all -d \ -v /mnt/d/workspace/April-AE/NLP_Module/Embedding_models:/app/NLP_Module/Embedding_models \ -v /mnt/d/workspace/April-AE/log:/app/log \ -v /mnt/d/workspace/April-AE/runs:/app/runs \ -e AGENTApril-AE -e GPU0 \ --name april-ae-train-bg \ april-ae:latest \ bash -lc python April.py --env_file \$ENV_FILE\ --agent \$AGENT\ --gpu \$GPU\; tail -f /dev/null - 或者把 TensorBoard和训练同容器内并行跑训练结束 tensorboard 仍在 docker run --gpus all -d \ -v /mnt/d/workspace/April-AE/NLP_Module/Embedding_models:/app/NLP_Module/Embedding_models \ -v /mnt/d/workspace/April-AE/log:/app/log \ -v /mnt/d/workspace/April-AE/runs:/app/runs \ -p 6666:6666 \ -e AGENTApril-AE -e GPU0 \ --name april-ae-train-and-tb \ april-ae:latest \ bash -lc python April.py --env_file \$ENV_FILE\ --agent \$AGENT\ --gpu \$GPU\ tensorboard --logdir runs --host 0.0.0.0 --port 6666; wait三、使用方式defload(self,path):ifself.use_state_norm:self.state_normNormalization(shapeStateEncoder.state_space# shapeStateEncoder.state_space, finetuneTrue)mean_checkpointpath/fstate_norm_mean.ptstd_checkpointpath/fstate_norm_std.ptmeantorch.load(mean_checkpoint)stdtorch.load(std_checkpoint)self.state_norm.running_ms.meanmean self.state_norm.running_ms.stdstd self.state_norm.running_ms.Sstd*std self.Policy.load(path)self.is_loaded_agentTruedockerrun--rm-it-vD:\workspace\April-AE\log:/app/log-vD:\workspace\April-AE\NLP_Module\Embedding_models:/app/NLP_Module/Embedding_models-eCKPT_TITLEApril-AE-Sep24_15-57-50-single/env-CVE-2019-9193-eEVAL_ENVsingle/env-CVE-2019-9193.json--entrypointbashapril-ae:latest-lcpython evaluate.py四、总结训练april-ae提供很多种场景scenarios在scenarios\single目录有很多配置项我们这里是拿env-CVE-2019-9193.json这个场景为例训练的模型。如果需要训练其他场景修改下面命令中ENV_FILE的值即可dockerrun--rm-it--gpusall-vD:/workspace/April-AE/NLP_Module/Embedding_models:/app/NLP_Module/Embedding_models-vD:/workspace/April-AE/log:/app/log-vD:/workspace/April-AE/runs:/app/runs-eENV_FILEsingle/env-CVE-2019-9193.json-eAGENTApril-AE-eGPU0--nameapril-ae april-ae:latest模型训练结束之后在log/checkpoints目录下会生成模型结果比如/April-AE-Sep24_15-57-50-single/April-AE-actor.pt /April-AE-Sep24_15-57-50-single/env-CVE-2019-9193/April-AE-critic_1.pt /April-AE-Sep24_15-57-50-single/env-CVE-2019-9193/April-AE-critic_2.pt /April-AE-Sep24_15-57-50-single/env-CVE-2019-9193/state_norm_mean.pt /April-AE-Sep24_15-57-50-single/env-CVE-2019-9193/state_norm_std.pt这些结果我们会在模型评估中用到。评估模型评估阶段我们会调用evaluate.py对模型结果做评估验证传入的参数中我们需要修改CKPT_TITLE的值为上面的/April-AE-Sep24_15-57-50-single/env-CVE-2019-9193文件夹的相对路径。EVAL_ENV的值为场景值即与模型训练阶段的传值保持一致ingle/env-CVE-2019-9193.json。dockerrun--rm-it-vD:\workspace\April-AE\log:/app/log-vD:\workspace\April-AE\NLP_Module\Embedding_models:/app/NLP_Module/Embedding_models-eCKPT_TITLEApril-AE-Sep24_15-57-50-single/env-CVE-2019-9193-eEVAL_ENVsingle/env-CVE-2019-9193.json--entrypointbashapril-ae:latest-lcpython evaluate.py运行后我们得到结果--Running April-AE agent-- eval_return915, steps6, success_rate1.0