From 255420e6223a97974dc5e50617a20896137fcb12 Mon Sep 17 00:00:00 2001 From: Marcin Kostrzewski Date: Thu, 16 Jun 2022 22:05:00 +0200 Subject: [PATCH] Add path traversal with infection --- network_attack_propagation.py | 46 +++++++++++++++++++++++++++++------ 1 file changed, 38 insertions(+), 8 deletions(-) diff --git a/network_attack_propagation.py b/network_attack_propagation.py index 0e1cad7..e1c6c29 100644 --- a/network_attack_propagation.py +++ b/network_attack_propagation.py @@ -25,10 +25,14 @@ class Edge: def as_tuple(self): return self.node_a, self.node_b, {'weight': self.weight} + def has_node(self, node: Node) -> bool: + return self.node_a is node or self.node_b is node + class Graph: - def __init__(self): + def __init__(self, use_weights=False): self.edges = [] + self.use_weights = use_weights def add_edge(self, edge: Edge): self.edges.append(edge) @@ -43,6 +47,31 @@ class Graph: nodes.add(edge.node_b) return nodes + def get_adjacent_nodes(self, node: Node) -> [(Node, int)]: + """ + :param node: Node to search for + :return: An array of tuples (node, weight) + """ + edges_with_node = filter(lambda ed: ed.has_node(node), self.edges) + nodes = set() + for e in edges_with_node: + if e.node_a is node: + nodes.add((e.node_b, e.weight)) + else: + nodes.add((e.node_a, e.weight)) + + return nodes + + def infect_step(self): + infected_nodes = list(filter(lambda n: n.is_infected, self.get_nodes())) + for node in infected_nodes: + adjacent_nodes = self.get_adjacent_nodes(node) + if self.use_weights: + to_be_infected = random.choices([n[0] for n in adjacent_nodes], weights=[n[1] for n in adjacent_nodes])[0] + else: + to_be_infected = random.choice([n[0] for n in adjacent_nodes]) + to_be_infected.is_infected = True + def update(num, layout, g_repr, ax, our_graph: Graph): """ @@ -50,9 +79,6 @@ def update(num, layout, g_repr, ax, our_graph: Graph): """ ax.clear() - for n in our_graph.get_nodes(): - n.is_infected = bool(random.getrandbits(1)) - colors = ['red' if n.is_infected else 'blue' for n in g_repr] sizes = [50 if n.is_infected else 1 for n in g_repr] nx.draw( @@ -67,6 +93,8 @@ def update(num, layout, g_repr, ax, our_graph: Graph): linewidths=40, ) + our_graph.infect_step() + def do_graph_animation(output_file_name: str, in_graph: Graph, frame_count: int, layout): g_repr = nx.Graph() @@ -80,7 +108,7 @@ def do_graph_animation(output_file_name: str, in_graph: Graph, frame_count: int, fig.set_figheight(15) anim = animation.FuncAnimation( - fig, update, frames=frame_count, fargs=(layout, g_repr, ax, in_graph) + fig, update, frames=frame_count, interval=500, fargs=(layout, g_repr, ax, in_graph) ) anim.save(output_file_name) @@ -88,9 +116,10 @@ def do_graph_animation(output_file_name: str, in_graph: Graph, frame_count: int, plt.show() -def bus_network(n=30) -> Graph: +def bus_network(n=30, infected_idx=0) -> Graph: network = Graph() nodes = [Node() for _ in range(n)] + nodes[infected_idx].is_infected = True edges = [Edge(nodes[i], nodes[i + 1], 1.0) for i in range(n - 1)] network.add_edges(edges) @@ -100,6 +129,7 @@ def bus_network(n=30) -> Graph: def ring_network(n=30) -> Graph: network = Graph() nodes = [Node() for _ in range(n)] + nodes[0].is_infected = True edges = [Edge(nodes[i], nodes[i + 1], 1.0) for i in range(n - 1)] end_edge = Edge(nodes[n - 1], nodes[0], 1.0) edges.append(end_edge) @@ -126,10 +156,10 @@ def main(): do_graph_animation('test.gif', network, 5, nx.spring_layout) bus = bus_network() - do_graph_animation('bus.gif', bus, 5, nx.spiral_layout) + do_graph_animation('bus.gif', bus, 20, nx.spiral_layout) ring = ring_network() - do_graph_animation('ring.gif', ring, 5, nx.circular_layout) + do_graph_animation('ring.gif', ring, 20, nx.circular_layout) if __name__ == "__main__":