Attack path planning method based on graph attention network and deep reinforcement learning

Abstract

This invention relates to an attack path planning method based on graph attention networks and deep reinforcement learning, belonging to the field of network security technology, and solves the problems of insufficient accuracy and efficiency in threat path mining in existing technologies. Specific steps include: obtaining the starting node and target node in the threat graph of the target network; extracting features from the threat graph using a graph attention network to obtain a topology-aware embedding representation of each network node; wherein the topology-aware embedding representation includes the features of the current node and its neighboring nodes; performing embedding mapping on the topology-aware embedding representations of all nodes to obtain a global state representation; based on the global state representation, exploring the network using threat paths to detect the optimal threat path from the starting node to the target node; and formulating a defense strategy based on the optimal threat path for network defense, thus improving the speed and quality of threat path mining.

Description

Technical Field

[0001] This invention relates to the field of network security technology, and in particular to an attack path planning method based on graph attention networks and deep reinforcement learning. Background Technology

[0002] Threat path discovery is a crucial component of proactive network defense systems. Its goal is to systematically explore and analyze vulnerabilities in the network environment by simulating the intrusion behavior of potential attackers, thereby revealing potential threat paths. This process not only helps defenders identify critical weaknesses in advance but also provides important cybersecurity data for the formulation of defense strategies and the rational allocation of threat defense resources. With the continuous increase in network scale and complexity, traditional path analysis methods based on human experience are no longer sufficient to meet the requirements of real-time performance and comprehensiveness. Therefore, automated and intelligent threat path discovery methods have gradually become a research focus.

[0003] In recent years, deep reinforcement learning, as a machine learning paradigm capable of autonomously learning optimal strategies through interaction with the environment, has been introduced into the threat path discovery problem and has achieved results in small-scale networks. However, when the network size increases, existing reinforcement learning-based methods reveal significant limitations:

[0004] 1) Insufficient representation of state information: Most methods use matrix or vector forms to model network topology, failing to fully explore the complex structural dependencies and semantic information between network nodes. This leads to biases in the agent's understanding of the environment, making it difficult to maintain good generalization ability in large-scale or heterogeneous network scenarios, ultimately affecting the feasibility and accuracy of path discovery.

[0005] 2) Low action selection efficiency: In large-scale threat graphs, due to the complex connections between network nodes, agents are prone to generating a large number of invalid or redundant actions during exploration, such as attempting to threaten unreachable network nodes or repeatedly selecting already visited network nodes. This not only increases computational costs but also significantly reduces training efficiency, leading to slow policy convergence or even instability. Summary of the Invention

[0006] Based on the above analysis, the embodiments of the present invention aim to provide an attack path planning method based on graph attention networks and deep reinforcement learning, in order to solve the problems of insufficient accuracy and efficiency in threat path mining in the prior art.

[0007] The objective of this invention is mainly achieved through the following technical solutions:

[0008] This invention provides an attack path planning method based on graph attention networks and deep reinforcement learning, comprising the following steps:

[0009] Obtain the starting node and target node from the threat graph of the target network;

[0010] The threat graph is feature extracted using a graph attention network to obtain a topology-aware embedding representation of each network node; wherein the topology-aware embedding representation includes the features of the current node and the features of its neighboring nodes.

[0011] Embedding mapping is performed on the topology-aware embedding representations of all nodes to obtain a global state representation. Based on the global state representation, the network is explored using a threat path based on reinforcement learning to detect the optimal threat path from the starting node to the target node.

[0012] Based on the optimal threat path, a defense strategy is formulated for network defense.

[0013] Furthermore, the topology-aware embedding representations of each network node are obtained, including:

[0014] A linear transformation is performed on the features of each network node in the threat map to obtain the initial feature representation of each network node;

[0015] Based on the initial feature representation of network node i and the initial feature representation of any of its first-order neighbor nodes, the attention score of any neighbor node to network node i is obtained; where, N is the total number of network nodes in the threat graph;

[0016] Based on the attention score, a masking term is introduced to obtain the attention weight of network node i;

[0017] The first-order neighbor node features of network node i are weighted and summed using the attention weights to obtain the topology-aware embedding representation of network node i.

