sshgraph.py

#!/usr/bin/env python3

import sys
import os
import dash
import dash_cytoscape as cyto
from dash import html, Input, Output, callback, dcc, State
from sshpubkeys import SSHKey
from collections import Counter


######## Functions ############

def check_key(line):
    key = SSHKey(line)
    try:
        key.parse()
    except InvalidKeyError as err:
        print("Invalid key:", err)
        sys.exit(1)
    except NotImplementedError as err:
        print("Invalid key type:", err)
        sys.exit(1)


def parse_publicssh_key(machinename):
        
        fullpath = os.path.join("publicssh_keys/", machinename) 
        
        try:
            machine_key_file = open(fullpath)
        except:
            # print("no publicssh key file for " + machinename)
            machine_key_file = ""
            machine_key_type = ""
            machine_key_data = ""
            machine_key_comment = machinename

            
        if machine_key_file:
            machine_key = machine_key_file.read()
            if machine_key and machine_key.strip() and not machine_key.startswith("#"):
                check_key(machine_key)
                machine_key_splitted = machine_key.split(" ", 2)
                machine_key_type = machine_key_splitted[0]
                machine_key_data = machine_key_splitted[1]
                try:
                    machine_key_comment = machine_key_splitted[2]
                except:
                    machine_key_comment = machinename
            else:
                machine_key_type = ""
                machine_key_data = ""
                machine_key_comment = machinename

        machine_identifier = machine_key_data + machine_key_comment

        return machine_identifier, machine_key_comment


def parse_authorized_keys_file(machine_key_file):

        check_key(machine_key_file)
        
        authorized_key_splitted = machine_key_file.split(" ", 2)
        authorized_key_type = authorized_key_splitted[0]
        authorized_key_data = authorized_key_splitted[1]
        try:
            authorized_key_comment = authorized_key_splitted[2]
        except:
            authorized_key_comment = "NULL"

        authorized_key_identifier = authorized_key_data+authorized_key_comment

        return authorized_key_identifier,authorized_key_comment



def add_node(id, label, classes):
    graph_elements = graph.elements + [
    {'data': {'id': id, 'label': label}, 'classes': classes}
    ]
    graph.elements = graph_elements

def add_connection(source, target, classes):
    graph_elements = graph.elements + [
    {'data': {'source': source, 'target': target}, 'classes': classes}
    ]
    graph.elements = graph_elements


def create_cytoscape_graph():
    graph = cyto.Cytoscape(
        id='graph',
        layout={'name': 'klay'},
        style={'width': '99vw',
            'height': '98vh',
            'z-index': '1'
            },
        elements=[],
        responsive=True,
        autoRefreshLayout=False,
        boxSelectionEnabled=True,
        stylesheet=[
        {
            'selector': 'edge',
            'style': {
                'curve-style': 'bezier'
            }
        },
        {
            'selector': '.auth_to_machine',
            'style': {
                'target-arrow-color': 'black',
                'target-arrow-shape': 'triangle'
            }
        },
        {
            'selector': '.machine',
            'style': {
                'content': 'data(label)',
                'background-color': 'black',
            }
        },
        {
            'selector': ':selected.machine',
            'style': {
                'content': 'data(label)',
                'background-color': 'blue',
            }
        },
        {
            'selector': '.auth_key',
            'style': {
                'content': 'data(label)'
            }
        }
    ]
    )

    return graph


def scan_keys_and_add_elements():

    # Hauptschleife, geht durch den Ordner authorized_keys, in welchem je Maschine eine Datei gibt.
    for root, dirs, files in os.walk('authorized_keys/'):
        
        for filename in files:

            machinename = filename

            machine_identifier, machine_key_comment = parse_publicssh_key(machinename)
            add_node(machine_identifier, machine_key_comment, 'machine')


            fullpath = os.path.join(root, filename) 
            with open(fullpath, "r") as authorized_keys_file:
                

                for idx, line in enumerate(authorized_keys_file.read().splitlines()):
                    if line.strip() and not line.startswith("#"):

                        authorized_key_identifier,authorized_key_comment = parse_authorized_keys_file(line)

                        add_node(authorized_key_identifier, authorized_key_comment, 'auth_key')
                        add_connection(authorized_key_identifier, machine_identifier, 'auth_to_machine')






def remove_auth_keys_which_are_machines():
    # Wenn durch die alle authorized_keys file gegangen wird, wird jeder Key erstmal als classes=auth_key eingetragen
    # Es besteht aber die Möglichkeit, dass einer dieser auth_keys bereits als classes=machine hinzugefügt wurde
    # Aber weil anhand der classes der Style geändert wird (Farben), werden manche als auth_keys gestyled
    # Deswegen schaut man nach, ob es diese Elemente mit auth_key bereits als classes=machine gibt
    # Zuerst die machine IDs sammeln:
    nur_machines = []
    for graph_element in graph.elements:
        if "id" in graph_element['data']:
            if graph_element['classes'] == 'machine':
                nur_machines.append(graph_element['data']['id'])
    # Dann vergleichen ob diese auth_key ID in der Machine ID List ist und wenn ja, dann den auth_key löschen
    for graph_element in graph.elements:
        if "id" in graph_element['data']:
            if graph_element['classes'] == 'auth_key':
                if graph_element['data']['id'] in nur_machines:
                    graph.elements.remove(graph_element)


