Automatic penetration testing method based on reinforcement learning

Abstract

The invention relates to an automatic penetration testing method based on reinforcement learning, and belongs to the field of computer and information science and technology network space security. The method comprises the following steps: firstly, constructing a state space, an action space and a reward function of reinforcement learning by utilizing parameters related to the environment to generate model training data; secondly, applying training data and vulnerability utilization component training to generate a reinforcement learning model; and finally, based on the vulnerability utilization component library, preferential vulnerability utilization information is recommended through a reinforcement learning model, and automation of the whole process of penetration testing is realized. The artificial intelligence technology is adopted for solving the problems of manual dependence of penetration testing, application solidification of the vulnerability utilization component and the like, the recommendation accuracy of the vulnerability utilization component is effectively improved, and the efficiency of penetration testing is remarkably improved.

Description

Technical Field

[0001] The present invention relates to an automated penetration testing method based on reinforcement learning, and belongs to the field of computer and information science and technology cyberspace security. Background Art

[0002] Cybersecurity incidents occur frequently around the world, such as data leaks and hacker attacks, and the economic losses caused by them have also increased significantly. Penetration testing is an assessment method that simulates the attacker's attack intentions and behaviors in a real environment to conduct risk assessment on the target system. It helps to discover potential vulnerabilities in the system and provide security managers with a basis for patching. It is of great significance in the active defense of cyberspace.

[0003] Existing penetration testing methods mainly rely on manual implementation, which requires a lot of manpower and is inefficient, and cannot meet the needs of large-scale and rapidly growing vulnerability risk investigation. Automated penetration testing technology can automatically analyze the network environment where the target system is located, discover and verify potential vulnerabilities and vulnerabilities of the target system, greatly reducing the cost of penetration testing. Therefore, automation technology has gradually been introduced into penetration testing tasks. However, existing automated penetration testing technologies mainly use the method of traversing vulnerability exploitation components and payloads, which is time-consuming and lacks scalability. In response to the above problems, existing research introduces artificial intelligence methods to convert the vulnerability exploitation process of penetration testing into an action selection process of reinforcement learning. For example, the Q-learning reinforcement learning model is used to learn the successful test experience of the exploitation component and record the state, and select the payload to exploit the available vulnerability in the next exploitable environment. However, this type of reinforcement learning algorithm has a small action space and a single attack method, and cannot select accurate actions in complex environments. In addition, some reinforcement learning algorithms transfer states in a random manner, which is separated from the connection with actions. This does not meet the requirements of determining the execution action based on the execution situation in the actual penetration testing process. In addition, there are differences between the actual algorithm and the definition of reinforcement learning, resulting in problems such as low action coverage and unstable convergence.

[0004] In summary, the existing automated penetration testing methods have the following problems: 1. The penetration means are limited, the action space is limited to a small number of simple attack commands, and the solidified attack behaviors are used in the actual use process, which is difficult to deal with the scene with a large number of penetration targets and multiple penetration means; 2. No state transition information is used, only random search transfer state is used, and there is a difference between the actual algorithm and the definition of reinforcement learning, resulting in low action coverage, long training time, unstable algorithm convergence, etc. Therefore, the present invention proposes an automated penetration system evolution method based on reinforcement learning. Summary of the invention

[0005] The purpose of the present invention is to propose an automated penetration system evolution method based on reinforcement learning to address the problems in automated penetration testing, such as limited penetration means, rigid attack behaviors, making it difficult to cope with scenarios with a large number of penetration targets and multiple penetration means, and the lack of state transition information resulting in low action coverage and unstable convergence.

[0006] The design principle of the present invention is: first, the state space, action space and reward function of reinforcement learning are constructed using parameters related to the environment to generate model training data; secondly, the training data and vulnerability exploitation component training are used to generate a reinforcement learning model; finally, based on the vulnerability exploitation component library, the reinforcement learning model is used to recommend preferred vulnerability exploitation information to achieve automation of the entire penetration test process.

[0007] The technical solution of the present invention is achieved through the following steps:

[0008] Step 1: Use environment-related parameters to construct the state space, action space, and reward function of reinforcement learning 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 permission 7-dimensional parameters in the test environment as state space variables to construct the state space.

[0010] Step 1.2, construct all vulnerability exploit components in the vulnerability exploit component library into an action space and record them using an index table.

[0011] Step 1.3, construct a reward function based on the running speed, threat level, complexity and other information of the vulnerability exploit component.

[0012] Step 2: Apply the training data and vulnerability exploitation components to generate a reinforcement learning model.