[0018] Furthermore, the threat path exploration network is a trained actor network based on proximal policy optimization, and the optimal threat path obtained includes:

[0019] The global state representation is combined with the network node index and target node index of the current time step in the threat graph to form the observation state of the current time step;

[0020] Based on neighborhood restrictions and self-loop restrictions, a pre-defined action mask constraint is set. The observation state is input into the actor network, and combined with the action mask constraint, the final probability distribution of the candidate actions is obtained; wherein, the action is the index of the next network node to be reached.

[0021] With the goal of maximizing the cumulative reward for reaching the target node, actions are selected from candidate actions based on the final probability distribution, and the optimal threat path is obtained after multiple rounds of action selection.

[0022] Furthermore, the cumulative reward is obtained by accumulating the instant rewards obtained after performing actions at each time step; wherein, the vulnerability severity score from the current network node to the next network node at the corresponding time step is used as the instant reward.

[0023] Furthermore, the neighborhood restriction means that the agent can only perform threat transfers along adjacent network nodes defined in the threat graph; the self-loop restriction means that the agent is not allowed to select the current network node as an action to be performed.

[0024] Furthermore, combining the action mask constraints, the final probability distribution of the candidate actions is obtained as follows:

[0025] The observation state is input into the actor network to obtain the first probability distribution of the candidate actions;

[0026] The action mask constraint is multiplied by the first probability distribution to make the probability of invalid actions zero. Then, the probability of the remaining possible actions is normalized to obtain the final probability distribution.

[0027] Furthermore, the threat path exploration network is trained based on the following joint loss function:

[0028] ,

[0029] in, Let the joint loss function be... Indicates policy-based The entropy regularization term, To minimize the mean squared error loss of the critic network, Let the loss function be the actor network. These represent the weighting coefficients, These represent the actor's network parameters and the critic's network parameters, respectively.

[0030] Furthermore, the threat graph includes network nodes and edges; wherein, the network nodes are network devices of the target network, and the edges are potential threat relationships between nodes.

[0031] Furthermore, based on the configuration information of the network nodes, and combined with a general vulnerability scoring system library, the vulnerability severity scores of adjacent nodes are obtained, and the vulnerability severity scores are used as the potential threat relationships.

[0032] Furthermore, the action mask is represented as follows:

[0033] ,

[0034] in, For time steps Time network nodes The corresponding action mask; For network nodes The set of first-order neighbors; The total number of network nodes in the threat graph.

[0035] Compared with the prior art, the present invention can achieve at least one of the following beneficial effects:

[0036] 1. This invention proposes to construct a threat graph of the target network, extract network topology information of the threat graph using graph attention network, and mine threat attack paths based on the extracted topology information using deep reinforcement learning technology. This improves the efficiency and quality of threat path mining while ensuring the adaptability of the mining method. Based on the mined threat paths, defense strategies are specified to promote the rapid and effective deployment of security defense measures.

[0037] 2. Utilize graph attention network technology to extract network node features, fully leverage the topological features of the threat graph, highlight key connections, and thus generate fine-grained representations of state information; combine domain expertise to filter out invalid or redundant actions, generate action masks, avoid invalid exploration, and improve training efficiency and stability.

[0038] 3. By combining the near-end strategy optimization algorithm and leveraging the advantages of deep reinforcement learning algorithm, it can achieve self-learning exploration of high-quality threat paths, enabling rapid discovery of threat paths and vulnerability location in large-scale network environments, thereby assisting network defense.

[0039] In this invention, the above-described technical solutions can be combined with each other to achieve more preferred combinations. Other features and advantages of this invention will be set forth in the following description, and some advantages may become apparent from the description or be learned by practicing the invention. The objects and other advantages of this invention can be realized and obtained from what is particularly pointed out in the description and drawings. Attached Figure Description

[0040] The accompanying drawings are for illustrative purposes only and are not intended to limit the invention. Throughout the drawings, the same reference numerals denote the same parts.

[0041] Figure 1 This is a flowchart of the efficient threat path detection method based on graph attention and reinforcement learning according to an embodiment of the present invention;

[0042] Figure 2 This is a schematic diagram of the threat path detection method according to an embodiment of the present invention;

