Multi-agent collaborative penetration path discovery method, system and medium based on deep hierarchical reinforcement learning

Abstract

The present invention proposes a multi-agent collaborative penetration path discovery method, system and medium based on deep hierarchical reinforcement learning. Based on hierarchical reinforcement learning, the penetration path discovery is constructed as a hierarchical Markov decision method. By decomposing the penetration test into the upper-layer intelligent agent's lateral movement target host selection and the lower-layer intelligent agent's single-host penetration process, the target host selection is used as the sub-goal of the lower-layer intelligent agent. Through the hierarchical training of the upper and lower layer networks, the problem of algorithm convergence difficulty caused by the explosion of the intelligent agent state space, the discrete high dimension of the action space and the sparse rewards in large-scale networks is effectively alleviated, thereby accelerating the efficiency of the penetration test and ensuring the high sensitivity of the test results.

Description

Technical Field

[0001] The present invention relates to the field of network security testing, and in particular to a multi-agent collaborative penetration path discovery method, system, and medium based on deep hierarchical reinforcement learning. Background Art

[0002] Penetration testing is a common network security testing method that simulates a hacker attack by exploiting system vulnerabilities to gain control of the target system network without disrupting its normal operation. This method allows security personnel to analyze the vulnerabilities and vulnerable nodes found in the test report and implement appropriate protective measures to make the system network more robust and less susceptible to hacker attacks.

[0003] Traditional penetration testing relies primarily on manual methods, relying on the experience and knowledge of security experts and the assistance of penetration tools. However, given the large and complex network systems of modern enterprises and organizations, relying solely on manual methods to conduct security testing is impractical and requires significant time and labor costs, increasing the cost of penetration testing. Therefore, utilizing automated and intelligent penetration testing technology to simulate real attackers with attack strategies, identify key attack paths within the target system network, and discover network weaknesses can significantly reduce penetration testing costs and improve their efficiency.

[0004] Reinforcement learning (RL) is a learning paradigm that effectively solves sequential decision-making problems by acquiring empirical knowledge through exploration and exploitation of unknown environments and learning optimal strategies for taking actions within them based on this experience. Compared to supervised and unsupervised learning, RL has the advantage of not relying on large amounts of static data. It can learn strategies that maximize reward based on experience, even with limited or no prior knowledge. It is therefore suitable for real-time and adversarial environments.

[0005] Penetration testing is a dynamic decision-making process that interacts with the environment in real time. Testers perform various attacks and, based on the feedback provided by the network environment, discover host vulnerabilities and configuration information, thereby identifying network vulnerabilities through analysis. Therefore, reinforcement learning can be used as an artificial intelligence method to simulate and develop penetration tests for network security, providing new insights for intelligent attack path discovery.

[0006] Existing penetration testing methods based on reinforcement learning usually formalize the penetration testing process as a Markov Decision Process (MDP). By designing the success probability of action execution, it can simulate to a certain extent the uncertainty of attackers in the target network system in the real world. By designing a reward and punishment mechanism, it guides the trained intelligent agent to select the best action based on the environmental state. GAIL-PT, a deep reinforcement learning intelligent penetration testing method based on imitation learning, combines generative adversarial imitation learning networks into penetration testing for the first time. It builds an expert knowledge base by pre-collecting the successful penetration experience of security experts, improves the overall penetration performance of the model through adversarial training of the generator and discriminator, and verifies it in real target ranges and simulated networks. In the existing technology, the automated penetration testing method based on NIG-AP introduces information entropy into the reward value design of reinforcement learning, and realizes penetration attack path discovery based on network information gain. This method uses network information to obtain rewards and guides the intelligent agent to choose the best action.

[0007] Based on existing penetration testing, attack path discovery refers to the process by which an attacker exploits a sequence of vulnerabilities in a target network to gain access to specific assets. Attack path discovery encompasses not only penetration paths between hosts but also the attack process within a single host. When attacking a single host, an agent cannot alter the state of the entire network. Only after successfully compromising that host can it, through lateral movement and attacking other hosts, change the state of the current network host.

