PenHeal: A Two-Stage LLM Framework for Automated Pentesting and Optimal Remediation

TL;DR

PenHeal presents a two‑stage LLM framework that automates both vulnerability discovery and remediation. By coupling a Pentest Module with a Remediation Module and guiding reasoning with Counterfactual Prompting and an external Instructor knowledge base, it achieves higher vulnerability coverage and more effective, cost‑efficient remediation than baselines on a Metasploitable2 benchmark. The work introduces a concrete evaluation framework with CVSS‑informed remediation scoring and a Group Knapsack optimization to select budget‑constrained actions. Results show a 31% gain in detection coverage, a 32% improvement in remediation effectiveness, and a 46% reduction in remediation costs, signaling substantial practical potential for autonomous cybersecurity workflows. Limitations include tool dependence, multi‑host scalability, and the need for end‑to‑end deployment considerations, which the authors propose addressing in future research.

Abstract

Recent advances in Large Language Models (LLMs) have shown significant potential in enhancing cybersecurity defenses against sophisticated threats. LLM-based penetration testing is an essential step in automating system security evaluations by identifying vulnerabilities. Remediation, the subsequent crucial step, addresses these discovered vulnerabilities. Since details about vulnerabilities, exploitation methods, and software versions offer crucial insights into system weaknesses, integrating penetration testing with vulnerability remediation into a cohesive system has become both intuitive and necessary. This paper introduces PenHeal, a two-stage LLM-based framework designed to autonomously identify and mitigate security vulnerabilities. The framework integrates two LLM-enabled components: the Pentest Module, which detects multiple vulnerabilities within a system, and the Remediation Module, which recommends optimal remediation strategies. The integration is facilitated through Counterfactual Prompting and an Instructor module that guides the LLMs using external knowledge to explore multiple potential attack paths effectively. Our experimental results demonstrate that PenHeal not only automates the identification and remediation of vulnerabilities but also significantly improves vulnerability coverage by 31%, increases the effectiveness of remediation strategies by 32%, and reduces the associated costs by 46% compared to baseline models. These outcomes highlight the transformative potential of LLMs in reshaping cybersecurity practices, offering an innovative solution to defend against cyber threats.

Figures (7)

• Figure 1: PenHeal Abstract Architecture: the Pentest Module is responsible for detecting vulnerabilities, and the Remediation Module suggests remediation strategies. Initially, the user inputs the target system's IP address, prompting the Pentest Module to conduct penetration testing and log vulnerabilities. Subsequently, the Remediation Module generates and presents remediation recommendations to the user based on these findings.
• Figure 2: The Pentest Module begins with the user providing the target system's IP address. The Planner then formulates the attack strategy, directing the Executor, which consults the Instructor to generate commands aimed at exploiting system vulnerabilities. These commands are analyzed by the Summarizer, and Counterfactual Prompting aids the Planner in refining the attack strategy by considering alternative vectors in response to new vulnerabilities. The blocks with OpenAI icons represent that they are OpenAI LLMs.
• Figure 3: Remediation Module Architecture: Upon receiving vulnerability information from the Pentest Module, additional attributes such as CVSS score are retrieved for the vulnerabilities publicly documented. For the others, the Estimator generates descriptors to assess their severity and attributes based on the available data. Then, based on the retrieved attributes, the Advisor proposes remediation strategies, which the Evaluator will determine their qualities. Finally, an optimal list of recommendations under the user-defined budget constraint is presented to the user.
• Figure 4: Comparison of Scores Achieved by Different Models
• Figure 5: Ablation Study of PenHeal