def dict_to_tuple(d):
    """
    Convert a dictionary to a hashable tuple by recursively converting nested dictionaries to frozensets.
    """
    return tuple((k, dict_to_tuple(v) if isinstance(v, dict) else v) for k, v in d.items())



def remove_duplicate_auth_keys():
    # Doppelte Einträge Löschen von den auth_keys
    # Für add_node mit Machine diese sind unique
    # Aber beim durchgehen aller authorized_keys files sind da meherere gleiche enthalten.Diese werden zwar im Graph nicht angezeigt, aber graph.elements ist halt größer
    # Deswgen die doppelten Einträge löschen
    # Anhand der unique id mit classes=auth_key jeweils nur 1x da lassen, alle anderen löschen
    only_auth_keys_list = []
    for graph_element in graph.elements:
        if "id" in graph_element['data']:
            if graph_element['classes'] == "auth_key":
                only_auth_keys_list.append(graph_element)
                graph.elements.remove(graph_element)

    unique_dicts = []
    seen = set()

    for d in only_auth_keys_list:
        dict_tuple = dict_to_tuple(d)
        if dict_tuple not in seen:
            unique_dicts.append(d)
            seen.add(dict_tuple)

    only_auth_keys_list = unique_dicts

    for ele in only_auth_keys_list:
        graph.elements.append(ele)

def count_connections_per_node_and_update_label_with_count():
    # Zählen der Verbindungen für jeden Node
    # bzw. einfach für beide identifier die sources und targets zählen und für diesen Node ist das die Insgesamt Count (Verbindungscounter)
    # Erstmal alle Connections holen
    only_sources = []
    only_targets = []
    for graph_element in graph.elements:
        if "source" in graph_element['data']:
                only_sources.append(graph_element['data']['source'])
                only_targets.append(graph_element['data']['target'])
    ges_list = only_sources + only_targets
    ges_counter_dict = Counter(ges_list)

    # now, for every identifier create and set new label (with count number in parethesis)
    for graph_element in graph.elements:
        if 'id' in graph_element['data']:
            graph_element['data']['label'] = graph_element['data']['label'] + " (" + str(ges_counter_dict[graph_element['data']['id']]) + ")"


def count_nodes_and_connections():

    anzahl_machines = 0
    anzahl_auth_keys = 0
    anzahl_connection = 0

    for graph_element in graph.elements:
        if "id" in graph_element['data']:
            if graph_element['classes'] == 'machine':
                anzahl_machines = anzahl_machines + 1

    for graph_element in graph.elements:
        if "id" in graph_element['data']:
            if graph_element['classes'] == 'auth_key':
                anzahl_auth_keys = anzahl_auth_keys + 1

    for graph_element in graph.elements:
        if "source" in graph_element['data']:
                anzahl_connection = anzahl_connection + 1

    return anzahl_machines,anzahl_auth_keys,anzahl_connection


def create_option_menu(anzahl_machines,anzahl_auth_keys,anzahl_connection):

    rescan = html.Div(
        children=[
            html.Button(
                'Rescan',
                id = 'rescan_button',
                n_clicks=0
            )
        ]
    )


    dropdown_update_layout = html.Div(
        children=[
            html.Label('Layout:'),
            dcc.Dropdown(
            id='dropdown_update_layout',
            value='preset',
            clearable=False,
            options=[
                {'label': name.capitalize(), 'value': name}
                for name in ['klay', 'grid', 'random', 'circle', 'cose', 'concentric', 'breadthfirst', 'cose-bilkent', 'cola', 'euler', 'spread', 'dagre']
                ]
            )
        ]
    )

    node_checklist = html.Div(
        children=[
            html.Label('Display:'),
            dcc.Checklist(
            id = 'node_checklist',
            options = [
                {'label': 'Machines', 'value': 'machines'},
                {'label': 'Auth_Keys', 'value': 'auth_keys'},
                {'label': 'Connections', 'value': 'connections'}],
            value = ['machines', 'auth_keys', 'connections']
            )
        ]
    )
    # node_elements = dcc.Checklist(
    #     id = 'node_elements',
    #     options = [
    #         {'label': 'Machines', 'value': 'machines'},
    #         {'label': 'Auth_Keys', 'value': 'auth_keys'},
    #         {'label': 'Connections', 'value': 'connections'}],
    #     value = ['machines', 'auth_keys', 'connections']
    # )

    counter_info = html.Div(
        children=[
            html.Label('Anzahl:'),
            html.Div(
                id='counter_info',
                children=[
                    html.Div(
                        id='anzahl_machines',
                        children=['Machines: ' + str(anzahl_machines)]
                    ),
                    html.Div(
                        id='anzahl_auth_keys',
                        children = ['Auth_Keys: ' + str(anzahl_auth_keys)]
                    ),
                    html.Div(
                        id='anzahl_connections',
                        children = ['Connections: ' + str(anzahl_connection)]
                    )
                ],
            ),
        ]
    )

    optionmenu = html.Div(
        id="optionmenu",
        children=[
            rescan,
            html.Br(),
            dropdown_update_layout,
            html.Br(),
            node_checklist,
            html.Br(),
            counter_info
            # html.Label('Elements:'),
            # node_elements,
        ], style={
            'position': "absolute",
            'top': '0px',
            'right': '0px',
            'z-index': '100',
            'width': '10vw',
            'margin': '0.75vw'
        }
    )

    return optionmenu