[0008] In recent years, research on attack path discovery based on reinforcement learning has become a hot topic in related fields. Figure 1 As shown, GAIL-PT is an existing automated penetration testing method based on generative adversarial imitation learning. It performs three main phases of automated penetration testing: 1. Construction of a penetration expert knowledge base; 2. Generative adversarial network training; and 3. GAIL-based automated penetration testing. First, state-action pairs from successful penetrations of pre-trained algorithm models are automatically collected to build an expert knowledge base. Second, state-action pairs and expert knowledge generated online by different deep reinforcement learning models are input into the GAIL-PT discriminator for training. Finally, the discriminator's output is used to discount rewards to guide the agent to generate actions with a higher penetration probability, thereby improving overall penetration performance.

[0009] However, existing intelligent penetration testing methods still face the following three major technical difficulties: The first is the state space explosion problem; the local state of each agent has multiple dimensions, and the state space of the planning problem will grow exponentially. The second major difficulty is the large-scale action space problem; as the network scale increases, the types and number of vulnerabilities also increase, causing the action space to continue to expand. The third difficulty is the sparse reward problem; as the network scale continues to expand, it will be difficult for the agent to successfully find sensitive targets, and the sparse reward problem will become increasingly prominent. Therefore, how to solve these problems in the intelligent penetration testing process and improve the penetration efficiency of agents in large-scale network environments is a current research focus. Summary of the Invention

[0010] In order to address the deficiencies in the prior art, the present invention provides a multi-agent collaborative penetration path discovery method, system and medium based on deep hierarchical reinforcement learning. Based on the idea of ​​hierarchical reinforcement learning, the penetration path discovery process is constructed as a hierarchical Markov decision process, wherein the upper-layer strategy performs penetration path discovery between hosts, and the lower-layer strategy performs attack path discovery for a single host. By decomposing the continuous state change process of the penetration test into an intermittent state change process, the penetration test problem is decomposed into two serial subtasks, which effectively alleviates the problems of agent state, action space explosion and sparse rewards, thereby accelerating training efficiency.

[0011] The first aspect of the present invention discloses a multi-agent collaborative penetration path discovery method based on deep hierarchical reinforcement learning, which is applied to multi-agent collaborative penetration testing and includes the following steps:

[0012] Step S1: The upper and lower agents obtain the initial state s of the infiltration environment t , and initialize the upper network parameters θ at the same time h and the lower layer network parameters θ l ;

[0013] Step S2: The initial state s obtained by the upper network input t , the upper-layer agent learns the upper-layer strategy π h , the upper agent is based on ε h -greedy greedy strategy selects the target host g t , the upper agent obtains the external reward accumulated in the N-step time interval starting from time t where r i out Represents the external reward obtained by the upper-layer agent in a single time step; Step S3: Determine the target host g t Is it penetrated? If not, execute step S4. If the target host g tIf it is infiltrated, execute step S8;

[0014] Step S4: The lower-level agent selects the target host g t As the penetration target, according to ε l -greedy greedy strategy generates penetration action a t ;

[0015] Step S5: The lower-level agent performs the penetration action a t , the upper agent obtains the external reward R in step S2 t out , and generate the next moment state s t+1 ;

[0016] Step S6: The lower-level agent obtains the internal reward r in , and store ((s t ,g t ),a t ,r in ,(s t+1 ,g t+1 ),v t ))Sequence to the lower experience pool D l ; where s t Indicates the current state of the lower-level agent, g t Indicates that the target host is selected, a t represents the penetration action generated by the lower-level agent at the current moment, s t+1 , g t+1 They represent the state of the lower-level agent and its selected target host at the next moment, respectively. t A flag indicating whether the host reached by the lower-level agent has been successfully infiltrated;

[0017] Step S7: Update the upper network parameters and the lower network parameters θ respectively h and θ l , and simultaneously update the external reward R of the upper agent out ←R out +r in ,quit;

[0018] Step S8: Target host g t Has been infiltrated, giving the upper-level agent an external reward R out =-10000;

[0019] Step S9: Store the sequence Supreme Experience Pool D h , where s t Represents the current state of the upper network input, the agent learns the upper strategy, and outputs the selected target host g through the strategy generatort , represents the external reward accumulated by the upper agent from the time step interval, C t Refers to the flag of whether the agent has compromised all sensitive hosts at the current moment, represented by True and False, s t ′ represents the next state of the agent; at the same time, the upper network parameters θ are updated h ,quit.

[0020] According to the method of the first aspect of the present invention, in step S9, the upper layer network parameter θ is updated. h The steps include:

[0021] Step S9.1: The upper agent obtains the optimal value function by maximizing the target reward function through formula (1):

[0022]