[0013] Step 2.1, use the information collection tool to obtain the current test environment state from the test environment and initialize the model state space.

[0014] In step 2.2, the reinforcement learning agent outputs the Q-values ​​of all possible actions according to the policy network.

[0015] Step 2.3, based on the reinforcement learning strategy and according to the output Q value, an action is selected from the action space, and the vulnerability exploitation module used for the action is used to attempt to attack the test environment.

[0016] Step 2.4, verify whether the attack on the test environment is successful. If the attack is 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 from this attack based on the reward function and the current state of the test environment.

[0018] Step 2.6, update the reinforcement learning agent according to the reward value until the reinforcement learning agent model has the ability to select the optimal attack strategy.

[0019] Step 3: Recommend optimal vulnerability exploit information based on the vulnerability exploit component library through a reinforcement learning model.

[0020] Step 3.1, parse the vulnerable asset list and convert it into the corresponding vector of model input.

[0021] In step 3.2, the model makes decisions on the input vector, outputs the recommended vector, and summarizes it to obtain the vulnerability exploitation matrix.

[0022] Step 4: Automate the entire penetration testing process.

[0023] Step 4.1, parse the vulnerability exploitation matrix and obtain the recommended vulnerability exploitation sequence table for each vulnerable asset.

[0024] Step 4.2: Based on the vulnerability exploit component library, query and obtain vulnerability exploit payload information.

[0025] Step 4.3, construct the encoded vulnerability exploit action and execute the exploit action.

[0026] Step 4.4, based on the data analysis results, recommend sessions or return failure information, summarize vulnerability exploitation information, and assess risks.

[0027] Beneficial Effects

[0028] Compared with traditional penetration testing methods, the present invention adopts a fully automated penetration testing design, uses multi-dimensional parameters to accurately construct the state space, action space and reward function in the reinforcement learning process, selects execution actions according to strategies and associates vulnerability exploitation components, effectively solves the problems of strong manual dependence and solidified attack behaviors in penetration tools, and adopts automated and intelligent risk assessment technical means to effectively improve the work efficiency of risk assessment.

[0029] Compared with traditional penetration testing methods, the present invention combines value-based deep reinforcement learning with policy-based deep reinforcement learning, and adopts a reinforcement model environment interaction method, which effectively alleviates the problems of low action coverage and unstable convergence of traditional automated penetration testing, improves the effect of intelligent penetration testing in large-scale scenarios, and greatly improves the efficiency of intelligent penetration. BRIEF DESCRIPTION OF THE DRAWINGS

[0030] Figure 1 This is a schematic diagram of the automated infiltration system evolution method based on reinforcement learning of the present invention. DETAILED DESCRIPTION

[0031] In order to better illustrate the purpose and advantages of the present invention, the implementation method of the present invention is further described in detail below with reference to examples.

[0032] The experimental data uses public vulnerability information to build a target range, including 9 types of vulnerabilities, including remote overflow, cross-site scripting, cross-site request forgery, SQL injection, directory information leakage, file upload, file inclusion, file reading, weak password, etc., with a total of 207 vulnerabilities. There are 180 vulnerabilities in the training set and 27 vulnerabilities in the test set.

[0033] The experiment uses precision (Precision) and test time (Time) to evaluate the results. The precision is calculated by taking the first 5 in the recommendation list (P@5). The calculation method is shown in formula (1):

[0034]

[0035] Among them, N is the total number of tested vulnerabilities, 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 configuration is: Intel(R) Core(TM) i9-9900 eight-core, CPU frequency 3.10GHz, 32GB RAM, operating system is 64-bit Win10; the server configuration is: GTX 1080Ti, RAM 256G, operating system is Linux Ubuntu 64-bit.

[0037] The specific process is:

[0038] Step 1: Use environment-related parameters to construct the state space, action space, and reward function of reinforcement learning to generate model training data.

[0039] Step 1.1, define the state space of 7-dimensional parameters in dictionary form, the specific form is: {'OS':'','Port':'','Service':'','Version':'','Product':'','Session':'','Access':''}, construct an index table for different parameters, and encode the strings in the dictionary into digital vectors through the parameter-index table.

[0040] In step 1.2, the vulnerability exploit action is implemented through the vulnerability exploit component in the vulnerability exploit component library. There is a one-to-one correspondence between the component and the digital serial number. The specific form is {1: no_action; 2: exploit / aix / local / ibstat_path; ...}. These components together constitute the action space of the abstract simulation environment of the penetration test. The dimension of the action space is consistent with the number of vulnerability exploit components in the vulnerability exploit component library.