[0043] Figure 3 This is a schematic diagram of the threat path detection method architecture according to an embodiment of the present invention. Detailed Implementation

[0044] Preferred embodiments of the present invention will now be described in detail with reference to the accompanying drawings, which form part of this application and are used together with the embodiments of the present invention to illustrate the principles of the present invention, but are not intended to limit the scope of the present invention.

[0045] Example 1

[0046] A specific embodiment of the present invention discloses an attack path planning method based on graph attention networks and deep reinforcement learning, such as... Figure 1 As shown, it includes the following steps:

[0047] Step S1: Obtain the starting node and target node in the threat graph of the target network;

[0048] Step S2: Use a graph attention network to extract features from the threat graph to obtain a topology-aware embedding representation of each network node; wherein the topology-aware embedding representation includes the features of the current node and the features of its neighboring nodes.

[0049] Step S3: Perform embedding mapping on the topology-aware embedding representations of all nodes to obtain the global state representation. Based on the global state representation, use the threat path exploration network based on reinforcement learning to detect the optimal threat path from the starting node to the target node.

[0050] Step S4: Based on the optimal threat path, formulate a defense strategy for network defense.

[0051] The above method extracts topology-aware state representations using graph attention networks, explores stable and efficient threat paths using deep reinforcement learning algorithms, and deploys defense strategies based on the discovered threat paths. This invention can improve computational efficiency while ensuring adaptability, enabling rapid discovery of threat paths and assisting in vulnerability localization and defense.

[0052] Specifically, the target network refers to the network environment to be probed, such as an enterprise wide area network. In step S1, a threat graph of the target network is constructed based on the network nodes and edges in the target network; wherein, network nodes include network devices, service layer nodes and / or state layer nodes, and edges represent potential threat relationships that can be exploited between nodes;

[0053] Furthermore, in addition to network devices such as hosts, routers, and switches at the physical network layer, network nodes can also be system service layer nodes, such as service ports, processes, and system accounts, as well as state layer nodes, such as permissions.

[0054] By independently evaluating all current network nodes and combining the node configuration information with the Common Vulnerability Scoring System (CVSS) library, the vulnerability scores of adjacent nodes are obtained, which serve as the potential threat relationships.

[0055] Specifically, in step S2, a graph attention network is used to extract features from the threat graph, generating topology-aware embedding representations of network nodes to capture dependencies and key connections between network nodes, such as... Figure 3 As shown, the specific process is as follows:

[0056] Step 21: Perform a linear transformation on the features of each network node in the threat graph to obtain the initial feature representation of each network node;

[0057] Step 22: Based on the initial feature representation of network node i and the initial feature representation of any of its first-order neighbor nodes, obtain the attention score of any neighbor node to network node i.

[0058] Step 23: Based on the attention score, a masking term is introduced to obtain the attention weight of network node i;

[0059] Step 24: Use the attention weights to perform a weighted summation of the first-order neighbor node features of network node i to obtain the topology-aware embedding representation of network node i.

[0060] For example, given a collection of Threat graph of individual network nodes The network node feature matrix is A shared linear transformation is applied to each network node feature in the network node feature matrix based on the following formula:

[0061] ,

[0062] in, For network nodes Node characteristics; Network nodes after linear transformation The initial feature representation; , The total number of network nodes in the threat graph; To share the linear transformation matrix; These represent the initial feature dimension of the node and the dimension after linear transformation, respectively.

[0063] Record network nodes The first-order neighbor set is For any network node , The unnormalized attention score is calculated based on the following formula:

[0064] ,

[0065] in, For network nodes For network nodes Attention score; Represents a learnable vector. Indicates series connection.

[0066] To ensure that aggregation occurs only on actual edges (feasible edges, i.e., those with threat relationships) in the threat graph, the following masking term is introduced. , represented as:

[0067] .

[0068] The masking term described above is used to evaluate whether a feasible edge exists between two nodes. If it does, the value of the edge is represented as the corresponding attention score. If not, it is set to negative infinity. All neighboring nodes of a given node are selected to prepare for embedding the features of neighboring nodes into the features of the current node to obtain the network node embedding representation. This ensures the correctness of the graph structure while avoiding invalid gradients.