[0023] Among them, π h Indicates the real-time policy generated by the upper host. represents the external reward accumulated by the upper agent from the time step interval, g t ' indicates that the next N time intervals are in s t+N The target host selected by the upper agent in the state, s t ′ represents the next state of the agent;

[0024] Step S9.2: Update the upper network parameters θ by deriving the upper network loss function shown in formula (2) h :

[0025]

[0026] in, Indicates sampling from the upper-level experience buffer pool; when the upper-level agent reaches the sub-goal or exceeds the set number of round steps T, the upper-level agent will reselect the target host.

[0027] According to the method of the first aspect of the present invention, step S7 includes:

[0028] Step S7.1: The lower-level agent calculates the optimal augmented state-action pair function using the DDQN algorithm, as shown in formula (3):

[0029]

[0030] Among them, π l Indicates the policy generated by the upper host, r t in Represents the internal reward of the lower-level agent at time t; each time the agent performs an action at The internal reward r will be generated based on whether the sub-goal is achieved in , a t+1 and g t+1 They represent the next action of the lower-level agent and the infiltration target selected at the next moment, and at the same time update the external reward R of the upper-level agent. out ←R out +r in ;

[0031] Step S7.2: Update the lower layer network parameters θ by deriving the lower layer network loss function shown in Equation (4) l :

[0032]

[0033] Among them, ((s t ,g t ),a t ,r in ,(s t+1 ,g t+1 ),v t ))~D l Represents sampling from the lower-level experience buffer pool, and α and β represent the parameters of the value function and advantage function in the DDQN network, respectively.

[0034] The second aspect of the present invention discloses a multi-agent collaborative penetration path discovery system based on deep hierarchical reinforcement learning, including a computer device for executing the steps of the multi-agent collaborative penetration path discovery method based on deep hierarchical reinforcement learning described in the first aspect.

[0035] The third aspect of the present invention discloses a computer-readable storage medium, which stores a computer program to implement the multi-agent collaborative penetration path discovery method based on deep hierarchical reinforcement learning described in the first aspect.

[0036] In summary, the solution proposed by the present invention has the following technical effects: based on the idea of ​​hierarchical reinforcement learning, the penetration path discovery process is constructed as a hierarchical Markov decision process, in which the upper-level strategy performs penetration path discovery between hosts, and the lower-level strategy performs attack path discovery on a single host. By decomposing the continuous state change process of the penetration test into an intermittent state change process, the penetration test problem is decomposed into two serial subtasks, thereby accelerating the test efficiency. In addition, through the lateral movement of the target host selection by the upper-level intelligent agent and the single-host penetration of the lower-level intelligent agent, the network host state can be changed in real time based on the vulnerability obtained in the attack, ensuring the timeliness of the penetration test. BRIEF DESCRIPTION OF THE DRAWINGS

[0037] In order to more clearly illustrate the specific embodiments of the present invention or the technical solutions in the prior art, the following briefly introduces the drawings required for use in the specific embodiments or the description of the prior art. Obviously, the drawings described below are some embodiments of the present invention. For ordinary technicians in this field, other drawings can be obtained based on these drawings without paying any creative work.

[0038] Figure 1 This is the flow chart of the GAIL-PT method;

[0039] Figure 2 This is a flow chart of the multi-agent collaborative penetration path discovery method based on deep hierarchical reinforcement learning of the present invention. DETAILED DESCRIPTION

[0040] To make the objectives, technical solutions, and advantages of the embodiments of the present invention more clear, the technical solutions in the embodiments of the present invention will be clearly and completely described below in conjunction with the accompanying drawings of the embodiments of the present invention. Obviously, the described embodiments are only part of the embodiments of the present invention, not all of the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without making creative efforts shall fall within the scope of protection of the present invention.

[0041] Based on the method proposed in this paper, a multi-agent collaborative penetration path discovery method based on deep hierarchical reinforcement learning is proposed for the hierarchical Markov decision process. This section decomposes the attack path discovery process into a lateral penetration process between hosts and a penetration attack process on a single host. Based on two deep reinforcement learning models with different network structures, host selection and attack action selection are completed respectively, achieving rapid penetration path discovery. Finally, this paper trains agents based on the proposed heterogeneous hierarchical reinforcement learning collaborative algorithm and conducts comparative experiments based on virtual network scenarios constructed using the open source network attack simulator NASim. By varying the network scale, the performance advantages and scalability of the proposed method are verified.

