88 lines
3.7 KiB
Python
88 lines
3.7 KiB
Python
import networkx as nx
|
|
import yaml
|
|
|
|
from networkx.drawing.nx_agraph import to_agraph
|
|
from utils.log import log
|
|
|
|
class Topology(object):
|
|
def __init__(self):
|
|
self.G = nx.MultiGraph()
|
|
|
|
def build(self):
|
|
'''
|
|
Summary:
|
|
Builds a network topology from the topology.yaml file into memory, stored in a networkx Graph object.
|
|
Once the networkx Graph object has been built, it's drawn and saved as topology.png.
|
|
|
|
Nodes: network devices
|
|
Node attributes: dict of attributes assigned to a specific node
|
|
Edges: links between network devices
|
|
Edge attributes: dict of attributes assigned to a specific edge
|
|
'''
|
|
log('Building topology...', 'info')
|
|
|
|
topology_data = yaml.load(open('topology/topology.yaml'), Loader=yaml.SafeLoader)
|
|
|
|
for topology_name, topology_vars in topology_data.items():
|
|
log(f'Topology name: {topology_name}', 'info')
|
|
|
|
for node, node_attributes in topology_vars['nodes'].items():
|
|
self.G.add_node(node)
|
|
self.G.nodes[node]['shape'] = 'box'
|
|
for attribute_name, attribute_value in node_attributes.items():
|
|
self.G.nodes[node][attribute_name] = attribute_value
|
|
|
|
for edge_number, edge_attributes in topology_vars['edges'].items():
|
|
self.G.add_edge(edge_attributes['a_end'], edge_attributes['b_end'])
|
|
nx.set_edge_attributes(
|
|
self.G,
|
|
{
|
|
(edge_attributes['a_end'],
|
|
edge_attributes['b_end'],
|
|
edge_attributes['edge_index']): edge_attributes
|
|
}
|
|
)
|
|
|
|
for node1, node2, edge_attributes in self.G.edges(data=True):
|
|
edge_attributes['label'] = f"""
|
|
{edge_attributes['a_end']} ({self.G.nodes[edge_attributes['a_end']]['os']})
|
|
{edge_attributes['a_end_ip']}, {edge_attributes['a_end_interface']}\n
|
|
{edge_attributes['b_end']} ({self.G.nodes[edge_attributes['b_end']]['os']})
|
|
{edge_attributes['b_end_ip']}, {edge_attributes['b_end_interface']}
|
|
"""
|
|
|
|
log('Drawing topology...', 'info')
|
|
vis = to_agraph(self.G)
|
|
vis.layout('dot')
|
|
vis_path = 'topology/topology.png'
|
|
vis.draw(vis_path)
|
|
|
|
log(f'Topology drawn, saved as: {vis_path}', 'info')
|
|
log('Topology build complete', 'info')
|
|
|
|
def get_client_calling_for_ip(self, giaddr, _type):
|
|
'''
|
|
Summary:
|
|
Gets the hostname and operating system of the client calling for an IP from the networkx Graph object.
|
|
|
|
Takes:
|
|
giaddr: Gateway or relay IP that the DHCP packet was received from
|
|
_type: Type of DHCP packet, if type is offer log the client name and OS
|
|
|
|
Returns:
|
|
client_device_name: hostname of the client device
|
|
client_device_os: operating system of the client device
|
|
'''
|
|
for node1, node2, edge_attributes in self.G.edges(data=True):
|
|
if giaddr in edge_attributes.values():
|
|
if giaddr == edge_attributes['a_end_ip']:
|
|
client_device_name, client_device_os = edge_attributes['b_end'], self.G.nodes[edge_attributes['b_end']]['os']
|
|
elif giaddr == edge_attributes['b_end_ip']:
|
|
client_device_name, client_device_os = edge_attributes['a_end'], self.G.nodes[edge_attributes['a_end']]['os']
|
|
|
|
if _type == 'offer':
|
|
log(f'Saw client: {client_device_name}', 'info')
|
|
log(f'OS: {client_device_os}', 'info')
|
|
|
|
return client_device_name, client_device_os
|