[0041] Step 1.3, the reward function includes three parts of rewards: vulnerability exploitation component speed reward, threat reward, and complexity reward. The speed reward of the vulnerability exploitation component can be calculated by recording the running time t of the component during the attack process. Combined with the average running time of the vulnerability exploitation component, the speed reward R1 = 1 / t can be constructed. The threat reward of the vulnerability exploitation component is evaluated by collecting information on the test environment status after the vulnerability exploitation component attacks through the information collection module. The threat reward R2 is shown in formula (2):

[0042]

[0043] The complexity modeling of the vulnerability 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 test requirements.

[0044] Step 2: Apply the training data and vulnerability exploitation components to generate a reinforcement learning model.

[0045] Step 2.1, scan the survival ports in the test environment, collect the operating system name, operating system version, service name, service version, port number corresponding to the service, etc. of the test environment, use string matching to determine the type, filter the service types that are not supported in the vulnerability exploit component library, and obtain the service list to be penetrated. In the training phase, first determine the OS parameters based on the operating system information obtained by the scan. The Port, Service, Version, and Product parameters will randomly select the ports and services in the test environment identified in 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, using the DQN network structure, the number of neurons in the input layer is consistent with the number of parameters in the state space, the number of neurons in the output layer is consistent with the dimension of the action space, and there are several hidden layers in the middle. The input layer inputs the current state of the test environment s i The output layer outputs the Q value of each action in the current state, and gives the value evaluation of each vulnerability exploitation component according to 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 vulnerability exploit component and its Q value according to the vulnerability exploit component index table, and select the component with the highest Q value as the attack action a used in this iteration. i , using the agent-selected exploit component to attack the test environment.

[0048] Step 2.4, record the time spent in the attack process, determine the type of communication established and the permissions obtained after the attack, update the Session and Access parameters, and reset the test environment if the attack is successful, and restore the current state to the initial state s0. If the attack fails, determine whether the attack has not been successful for a long time. If the maximum number of attempts has not been reached, continue to calculate the reward update 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 parameters will not be updated. Arrange all state parameters to obtain the new state s i+1 .

[0049] Step 2.5, use the information collected in step 1.3 to calculate the speed reward and threat reward, query the code length of the corresponding vulnerability exploit component from the vulnerability exploit component library, calculate the complexity reward, and sum the three rewards to calculate the reward value r obtained by this round of attack actions. i .

[0050] Step 2.6, a fixed-size memory bank D is used to record experience. The experience is represented as [s i ,a i ,r i ,s i+1 ]. In the DQN model, there are target networks Q′ and Q networks Q, which have exactly the same structure. The Q network is used to calculate the Q value of the strategy selection. The Q network is updated during each training. The target network is used to calculate the Q value of the next state, thereby reducing the fluctuation of the target value. In each round of model update, some data is randomly selected from the memory bank and the current state s is updated. i Input Q network to predict the current action a i The Q value of the next state s i+1 Enter 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, Q(s i ,a i ; θ) is the Q value calculated by the Q network, γ is the attenuation coefficient, θ is the Q network parameter, and θ′ is the target network parameter. The loss function is calculated using the gradient descent method to update the Q network parameter θ. When the number of training steps reaches a preset fixed value, the target network parameter θ′ is updated until the model reaches the predetermined number of training rounds.

[0051] Step 3: Recommend optimal vulnerability exploit information based on the vulnerability exploit component library through a reinforcement learning model.

[0052] Step 3.1, extract the key attributes of each asset with a vulnerability, including asset name, asset IP, vulnerability name, and vulnerability type.

[0053] Step 3.2: Outlier screening and processing to ensure the fault tolerance of the intelligent scheduling model. Detect and identify the non-standard IP values ​​and empty vulnerability names in the feature vector list.

[0054] Step 3.3, feature encoding and normalization of the vulnerability feature vector list after outlier processing. The input is mainly text type data, the feature information in the vulnerability feature vector is text encoded, and the vector is normalized for processing by the decision algorithm model.

[0055] Step 3.4, for each 4-dimensional vector (including asset name, asset IP, survivable vulnerability name array, survivable vulnerability type array), generate a corresponding vulnerability exploitation planning decision vector (including asset name, asset IP, recommended vulnerability exploitation sequence list, a total of 3 dimensions). Among them, the recommended vulnerability exploitation sequence list is ranked based on comprehensive consideration of vulnerability severity, utilization complexity and other information.