[0042] like Figure 2 The figure shows a flow chart of the multi-agent collaborative penetration path discovery method based on deep hierarchical reinforcement learning. The specific execution process is as follows:

[0043] Step 1: The agent obtains the initial state s of the infiltration environment t , and initialize the upper network parameters θ at the same time h , the lower layer network parameters θ l ;

[0044] Step 2: The goal of the upper layer task is to complete the lateral penetration between hosts and generate a penetration path. The upper layer network input current state s t, the agent learns the upper-level strategy π h , output the selected target host g through the policy generator t , the upper agent obtains the immediate external reward accumulated over the N-step time interval starting from time t

[0045] Step 3: Determine whether the target host has been penetrated, V t Indicates the penetration flag of the target host; if V t = True, indicating that the selected target host has been infiltrated, proceed to step 4. If the target host has not been infiltrated, V t =False, skip to step 7;

[0046] Step 4: Assign external reward R out = -10000 and store the sequence Supreme Experience Pool D h , where C t Refers to the flag indicating whether the agent has compromised all sensitive hosts at the current moment, represented by True and False;

[0047] Step 5: The upper agent obtains the optimal value function by maximizing the target reward function through formula (1):

[0048]

[0049] Among them, g t ′ represents the target host selected by the upper agent in the next time interval, s t ′ represents the next state of the agent;

[0050] Step 6: Update the upper network parameters θ by deriving the upper network loss function shown in formula (1) h :

[0051]

[0052] in, Indicates sampling from the upper-level experience buffer pool. When the upper-level agent reaches the sub-goal or exceeds the set number of round steps T, the upper-level agent will reselect the target host;

[0053] Step 7: The lower-level task is to complete the attack action planning for a single host and achieve penetration of the host. The lower-level agent selects the controlled host g through the upper-level agent. t As the penetration target, according to ε l -greedy greedy strategy generates penetration action a t ;

[0054] Step 8: Execute action a t, the upper agent obtains external reward R out And generate the next moment state s t+1 , while the lower-level agent obtains the internal reward r in , storage sequence ((s t ,g t ),a t ,r in ,(s t+1 ,g t+1 ),v t )) to the lower experience pool D l ;

[0055] Step 9: The lower-level agent uses the DDQN algorithm to calculate the optimal augmented state-action pair value function, as shown in formula (3):

[0056]

[0057] Among them, r t in Represents the internal reward of the lower-level agent at time t. Each time the agent performs an action a t The internal reward r will be generated based on whether the sub-goal is achieved in , and update the external reward R of the upper agent at the same time out ←R out +r in ;

[0058] Step 10: Update the lower layer network parameters θ by deriving the lower layer network loss function shown in formula (3) l :

[0059]

[0060] Among them, ((s t ,g t ),a t ,r in ,(s t+1 ,g t+1 ),v t ))~D l Represents sampling from the lower-level experience buffer pool, and α and β represent the unique parameters of the value function and advantage function in the DDQN network, respectively.

[0061] In addition, an embodiment of the present invention also provides a multi-agent collaborative penetration path discovery system based on deep hierarchical reinforcement learning, including a computer device, which is used to execute the steps of the aforementioned multi-agent collaborative penetration path discovery method based on deep hierarchical reinforcement learning.

[0062] An embodiment of the present invention also provides a computer-readable storage medium, which stores a computer program to implement the aforementioned multi-agent collaborative penetration path discovery method based on deep hierarchical reinforcement learning.

[0063] In summary, this invention decomposes penetration testing into the upper-layer agent's lateral movement target host selection and the lower-layer agent's single-host penetration process. By making target host selection a sub-goal for the lower-layer agent, the layered training of the upper and lower-layer networks effectively alleviates the problem of agent state space explosion, discrete high-dimensional action space, and sparse rewards in large-scale networks, which leads to algorithm convergence difficulties. This is a new approach to intelligent penetration testing based on deep reinforcement learning.

[0064] Please note that the technical features of the above embodiments can be combined arbitrarily. In order to make the description concise, not all possible combinations of the technical features in the above embodiments are described. However, as long as there is no contradiction in the combination of these technical features, they should be considered to be within the scope of this specification. The above embodiments only express several implementation methods of the present application. The description is relatively specific and detailed, but it cannot be understood as a limitation on the scope of the invention patent. It should be pointed out that for ordinary technicians in this field, without departing from the concept of this application, several variations and improvements can be made, which all fall within the scope of protection of this application. Therefore, the scope of protection of the patent in this application shall be based on the attached claims.