[0069] For network nodes The unnormalized attention score is normalized based on the following formula to obtain the attention weight. :

[0070] .

[0071] Network nodes are obtained by weighting and summing the neighborhood features using attention weights. The updated representation (i.e., the topology-aware embedding representation) is as follows:

[0072] ,

[0073] in, For network nodes The update indicates; This represents a non-linear activation function.

[0074] Through the above steps, the graph attention network can assign adaptive weights to different neighbors, realize joint encoding of key local structures and cross-level global dependencies, and thus generate a topology-aware network node embedding representation for threat path mining.

[0075] Specifically, the threat path exploration network in step S3 is a trained actor network based on proximal policy optimization, such as... Figure 2 As shown, the specific steps to obtain the optimal threat path include:

[0076] Step 311: Perform embedding mapping on the topology-aware embedding representations of all nodes to obtain the global state representation;

[0077] Step 312: Combine the global state representation with the network node index and target node index of the current time step in the threat graph to form the observation state of the current time step;

[0078] Step 313: Based on neighborhood restrictions and self-loop restrictions, preset action mask constraints, input the observation state into the actor network, and combine the action mask constraints to obtain the final probability distribution of candidate actions; wherein, the action is the index of the next network node to be reached;

[0079] Combining the action mask constraints, the final probability distribution of candidate actions is obtained as follows:

[0080] The observation state is input into the actor network to obtain the first probability distribution of the candidate actions;

[0081] The action mask constraint is multiplied by the first probability distribution to make the probability of invalid actions zero. Then, the probability of the remaining possible actions is normalized to obtain the final probability distribution.

[0082] It should be noted that the specific process of generating action masks based on domain-specific knowledge is as follows:

[0083] First, the action space is defined as follows: The action space consists of the indices of all network nodes in the threat graph, with each index corresponding to an optional action. The action space is represented as follows: .

[0084] Then, the domain expertise involves: the domain expertise includes: 1. Neighborhood restriction rule: the agent can only transfer threats along neighboring network nodes defined in the threat graph; 2. Self-loop restriction rule: the agent is not allowed to select the current node as an action to be performed.

[0085] Finally, based on domain expertise, at each time step, the current network node is generated. The corresponding action mask is represented as:

[0086]

[0087] in, For time steps Time based on the current network node The obtained action mask.

[0088] Through the above process, action masks are generated based on domain expertise and pre-defined parameters. Pruning the candidate action space filters out invalid or redundant actions. This pruning utilizes domain expertise (expert knowledge) to ensure that large-scale networks can still operate efficiently.

[0089] Step 314: With the goal of maximizing the cumulative reward for reaching the target node, select from candidate actions based on the final probability distribution; gradually construct the optimal threat path through multiple rounds of action selection; wherein, the cumulative reward is obtained by accumulating the immediate rewards obtained after executing actions at each time step; and is determined by the actions output in each state observation state. The optimal threat path is formed by combining these paths in sequence. ,in Indicates the final time step.

[0090] It should be noted that the threat path exploration strategy is obtained by training the actor network and critic network based on the proximal strategy optimization. The threat path exploration strategy is the trained actor network. The specific process of training the actor network and critic network to obtain the threat path mining strategy is as follows:

[0091] Step 321, at the time step The intelligent agent obtains a set of network node embedding representations by encoding the graph attention network described in step S2. The global state representation is obtained by performing embedding mapping on the embedding set of all network nodes. Finally, the global state is represented, and the network state observed by the agent is formed by combining the current node index and the target node index. .

[0092] Step 322, in the current network state Below, the probability distribution of candidate actions output by the actor network in the near-end policy optimization algorithm. At that time, the action mask and the probability distribution are used. After multiplying to ensure the probability of invalid actions is 0, the remaining possible actions are normalized to obtain the final action probability distribution. .like Figure 3 As shown, the specific process is as follows:

[0093] ,

[0094] in, This represents the current network state. Next, action selection The probability of. This indicates that network node j will be the next threat node.

[0095] Step 323, Critics Network Output: Critics Network Output the current state Value estimation:

[0096] ,