def create_callback_dropdown_update_layout():

    @callback(Output('graph', 'layout'),
                    Input('dropdown_update_layout', 'value'))
    def update_layout(layout):
            return {
                'name': layout,
                'animate': True
            }


def create_callback_checkbox_node_display():

    @callback(Output('graph', 'stylesheet'),
                    Input('node_checklist', 'value'))
    def update_node_display(value):
            for graph_style in graph.stylesheet:
                if not 'machines' in value:
                    if graph_style['selector'] == '.machine':
                        style = {'style': {'display': 'none'}}
                        graph_style.update(style)
                if 'machines' in value:
                    if graph_style['selector'] == '.machine':
                        style = {'style': {'content': 'data(label)', 'background-color': 'black'}}
                        graph_style.update(style)
                if not 'auth_keys' in value:
                    if graph_style['selector'] == '.auth_key':
                        style = {'style': {'display': 'none'}}
                        graph_style.update(style)
                if 'auth_keys' in value:
                    if graph_style['selector'] == '.auth_key':
                        style = {'style': {'content': 'data(label)'}}
                        graph_style.update(style)
                if not 'connections' in value:
                    if graph_style['selector'] == '.auth_to_machine':
                        style = {'style': {'display': 'none'}}
                        graph_style.update(style)
                if 'connections' in value:
                    if graph_style['selector'] == '.auth_to_machine':
                        style = {'style': {'content': 'data(label)', 'target-arrow-color': 'black', 'target-arrow-shape': 'triangle'}}
                        graph_style.update(style)
            return(graph.stylesheet)


def create_callback_checkbox_element_update():

    machine_elements = []

    @callback(Output('graph', 'elements'),
                Input('node_elements', 'value'),
                State('graph', 'elements'))
    def update_elements(value, elements):

        if not 'machines' in value:
            for element in elements:
                if "id" in element['data']:
                    if element['classes'] == 'machine':
                        elements.remove(element)
                        machine_elements.append(element)
        
    
        if 'machines' in value:
            for element in elements:
                if "source" in element['data']:
                    pass
                    # print(element)
            # print(machine_elements)
            
            # elements = elements + machine_elements
            # elements = all_elements

        # print(value)

        return elements




###################### MAIN ########################

cyto.load_extra_layouts()
app = dash.Dash(__name__)
graph = create_cytoscape_graph()
scan_keys_and_add_elements()
remove_auth_keys_which_are_machines()
remove_duplicate_auth_keys()
count_connections_per_node_and_update_label_with_count()
anzahl_machines,anzahl_auth_keys,anzahl_connection = count_nodes_and_connections()
optionmenu = create_option_menu(anzahl_machines,anzahl_auth_keys,anzahl_connection)
create_callback_dropdown_update_layout()
create_callback_checkbox_node_display()
# create_callback_checkbox_element_update()


@callback(Output('graph', 'elements'),
          Output('counter_info', 'children'),
              Input('rescan_button', 'n_clicks'))
def update_elements(n_clicks):
    graph.elements = []
    # optionmenu
    scan_keys_and_add_elements()
    remove_auth_keys_which_are_machines()
    remove_duplicate_auth_keys()
    count_connections_per_node_and_update_label_with_count()
    anzahl_machines,anzahl_auth_keys,anzahl_connection = count_nodes_and_connections()
    children=[
    html.Div(
        id='anzahl_machines',
        children=['Machines: ' + str(anzahl_machines)]
    ),
    html.Div(
        id='anzahl_auth_keys',
        children = ['Auth_Keys: ' + str(anzahl_auth_keys)]
    ),
    html.Div(
        id='anzahl_connections',
        children = ['Connections: ' + str(anzahl_connection)]
    )
    ]

    return graph.elements, children




app.layout = html.Div([
    graph,
    optionmenu
])

if __name__ == '__main__':
    app.run_server(debug=False, host="0.0.0.0")