Claims

1. A collaborative penetration testing method based on heterogeneous layered reinforcement learning, applied to multi-agent collaborative penetration testing, characterized by: The steps include: Step S1: The upper and lower agents obtain the initial state s of the infiltration environment t , and initialize the upper network parameters θ at the same time h and the lower layer network parameters θ l ; Step S2: The initial state s obtained by the upper network input t , the upper-layer agent learns the upper-layer strategy π h , the upper agent is based on ε h -greedy greedy strategy selects the target host g t , the upper agent obtains the external reward accumulated in the N-step time interval starting from time t in Represents the external reward obtained by the upper-level agent at a single time step; Step S3: Determine the target host g t Is it penetrated? If not, execute step S4. If the target host g t If it is infiltrated, execute step S8; Step S4: The lower-level agent selects the target host g t As the penetration target, according to ε l -greedy greedy strategy generates penetration action a t ; Step S5: The lower-level agent performs the penetration action a t , the upper agent obtains the external reward in step S2 And generate the next moment state s t+1 ; Step S6: The lower-level agent obtains the internal reward r in , and store ((s t , g t ),a t , r in ,(s t+1 , g t+1 ),v t ))Sequence to the lower experience pool D l ; where s t Indicates the current state of the lower-level agent, g t Indicates that the target host is selected, a t represents the penetration action generated by the lower-level agent at the current moment, s t+1 , g t+1 They represent the state of the lower-level agent and its selected target host at the next moment, respectively. t A flag indicating whether the host reached by the lower-level agent has been successfully infiltrated; Step S7: Update the upper network parameters and the lower network parameters θ respectively h and θ l , and simultaneously update the external reward R of the upper agent out ←R out +r in ,quit; Step S8: Target host g t Has been infiltrated, giving the upper-level agent an external reward R out =-10000; Step S9: Store the sequence Supreme Experience Pool D h , where s t Represents the current state of the upper network input, the agent learns the upper strategy, and outputs the selected target host g through the strategy generator t , represents the external reward accumulated by the upper agent from time t, C t Refers to the flag of whether the agent has compromised all sensitive hosts at the current moment, represented by True and False, s′ t It represents the next state of the agent; at the same time, the upper network parameters θ are updated h ,quit; Wherein, the step S7 includes: Step S7.1: The lower-level agent calculates the optimal augmented state-action pair function using the DDQN algorithm, as shown in formula (3): Among them, π l Indicates the policy generated by the upper host. Represents the internal reward of the lower-level agent at time t; each time the agent performs an action a t The internal reward r will be generated based on whether the sub-goal is achieved in , a t+1 and g t+1 They represent the next action of the lower-level agent and the infiltration target selected at the next moment, and at the same time update the external reward R of the upper-level agent. out ←R out +r in ; Step S7.2: Update the lower layer network parameters θ by deriving the lower layer network loss function shown in Equation (4) l : Among them, ((s t , g t ),α t , r in ,(s t+1 , g t+1 ),v t ))~D1 represents sampling from the lower-level experience buffer pool, and α and β represent the parameters of the value function and advantage function in the DDQN network respectively; Among them, in step S9, the upper layer network parameter θ is updated h The steps include: Step S9.1: The upper agent obtains the optimal value function by maximizing the target reward function through formula (1): Among them, π h Indicates the real-time policy generated by the upper host. represents the external reward accumulated by the upper agent from time t, g′ t Indicates that the next N time intervals are in s t+N The target host selected by the upper agent in the state, s′ t It represents the next state of the agent; Step S9.2: Update the upper network parameters θ by deriving the upper network loss function shown in formula (2) h : in, Indicates sampling from the upper-level experience buffer pool; when the upper-level agent reaches the sub-goal or exceeds the set number of round steps T, the upper-level agent will reselect the target host.

2. A collaborative penetration testing system based on heterogeneous layered reinforcement learning, including computer equipment, characterized by: The computer device is used to execute the steps of the collaborative penetration testing method based on heterogeneous layered reinforcement learning described in claim 1.

3. A computer-readable storage medium, characterized in that: The computer-readable storage medium stores a computer program to implement the collaborative penetration testing method based on heterogeneous layered reinforcement learning as described in claim 1.

Images