[0097] in, This represents the parameters of the critic network. Indicates the discount factor. Indicates at time step The immediate reward is represented by the CVSS score of the exploitable vulnerability from the current node to the next node. Additionally, if the next node is the target node, the reward is set to... When the state transitions to the next node, the reward from the current node to the next node will be updated to 0 because the vulnerability between the current node and the next node has been exploited.

[0098] Step 324: Update the actor network based on the following truncated replacement objective function:

[0099] ,

[0100] in, This indicates a truncation hyperparameter used to limit the update range of the actor network; Parameters representing the actor network; This represents the generalized advantage estimate; This represents the ratio of the new strategy to the old strategy, where the old strategy is defined as... The new strategy is ,but It can be represented as:

[0101] ,

[0102] It can be represented as: ,

[0103] in, The parameter is used to balance variance and bias.

[0104] Actor Network through Perform gradient ascent to update parameters The specific process can be represented as follows:

[0105] ,

[0106] in, This indicates the update step size of the actor's network parameters.

[0107] Step 325, Critics Network Update: The critic network is updated by minimizing the mean squared error, and its loss function can be expressed as:

[0108] ,

[0109] in, .parameter Updating via gradient descent can be represented as follows:

[0110] ,

[0111] in, This indicates the update step size of the commentator network.

[0112] Step 326, Joint Update: By combining the actor network loss function and the critic network loss function, we can obtain the joint optimization objective function, expressed as:

[0113] ,

[0114] in, Indicates policy-based The entropy regularization term is used to encourage exploration; This represents the weighting coefficient, used to balance the weights of the various loss terms; The joint loss function; The minimum mean squared error loss of the critic network is... Let be the loss function of the actor network.

[0115] Through the above process, the critic network provides state value assessment, which helps the actor network output the optimal action distribution under the constraint of action mask, thereby achieving efficient threat path mining by combining network node embedding and domain expertise.

[0116] Specifically, in step S4, after detecting the optimal threat path through simulated attacks, the optimal path is decomposed in reverse order to form a closed-loop strategy of "path-vulnerability-defense". Critical edges, critical vulnerabilities, or critical hosts are obtained through the optimal threat path, and dynamic defense strategies are specified based on the decomposed content to prioritize network defense for the threat paths most easily chosen by attackers.

[0117] Compared with existing technologies, this embodiment provides an efficient threat path detection method based on graph attention and reinforcement learning. It extracts topology-aware state representations through graph attention networks and combines them with an action masking mechanism based on expert knowledge to ensure the effectiveness of action selection. Finally, it detects stable and efficient threat paths within the framework of deep reinforcement learning algorithms and near-end policy optimization algorithms. This method improves computational efficiency while ensuring adaptability, enabling rapid discovery of threat paths and assisting in vulnerability location and defense.

[0118] Example 2

[0119] Another specific embodiment of the present invention discloses an attack path planning device based on graph attention networks and deep reinforcement learning, comprising:

[0120] The threat graph construction module is used to construct a threat graph of the target network and obtain the starting node and target node for threat detection in the threat graph;

[0121] The feature extraction module is used to extract features from the threat graph using a graph attention network to obtain a topology-aware embedding representation of each network node; wherein the topology-aware embedding representation includes the features of the current node and the features of its neighboring nodes.

[0122] The threat path detection module is used to perform embedding mapping on the topology-aware embedding representation of all nodes to obtain a global state representation. Based on the global state representation, the network is explored using threat paths to obtain the optimal threat path from the starting node to the target node.

[0123] The defense module is used to formulate defense strategies for network defense based on the optimal threat path.

[0124] The device can perform threat path detection and network defense based on the detected threat paths, according to any of the methods described in Embodiment 1. Related aspects can be referenced from each other, and are not repeated in this embodiment.

[0125] Compared with existing technologies, this embodiment provides a high-efficiency threat path detection device based on graph attention and reinforcement learning. Through the cooperation of various modules, it can not only identify key weaknesses in the network in advance, but also provide important basis for the formulation of defense strategies and the rational allocation of threat defense resources.