[0056] In step 3.5, the decision vector list generated in the model decision phase is concatenated and summarized 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, preprocess the input information and data, map and parse the recommended exploit vector matrix output by the scheduling engine, sort out the corresponding analysis for each asset, and obtain its vulnerability exploitation priority list. Combined with the attack path, determine the specific feasible vulnerability exploitation path according to the specific mark or rule in the vulnerability exploitation vector.

[0059] Step 4.2, according to the vulnerability exploitation path and plan, retrieve the vulnerability exploitation component based on the vulnerability name in the vulnerability exploitation component library, and construct the corresponding exploitation action. Utilize the vulnerability association information in the plan, combine the characteristics of the attack target, build a series of runtime detection payloads, and generate an effective exploitation action sequence.

[0060] 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, response content, etc. to determine whether the target system is affected by the exploit or vulnerability.

[0061] Step 4.4, summarize the vulnerability exploitation results and generate a system-level risk assessment report or output to assist security decision-making and vulnerability remediation. Specific risk assessment results may include: (1) Sessions created after successful vulnerability exploitation, which are created after the vulnerability exploitation action is executed, such as obtaining the target Shell to create a Session after executing the RCE exploit component; (2) Vulnerability threat level information; (3) Vulnerability type information.

[0062] Test results: This experiment used an automated penetration system evolution method based on reinforcement learning to conduct penetration tests on 25 vulnerabilities including 9 vulnerability types, and compared them with three traditional automated penetration testing methods. The results showed that the present invention can greatly improve the success rate of penetration testing and is faster, see Table 1.

[0063] Table 1. Experimental results of automated penetration system evolution method based on reinforcement learning

[0064]

[0065]

[0066] The specific description above further illustrates the purpose, technical solutions and beneficial effects of the invention in detail. It should be understood that the above 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 in the scope of protection of the present invention.

Claims

