An Automated Penetration Testing Method Based on Reinforcement Learning
By constructing a state space and action space for reinforcement learning and training the model in conjunction with vulnerability exploitation components, the entire penetration testing process is automated, solving the problems of limited penetration methods and unstable state transitions, and improving the efficiency and success rate of penetration testing.
Patent Information
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- BEIJING INST OF TECH
- Filing Date
- 2024-10-14
- Publication Date
- 2026-05-26
AI Technical Summary
Existing automated penetration testing methods have limited penetration techniques and a restricted action space, making them unable to handle multiple penetration targets. Furthermore, unstable state transitions result in low action coverage and unstable convergence, failing to meet the needs of large-scale and rapid vulnerability risk investigation.
We construct the state space, action space, and reward function of reinforcement learning, train and generate a reinforcement learning model using vulnerability exploitation components, and execute actions through policy selection to achieve full automation of the penetration testing process.
It improves the efficiency and success rate of penetration testing, solves the problems of strong reliance on manual intervention and fixed attack behaviors in traditional methods, and enhances the effectiveness of intelligent penetration testing in large-scale scenarios.
Smart Images
Figure CN119449373B_ABST
Abstract
Description
Technical Field
[0001] This invention relates to an automated penetration testing method based on reinforcement learning, belonging to the field of cyberspace security in computer and information science and technology. Background Technology
[0002] Cybersecurity incidents are frequent worldwide, with data breaches and hacker attacks occurring frequently, resulting in a significant increase in economic losses. Penetration testing is an assessment method that simulates attackers' intentions and behaviors in a real-world environment to evaluate the risk of a target system. It helps to discover potential vulnerabilities in the system, providing security managers with a basis for patching, and is of great significance in proactive cybersecurity defense.
[0003] Existing penetration testing methods primarily rely on manual execution, which is labor-intensive and inefficient, failing to meet the demands of large-scale, rapidly growing vulnerability risk assessment. Automated penetration testing technology, on the other hand, can automatically analyze the network environment of a target system, discovering and verifying potential vulnerabilities and weaknesses, significantly reducing the cost of penetration testing. Therefore, automation technology is increasingly being incorporated into penetration testing tasks. However, existing automated penetration testing techniques mainly employ a method of traversing exploit components and payloads, which is time-consuming and lacks scalability. To address these issues, current research introduces artificial intelligence methods, transforming the vulnerability exploitation process in penetration testing into a reinforcement learning action selection process. For example, Q-learning reinforcement learning models learn successful testing experiences of exploit components and record their states, selecting a payload to exploit available vulnerabilities in the next exploitable environment. However, these reinforcement learning algorithms have a limited action space and limited attack methods, making it difficult to select accurate actions in complex environments. Furthermore, some reinforcement learning algorithms use random state transitions, detaching them from action-based decisions, which does not meet the requirement of determining actions based on execution status in actual penetration testing. Moreover, the actual algorithm differs from the definition of reinforcement learning, leading to problems such as low action coverage and unstable convergence.
[0004] In summary, existing automated penetration testing methods have the following problems: 1. Limited penetration methods, with the action space confined to a small number of simple attack commands. In practical use, the fixed attack behaviors are difficult to cope with scenarios with a large number of penetration targets and multiple penetration methods; 2. Lack of state transition information, relying solely on random search for state transitions. The actual algorithm differs from the definition of reinforcement learning, resulting in low action coverage, excessively long training time, and unstable algorithm convergence. Therefore, this invention proposes an evolutionary method for automated penetration testing systems based on reinforcement learning. Summary of the Invention
[0005] The purpose of this invention is to address the problems in automated penetration testing, such as limited penetration methods, rigid attack behaviors that make it difficult to cope with scenarios with a large number of penetration targets and multiple penetration methods, and low action coverage and unstable convergence due to the lack of utilization of state transition information. The invention proposes an evolutionary method for automated penetration testing systems based on reinforcement learning.
[0006] The design principle of this invention is as follows: First, the state space, action space, and reward function of reinforcement learning are constructed using environment-related parameters to generate model training data; second, the reinforcement learning model is trained using the training data and vulnerability exploitation components; finally, based on the vulnerability exploitation component library, the reinforcement learning model recommends and optimizes vulnerability exploitation information to achieve automation of the entire penetration testing process.
[0007] The technical solution of the present invention is achieved through the following steps:
[0008] Step 1: Construct the state space, action space, and reward function of reinforcement learning using environment-related parameters to generate model training data.
[0009] Step 1.1: Extract the target host operating system, service name, service version, service running port, product name, session type, and permissions (7 dimensions) from the test environment as state space variables to construct the state space.
[0010] Step 1.2: Construct all exploit components in the exploit component library into an action space and record them using an index table.
[0011] Step 1.3: Construct a reward function based on information such as the running speed, threat level, and complexity of the exploit component.
[0012] Step 2: Use the training data and exploit components to train and generate a reinforcement learning model.
[0013] Step 2.1: Use information gathering tools to obtain the current test environment state from the test environment and initialize the model state space.
[0014] Step 2.2: The reinforcement learning agent outputs the Q-values of all possible actions based on the policy network.
[0015] Step 2.3: Based on the reinforcement learning strategy and the output Q-value, select actions from the action space, and use the corresponding exploit module to attempt to attack the test environment.
[0016] Step 2.4: Verify whether the attack on the test environment was successful. If the attack was successful, reset the test environment and initialize the state space; otherwise, update the current state.
[0017] Step 2.5: Calculate the reward value that can be obtained for this attack based on the reward function and the current state of the test environment.
[0018] Step 2.6: Update the reinforcement learning agent based on the reward value until the reinforcement learning agent model has the ability to select the optimal attack strategy.
[0019] Step 3: Based on the vulnerability exploitation component library, a reinforcement learning model is used to recommend the best vulnerability exploitation information.
[0020] Step 3.1: Parse the list of vulnerable assets and convert it into the corresponding vector for the model input.
[0021] Step 3.2: The model makes decisions on the input vectors, outputs recommended vectors, and summarizes them to obtain the vulnerability exploitation matrix.
[0022] Step 4: Automate the entire penetration testing process.
[0023] Step 4.1: parse the vulnerability exploitation matrix to obtain a recommended vulnerability exploitation order table for each vulnerable asset.
[0024] Step 4.2: Based on the vulnerability exploitation component library, query and obtain vulnerability exploitation payload information.
[0025] Step 4.3: Construct the code vulnerability exploit and execute the exploit.
[0026] Step 4.4: Based on the data analysis results, suggest sessions or return failure information, summarize vulnerability exploitation information, and assess risks.
[0027] Beneficial effects
[0028] Compared to traditional penetration testing methods, this invention adopts a fully automated penetration testing design. It uses multi-dimensional parameters to accurately construct the state space, action space, and reward function in the reinforcement learning process, selects and executes actions according to the strategy, and associates them with vulnerability exploitation components. This effectively solves the problems of strong reliance on manual intervention and fixed attack behaviors in penetration testing tools. It also adopts automated and intelligent risk assessment techniques to effectively improve the efficiency of risk assessment.
[0029] Compared to traditional penetration testing methods, this invention integrates value-based deep reinforcement learning and policy-based deep reinforcement learning, and adopts a reinforcement model-environment interaction method. This effectively alleviates the problems of low action coverage and unstable convergence in traditional automated penetration testing, improves the effect of intelligent penetration testing in large-scale scenarios, and greatly improves the efficiency of intelligent penetration testing. Attached Figure Description
[0030] Figure 1 This is a schematic diagram of the evolutionary method for automated penetration systems based on reinforcement learning, as proposed in this invention. Detailed Implementation
[0031] To better illustrate the purpose and advantages of the present invention, the implementation methods of the present invention will be further described in detail below with reference to examples.
[0032] The experimental data used publicly available vulnerability information to build the target range, including 207 vulnerabilities across 9 vulnerability types: remote overflow, cross-site scripting, cross-site request forgery, SQL injection, directory information leakage, file upload, file inclusion, file reading, and weak passwords. The training set contained 180 vulnerabilities, and the test set contained 27 vulnerabilities.
[0033] The experiment used precision and testing time to evaluate the results. Precision was calculated using the first five items from the recommended list (P@5), and the calculation method is shown in formula (1):
[0034]
[0035] Where N is the total number of vulnerabilities tested, and T is the number of vulnerabilities that can be correctly exploited in the top 5 recommended exploit components.
[0036] This experiment was conducted on a computer and a GPU server. The computer was configured with an Intel(R) Core(TM) i9-9900 eight-core processor with a CPU frequency of 3.10GHz, 32GB of RAM, and a 64-bit Win10 operating system. The server was configured with a GTX 1080Ti processor, 256GB of RAM, and a 64-bit Linux Ubuntu operating system.
[0037] The specific process is as follows:
[0038] Step 1: Construct the state space, action space, and reward function of reinforcement learning using environment-related parameters to generate model training data.
[0039] Step 1.1: Define a 7-dimensional state space for the parameters in dictionary form, specifically: {'OS':'','Port':'','Service':'','Version':'','Product':'','Session':'','Access':''}. Construct an index table for different parameters, and encode the strings in the dictionary into numeric vectors through the parameter-index table.
[0040] Step 1.2: The exploitation action is implemented through the exploitation components in the exploitation component library. There is a one-to-one correspondence between the component and the numerical sequence number, specifically in the form of {1: no_action; 2: exploit / aix / local / ibstat_path; ...}. These components together constitute the action space of the penetration testing abstract simulation environment. The dimension of the action space is consistent with the number of exploitation components in the exploitation component library.
[0041] Step 1.3, the reward function includes three parts: speed reward of the exploit component, threat reward, and complexity reward. The speed reward of the exploit component can be calculated by recording the running time t during the component's attack process. Combined with the average running time of the exploit component, the speed reward R1 = 1 / t can be constructed. The threat reward of the exploit component is evaluated by the information collection module based on the test environment status after the exploit component's attack. The threat reward R2 is shown in formula (2):
[0042]
[0043] The complexity modeling of the exploit component is calculated by calculating the code length l of the exploit component. The complexity reward R3 = 1 / l, and the final reward function R = λ1R1 + λ2R2 + λ3R3, where λ1, λ2, and λ3 are the reward weight coefficients of the three types of rewards, which can be adjusted according to the actual penetration testing needs.
[0044] Step 2: Use the training data and exploit components to train and generate a reinforcement learning model.
[0045] Step 2.1: Scan the test environment for live ports, collecting information such as the operating system name, operating system version, service name, service version, and corresponding port number. Use string matching for type determination, filtering out service types not supported by the vulnerability exploitation component library to obtain a list of services to be used for penetration testing. During the training phase, first determine the OS parameters based on the scanned operating system information. The Port, Service, Version, and Product parameters will be randomly selected from ports and services identified in the test environment during the information collection phase. Finally, set the Session and Access parameters to 0. This state is set to the initial state s0.
[0046] Step 2.2 employs a DQN network structure. The number of neurons in the input layer matches the number of parameters in the state space, and the number of neurons in the output layer matches the dimension of the action space. Several hidden layers are present in between. The input layer takes the current state s of the test environment as input. i The output layer outputs the Q-value of each action in the current state and provides a value assessment of each vulnerability exploitation component for the current network environment. The higher the Q-value, the higher the value of the vulnerability exploitation component. The attack action and its Q-value are the scheduling strategy given by the reinforcement learning agent model.
[0047] Step 2.3: Restore the correspondence between each exploit component and its Q-value according to the exploit component index table, and select the component with the highest Q-value as the attack action a to be used in this iteration. i The exploit component selected by the agent was used to attack the test environment.
[0048] Step 2.4: Record the time consumed during the attack process, determine the type of communication established after the attack and the permissions obtained, update the Session and Access parameters. If the attack is successful, reset the test environment and restore the current state to the initial state s0. If the attack fails, determine whether the attack has failed for an extended period. If the maximum number of attempts has not been reached, continue calculating the reward and updating the model; otherwise, randomly update the Port, Service, Version, and Product parameters until all ports and services have been tested. Since the operating system of the test environment is unique, the OS parameter will not be updated. Organize all state parameters to obtain the new state s. i+1 .
[0049] Step 2.5: Using the information collected in Step 1.3, calculate the speed reward and threat reward. Query the code length of the corresponding exploit component from the exploit component library and calculate the complexity reward. Sum the three rewards in a weighted manner to calculate the reward value r obtained for this round of attack. i .
[0050] Step 2.6: A fixed-size memory bank D is used to record experiences, and the experience is represented in the form of [s]. i ,a i ,r i ,s i+1 The DQN model consists of a target network Q′ and a Q network Q, both with identical structures. The Q network is used to calculate the Q-value for policy selection and is updated during each training iteration. The target network is used to calculate the Q-value for the next state, thereby reducing fluctuations in the target value. During each model update, a portion of data is randomly selected from the memory bank, and the current state s is... i Predict the current action a in the Q network. i The Q value will determine the next state s. i+1 The input is the maximum Q-value that can be obtained by predicting the next state in the target network. The loss function is: in The Q value calculated for the target network is Q(s). i ,a i Let θ be the Q-value calculated by the Q-network, γ be the decay coefficient, θ be the Q-network parameters, and θ′ be the target network parameters. The loss function is calculated using gradient descent, and the Q-network parameters θ are updated accordingly. When the number of training steps reaches a preset fixed value, the target network parameters θ′ are updated until the model reaches the predetermined number of training epochs.
[0051] Step 3: Based on the vulnerability exploitation component library, a reinforcement learning model is used to recommend the best vulnerability exploitation information.
[0052] Step 3.1: Extract the key attributes of each vulnerable asset, including asset name, asset IP, vulnerability name, and vulnerability type.
[0053] Step 3.2, outlier filtering and processing, to ensure the fault tolerance performance of the intelligent scheduling model. This involves detecting and identifying non-standard IP values and null vulnerability names in the feature vector list.
[0054] Step 3.3 involves feature encoding and normalization of the vulnerability feature vector list after outlier processing. The input is primarily text-based data. The feature information in the vulnerability feature vectors is encoded into text, and the vectors are normalized for processing by the decision algorithm model.
[0055] Step 3.4: For each 4-dimensional vector (containing asset name, asset IP, array of live vulnerability names, and array of live vulnerability types), generate a corresponding vulnerability exploitation planning decision vector (containing asset name, asset IP, and a recommended vulnerability exploitation order list, totaling 3 dimensions). The recommended vulnerability exploitation order list is ranked based on a comprehensive consideration of vulnerability severity, exploitation complexity, and other information.
[0056] Step 3.5: Concatenate and summarize the list of decision vectors generated in the model decision-making stage to obtain the result vector of vulnerability exploitation planning, where each dimension represents a specific decision or planning indicator.
[0057] Step 4: Automate the entire penetration testing process.
[0058] Step 4.1: In the exploitation phase, the input information and data are preprocessed. The recommended exploit vector matrix output by the scheduling engine is mapped and parsed to obtain the corresponding parsing for each asset, resulting in a vulnerability exploitation priority list. Based on the attack path and specific markers or rules in the vulnerability exploitation vector, a specific feasible exploitation path is determined.
[0059] Step 4.2: Based on the vulnerability exploitation path and plan, retrieve vulnerability exploitation components from the vulnerability exploitation component library based on the vulnerability name, and construct the corresponding exploitation actions. Utilize the vulnerability association information in the plan, combined with the characteristics of the attack target, to construct a series of runtime probe payloads and generate an effective exploitation action sequence.
[0060] Step 4.3: For each exploit, capture and analyze the corresponding application response information. Examine and verify the HTTP status code, response headers, response content, etc., to determine whether the target system has been exploited or affected by vulnerabilities.
[0061] Step 4.4: Summarize the vulnerability exploitation results and generate a system-level risk assessment report or output to assist in security decision-making and vulnerability remediation processes. Specific risk assessment results may include: (1) the session created after successful vulnerability exploitation, which is created after the vulnerability exploitation action is executed, such as obtaining a target shell and creating a session after executing the RCE exploitation component; (2) vulnerability threat level information; (3) vulnerability type information.
[0062] Test Results: This experiment used an automated penetration testing system evolution method based on reinforcement learning to perform penetration testing on 25 vulnerabilities of 9 vulnerability types. The results were compared with three traditional automated penetration testing methods. The results show that the present invention can significantly improve the success rate of penetration testing and is faster, as shown in Table 1.
[0063] Table 1. Experimental results of the evolutionary method for automated penetration systems based on reinforcement learning
[0064]
[0065]
[0066] The above detailed description further illustrates the purpose, technical solution, and beneficial effects of the invention. It should be understood that the above description is only a specific embodiment of the present invention and is not intended to limit the scope of protection of the present invention. Any modifications, equivalent substitutions, improvements, etc., made within the spirit and principles of the present invention should be included within the scope of protection of the present invention.
Claims
1. An automated penetration testing method based on reinforcement learning, characterized in that... The method includes the following steps: Step 1: Construct the state space, action space, and reward function for reinforcement learning using environment-related parameters to generate model training data; Step 1.1: Define a 7-dimensional state space using a dictionary format, specifically: {'OS':'', 'Port':'', 'Service':'', 'Version':'', 'Product':'', 'Session':'', 'Access':''}. The meanings of each parameter are as follows: target host operating system, service running port, service name, service version, product name, session type, and permissions. Construct an index table for different parameters, and encode the strings in the dictionary into numeric vectors using the parameter-index table. Step 1.2: The exploit action is implemented using exploit components from the exploit component library. A one-to-one correspondence is established between the component and its numerical identifier: {numerical identifier: component name}, specifically in the form {1: no_action; 2: exploit / aix / local / ibstat_path;} }, no_action means no exploit action is executed, exploit / aix / local / ibstat_path represents the identifier name of an exploit component in the exploit component library. These components together constitute the action space of the penetration testing abstract simulation environment. The dimension of the action space is consistent with the number of exploit components in the exploit component library. Step 1.3: The reward function comprises three parts: exploit component speed reward, threat level reward, and complexity reward. The exploit component speed reward is calculated by recording the runtime t during the component's attack process and combining it with the average runtime of the exploit component to construct the speed reward. The threat reward for exploit components is determined by the information gathering module, which assesses the state of the test environment after an exploit component attack. See formula (1): Complexity modeling of exploit components is achieved by calculating the code length *l* of the exploit component, and the complexity reward is calculated accordingly. Final reward function ,in , , These are the reward weighting coefficients for the three types of rewards, which can be adjusted according to the actual penetration testing needs; Step 2: Use the training data and exploit components to train and generate a reinforcement learning model; Step 2.1: Scan for live ports in the test environment, collect information such as operating system name, operating system version, service name, service version, and corresponding port number. Use string matching for type determination, filter out service types not supported by the vulnerability exploitation component library, and obtain a list of services to be used for penetration testing. During the training phase, first determine the OS parameters based on the scanned operating system information. The Port, Service, Version, and Product parameters will be randomly selected from ports and services identified in the test environment during the information collection phase. Finally, set the Session and Access parameters to 0; this state is set to the initial state. ; Step 2.2: A DQN network structure is adopted, with the number of neurons in the input layer matching the number of parameters in the state space, and the number of neurons in the output layer matching the dimension of the action space, with several hidden layers in between; the input layer inputs the current state of the test environment. The output layer outputs the Q-value of each action in the current state and provides a value assessment of each vulnerability exploitation component for the current network environment. The higher the Q-value, the higher the value of the vulnerability exploitation component. The attack action and its Q-value are the scheduling strategy given by the reinforcement learning agent model. Step 2.3: Restore the correspondence between each exploit component and its Q-value according to the exploit component index table, and select the component with the highest Q-value as the attack action to be used in this iteration. The test environment was attacked using the exploit component selected by the agent; Step 2.4: Record the time taken during the attack process, determine the type of communication established after the attack and the permissions obtained, update the Session and Access parameters, and if the attack is successful, reset the test environment and restore the current state to the initial state. If the attack fails, determine if the attack has been unsuccessful for an extended period. If the maximum number of attempts has not been reached, continue calculating the reward and updating the model; otherwise, randomly update the Port, Service, Version, and Product parameters until all ports and services have been tested. Since the operating system in the test environment is unique, the OS parameter will not be updated. All state parameters are then processed to obtain the new state. ; Step 2.5: Using the information collected in Step 1.3, calculate the speed reward and threat reward. Query the code length of the corresponding exploit component from the exploit component library and calculate the complexity reward. Sum the three rewards in a weighted manner to calculate the reward value obtained from this round of attack. ; Step 2.6, using a fixed-size memory bank Used for recording experience, the form of which is [ , , The DQN model contains a target network. And Q Network Both have identical structures. The Q-network is used to calculate the Q-value for policy selection and is updated during each training iteration. The target network is used to calculate the Q-value for the next state, thereby reducing fluctuations in the target value. During each model update, a portion of data is randomly selected from the memory bank, and the current state is... Predict the current action in the Q network. The Q value will determine the next state. The input target network is the maximum Q-value that can be obtained by predicting the next state; the loss function is: ,in The Q value calculated for the target network. The Q-value is calculated using the Q-network. The attenuation coefficient is... For Q network parameters, For target network parameters; The loss function is calculated using gradient descent, and the Q-network parameters are updated. When the number of training steps reaches a preset fixed value, update the target network parameters. Continue until the model reaches the predetermined number of training rounds; Step 3: Based on the vulnerability exploitation component library, a reinforcement learning model is used to recommend the best vulnerability exploitation information; Step 3.1: Extract the key attributes of each vulnerable asset, including asset name, asset IP, vulnerability name, and vulnerability type; Step 3.2, outlier filtering and processing to ensure the fault tolerance performance of the intelligent scheduling model; detect and identify non-standard IP values and null vulnerability names in the feature vector list; Step 3.3: Perform feature encoding and normalization on the vulnerability feature vector list after outlier processing; the input is text data, the feature information in the vulnerability feature vector is text encoded, and the vector is normalized for processing by the decision algorithm model; Step 3.4: For each 4-dimensional vector, the four dimensions are asset name, asset IP, array of live vulnerability names, and array of live vulnerability types, respectively. Generate a corresponding 3-dimensional vulnerability exploitation planning decision vector, with the three dimensions being asset name, asset IP, and a recommended vulnerability exploitation order list. The recommended vulnerability exploitation order list is ranked based on a comprehensive consideration of vulnerability severity and exploitation complexity information. Step 3.5: Concatenate and summarize the list of decision vectors generated in the model decision-making stage to obtain the result vector of vulnerability exploitation planning, where each dimension represents a specific decision or planning indicator; Step 4: Automate the entire penetration testing process; Step 4.1: In the exploitation phase, the input information and data are preprocessed, the recommended exploit vector matrix output by the scheduling engine is mapped and parsed, and the corresponding parsing for each asset is obtained to obtain its vulnerability exploitation priority list; combined with the attack path, the specific feasible vulnerability exploitation path is determined according to the specific markers or rules in the vulnerability exploitation vector. Step 4.2: Based on the vulnerability exploitation path and plan, retrieve vulnerability exploitation components from the vulnerability exploitation component library based on the vulnerability name, and construct corresponding exploitation actions; utilize the vulnerability association information in the plan, combine it with the characteristics of the attack target, construct a series of runtime probe payloads, and generate an effective exploitation action sequence. Step 4.3: For each exploit action, capture and analyze the corresponding application response information; check and verify the HTTP status code, response header information, and response content to determine whether the target system has been exploited or affected by vulnerabilities. Step 4.4, summarize the vulnerability exploitation results and generate a system-level risk assessment report or output to assist in security decision-making and vulnerability remediation process; specific risk assessment results include: (1) Session created after successful vulnerability exploitation: the session is created after the vulnerability exploitation action is executed, and the target shell is obtained after the RCE exploitation component is executed to create the session; (2) vulnerability threat level information; (3) vulnerability type information.