[0126] Those skilled in the art will understand that all or part of the processes of the methods described in the above embodiments can be implemented by a computer program instructing related hardware, and the program can be stored in a computer-readable storage medium. The computer-readable storage medium may be a disk, optical disk, read-only memory, or random access memory, etc.

[0127] The above description is only a preferred embodiment of the present invention, but the scope of protection of the present invention is not limited thereto. Any changes or substitutions that can be easily conceived by those skilled in the art within the scope of the technology disclosed in the present invention should be included within the scope of protection of the present invention.

Claims

1. An attack path planning method based on graph attention networks and deep reinforcement learning, characterized in that, Includes the following steps: Obtain the starting node and target node from the threat graph of the target network; The threat graph is feature extracted using a graph attention network to obtain a topology-aware embedding representation of each network node; wherein the topology-aware embedding representation includes the features of the current node and the features of its neighboring nodes. Embedding mapping is performed on the topology-aware embedding representations of all nodes to obtain a global state representation. Based on the global state representation, the network is explored using a threat path based on reinforcement learning to detect the optimal threat path from the starting node to the target node. Based on the optimal threat path, a defense strategy is formulated for network defense.

2. The method according to claim 1, characterized in that, The topology-aware embedding representations of each network node are obtained as follows: A linear transformation is performed on the features of each network node in the threat map to obtain the initial feature representation of each network node; Based on the initial feature representation of network node i and the initial feature representation of any of its first-order neighbor nodes, the attention score of any neighbor node to network node i is obtained; where, N is the total number of network nodes in the threat graph; Based on the attention score, a masking term is introduced to obtain the attention weight of network node i; The first-order neighbor node features of network node i are weighted and summed using the attention weights to obtain the topology-aware embedding representation of network node i.

3. The method according to claim 1 or 2, characterized in that, The threat path exploration network is a trained actor network based on proximal policy optimization, and the optimal threat path obtained includes: The global state representation is combined with the network node index and target node index of the current time step in the threat graph to form the observation state of the current time step; Based on neighborhood restrictions and self-loop restrictions, a pre-defined action mask constraint is set. The observation state is input into the actor network, and combined with the action mask constraint, the final probability distribution of the candidate actions is obtained; wherein, the action is the index of the next network node to be reached. With the goal of maximizing the cumulative reward for reaching the target node, actions are selected from candidate actions based on the final probability distribution, and the optimal threat path is obtained after multiple rounds of action selection.

4. The method according to claim 3, characterized in that, The cumulative reward is obtained by accumulating the instant rewards obtained after each time step of the action; wherein, the vulnerability severity score from the current network node to the next network node at the corresponding time step is used as the instant reward.

5. The method according to claim 3, characterized in that, The neighborhood restriction means that the agent can only perform threat transfers along neighboring network nodes defined in the threat graph; the self-loop restriction means that the agent is not allowed to select the current network node as an action to be performed.

6. The method according to claim 3, characterized in that, Combining the action mask constraints, the final probability distribution of candidate actions is obtained as follows: The observation state is input into the actor network to obtain the first probability distribution of the candidate actions; The action mask constraint is multiplied by the first probability distribution to make the probability of invalid actions zero. Then, the probability of the remaining possible actions is normalized to obtain the final probability distribution.

7. The method according to claim 3, characterized in that, The threat path exploration network is trained based on the following joint loss function: ， in, Let the joint loss function be... Indicates policy-based The entropy regularization term, To minimize the mean squared error loss of the critic network, Let the loss function be the actor network. These represent the weighting coefficients, These represent the actor's network parameters and the critic's network parameters, respectively.

8. The method according to any one of claims 1-2 and 4-7, characterized in that, The threat graph includes network nodes and edges; wherein, the network nodes are network devices of the target network, and the edges are potential threat relationships between nodes.

9. The method according to claim 8, characterized in that, Based on the configuration information of the network nodes, and combined with a general vulnerability scoring system library, the vulnerability severity scores of adjacent nodes are obtained, and the vulnerability severity scores are used as the potential threat relationships.

10. The method according to claim 3, characterized in that, The action mask is represented as follows: ， in, For time steps Time network nodes The corresponding action mask; For network nodes The set of first-order neighbors; The total number of network nodes in the threat graph.

Images