1. An automated penetration testing method based on reinforcement learning, characterized in that The method comprises the following steps: Step 1: Use environment-related parameters to construct the state space, action space, and reward function of reinforcement learning to generate model training data; Step 1.1, define the state space of 7-dimensional parameters in the form of a dictionary, the specific form is: {'OS':'','Port':'','Service':'','Version':'','Product':'','Session':'','Access':''}, construct an index table for different parameters, and encode the strings in the dictionary into digital vectors through the parameter-index table; Step 1.2, the vulnerability exploit action is implemented by the vulnerability exploit component in the vulnerability exploit component library. There is a one-to-one correspondence between the component and the digital serial number, and the specific form is {1: no_action; 2: exploit / aix / local / ibstat_path; ...}. These components together constitute the action space of the abstract simulation environment of the penetration test. The dimension of the action space is consistent with the number of vulnerability exploit components in the vulnerability exploit component library. Step 1.3, the reward function includes three parts of rewards: vulnerability exploitation component speed reward, threat reward, and complexity reward. The speed reward of the vulnerability exploitation component can be calculated by recording the running time t of the component during the attack process. Combined with the average running time of the vulnerability exploitation component, the speed reward R1 = 1 / t can be constructed. The threat reward of the vulnerability exploitation component is evaluated by collecting information on the test environment status after the vulnerability exploitation component attacks through the information collection module. The threat reward R2 is shown in formula (1): The complexity modeling of vulnerability exploit components is calculated by calculating the code length l of the exploit components, 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 test requirements; Step 2: Apply training data and vulnerability exploitation components to generate a reinforcement learning model; Step 2.1, scan the surviving ports in the test environment, collect the operating system name, operating system version, service name, service version, port number corresponding to the service, etc. of the test environment, use string matching to determine the type, filter the service types that are not supported in the vulnerability exploit component library, and obtain the service list to be tested for penetration testing; in the training phase, first determine the OS parameters based on the operating system information obtained by the scan, and the Port, Service, Version, and Product parameters will randomly select the ports and services in the test environment identified in the information collection phase; finally, set the Session and Access parameters to 0; this state is set to the initial state s0; Step 2.2, using the DQN network structure, the number of neurons in the input layer is consistent with the number of parameters in the state space, the number of neurons in the output layer is consistent with the dimension of the action space, and there are several hidden layers in the middle; the input layer inputs the current state of the test environment s i , the output layer outputs the Q value of each action in the current state, and gives the value evaluation of each vulnerability exploitation component according to 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 vulnerability exploit component and its Q value according to the vulnerability exploit component index table, and select the component with the highest Q value as the attack action a used in this round of iteration. i ,attack the test environment using the agent-selected exploit component; Step 2.4, record the time spent in the attack process, determine the type of communication established and the permissions obtained after the attack, update the Session and Access parameters, reset the test environment if the attack is successful, and restore the current state to the initial state s0; if the attack fails, determine whether the attack has not been successful for a long time, and continue to calculate the reward update model if the maximum number of attempts has not been reached, 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 parameters will not be updated; sort out all state parameters to obtain the new state s i+1 ; Step 2.5, use the information collected in step 1.3 to calculate the speed reward and threat reward, query the code length of the corresponding vulnerability exploit component from the vulnerability exploit component library, calculate the complexity reward, and sum the three rewards to calculate the reward value r obtained by this round of attack actions. i ; Step 2.6, a fixed-size memory bank D is used to record experience. The experience is represented as [s i ,a i ,r i ,s i+1 ]; In the DQN model, there are target networks Q′ and Q networks Q, which have exactly the same structure. The Q network is used to calculate the Q value of the strategy selection. The Q network is updated during each training; the target network is used to calculate the Q value of the next state, thereby reducing the fluctuation of the target value; in each round of model update, some data is randomly selected from the memory bank and the current state s is updated. i Input Q network to predict the current action a i The Q value of the next state s i+1 Enter 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, Q(s i ,a i ; θ) is the Q value calculated by the Q network, γ is the attenuation coefficient, θ is the Q network parameter, and θ′ is the target network parameter; Use the gradient descent method to calculate the loss function and update the Q network parameter θ; when the number of training steps reaches the preset fixed value, update the target network parameter θ′ until the model reaches the predetermined number of training rounds; Step 3: Recommend optimal vulnerability exploit information based on the vulnerability exploit component library through a reinforcement learning model; Step 3.1, extract the key attributes of each asset with a vulnerability, including asset name, asset IP, vulnerability name, and vulnerability type; Step 3.2: Outlier screening and processing to ensure the fault tolerance of the intelligent scheduling model; detect and identify the non-standard IP values ​​and empty vulnerability names in the feature vector list; Step 3.3, feature encoding and normalization processing are performed on the vulnerability feature vector list after outlier processing; the input is mainly text type 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 (including asset name, asset IP, array of surviving vulnerability names, array of surviving vulnerability types), generate a corresponding vulnerability exploitation planning decision vector (including asset name, asset IP, recommended vulnerability exploitation sequence list, a total of 3 dimensions); among which, the recommended vulnerability exploitation sequence list is ranked based on comprehensive consideration of vulnerability severity, exploitation complexity and other information; Step 3.5, concatenate and summarize the decision vector list generated in the model decision phase 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, pre-process the input information and data, map and parse the recommended exploit vector matrix output by the scheduling engine, sort out and obtain the corresponding analysis for each asset, and obtain its vulnerability exploitation priority list; combine the attack path, and determine the specific feasible vulnerability exploitation path according to the specific mark or rule in the vulnerability exploitation vector; Step 4.2, according to the vulnerability exploitation path and plan, retrieve the vulnerability exploitation component from the vulnerability exploitation component library based on the vulnerability name, and construct the corresponding exploitation action; use the vulnerability association information in the plan, combined with the characteristics of the attack target, to build a series of runtime detection 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, response content, etc. to determine whether the target system is affected by the exploit or vulnerability; Step 4.4, summarize the vulnerability exploitation results and generate a system-level risk assessment report or output to assist security decision-making and vulnerability remediation processes; specific risk assessment results may include: (1) sessions created after successful vulnerability exploitation, which are created after the vulnerability exploitation action is executed, such as obtaining the target Shell to create a Session after executing the RCE exploit component; (2) vulnerability threat level information; (3) vulnerability type information.

2. The automated penetration testing method based on reinforcement learning according to claim 1, characterized in that: The modeling method of the penetration testing process in step 1, that is, defining the parameterized state space in dictionary form, constructing the action space of the penetration testing abstract simulation environment that corresponds one-to-one to the vulnerability exploitation components, defining the vulnerability exploitation component speed reward, threat degree reward, and complexity reward, and using R = λ1R1+λ2R2+λ3R3 as the final reward function.

3. The automated penetration testing method based on reinforcement learning according to claim 1, characterized in that: Step 2.6: Set the current state s i Input Q network to predict the current action a i The Q value of the next state s i+1 Enter the maximum Q value that can be obtained by predicting the next state in the target network, and the loss function is: in The Q value calculated for the target network, Q(s i ,a i ; θ) is the Q value calculated by the Q network, γ is the attenuation coefficient, θ is the Q network parameter, θ′ is the target network parameter, and the loss function is calculated using the gradient descent method.

Images