From 106c3c8cd8059844f446bc92b0c8efebccec59e6 Mon Sep 17 00:00:00 2001 From: Micky <60691199+AliMickey@users.noreply.github.com> Date: Thu, 26 Dec 2024 13:46:44 +1100 Subject: [PATCH] implement ssh key support --- app/functions/netmiko.py | 77 ++++++++++++++++++++++++++++++++++++++++ app/functions/utils.py | 64 +-------------------------------- app/views/main.py | 3 +- 3 files changed, 80 insertions(+), 64 deletions(-) create mode 100644 app/functions/netmiko.py diff --git a/app/functions/netmiko.py b/app/functions/netmiko.py new file mode 100644 index 0000000..a621fec --- /dev/null +++ b/app/functions/netmiko.py @@ -0,0 +1,77 @@ +import logging, os +from netmiko import ConnectHandler, NetmikoTimeoutException, NetmikoAuthenticationException + +logger = logging.getLogger(__name__) + +# Establish connection to network device +def establish_connection(device_config): + try: + return ConnectHandler(**device_config) + except Exception as e: + logger.error(f"Failed to establish connection to {device_config['host']}: {e}") + raise + + +# Execute command on network device +def execute_command(device, command_format, target, ip_version): + device_config = { + 'device_type': device['type'], + 'host': device['host'], + 'port': device['port'], + 'username': device['username'], + 'timeout': 10, + 'session_timeout': 60, + 'conn_timeout': 10, + 'auth_timeout': 10, + } + + # Use SSH key if provided + if "ssh_key" in device: + key_path = os.path.join("/instance/ssh-keys", device['ssh_key']) + + if not os.path.exists(key_path): + logger.error(f"SSH file not found: {key_path} for {device['host']}") + return {'error': True, 'message': 'Authentication failed'} + + device_config['use_keys'] = True + device_config['key_file'] = os.path.join("/instance/ssh-keys", device['ssh_key']) + + else: + device_config['password'] = device['password'] + + command_timeout = 30 + + try: + with establish_connection(device_config) as connection: + # Format the command + command = str(command_format.format(ip_version=ip_version, target=target).strip()) + + # Execute the command + output = connection.send_command( + command, + read_timeout=command_timeout, + strip_command=True, + strip_prompt=True, + max_loops=int(command_timeout * 10) + ) + + # Clean output + output = output.strip() if output else "" + + if not output: + logger.error(f"No response from {device['host']}") + return {'error': True, 'message': 'No response from device'} + + return {'error': False, 'message': output} + + except NetmikoTimeoutException as e: + logger.error(f"Timeout error on {device['host']}: {e}") + return {'error': True, 'message': 'Timeout error'} + + except NetmikoAuthenticationException as e: + logger.error(f"Authentication failed for {device['host']}: {e}") + return {'error': True, 'message': 'Authentication failed'} + + except Exception as e: + logger.error(f"An unexpected error occurred on {device['host']}: {e}") + return {'error': True, 'message': 'Unexpected error'} diff --git a/app/functions/utils.py b/app/functions/utils.py index 6dd1930..0d108bb 100644 --- a/app/functions/utils.py +++ b/app/functions/utils.py @@ -1,6 +1,5 @@ import os, yaml, logging, requests from functools import wraps -from netmiko import ConnectHandler, NetmikoTimeoutException, NetmikoAuthenticationException from flask import request logger = logging.getLogger(__name__) @@ -70,65 +69,4 @@ def send_webhook(webhook_url, text_data): def get_client_ip(): if not request.headers.getlist("X-Forwarded-For"): return request.remote_addr - return request.headers.getlist("X-Forwarded-For")[0] - - -# Establish connection to network device -def establish_connection(device_config): - try: - return ConnectHandler(**device_config) - except Exception as e: - logger.error(f"Failed to establish connection to {device_config['host']}: {e}") - raise - - -# Execute command on network device -def execute_command(device, command_format, target, ip_version): - device_config = { - 'device_type': device['type'], - 'host': device['host'], - 'port': device['port'], - 'username': device['username'], - 'password': device['password'], - 'timeout': 10, - 'session_timeout': 60, - 'conn_timeout': 10, - 'auth_timeout': 10, - } - - command_timeout = 30 - - try: - with establish_connection(device_config) as connection: - # Format the command - command = str(command_format.format(ip_version=ip_version, target=target).strip()) - - # Execute the command - output = connection.send_command( - command, - read_timeout=command_timeout, - strip_command=True, - strip_prompt=True, - max_loops=int(command_timeout * 10) - ) - - # Clean output - output = output.strip() if output else "" - - if not output: - logger.error(f"No response from {device['host']}") - return {'error': True, 'message': 'No response from device'} - - return {'error': False, 'message': output} - - except NetmikoTimeoutException as e: - logger.error(f"Timeout error on {device['host']}: {e}") - return {'error': True, 'message': 'Timeout error'} - - except NetmikoAuthenticationException as e: - logger.error(f"Authentication failed for {device['host']}: {e}") - return {'error': True, 'message': 'Authentication failed'} - - except Exception as e: - logger.error(f"An unexpected error occurred on {device['host']}: {e}") - return {'error': True, 'message': 'Unexpected error'} + return request.headers.getlist("X-Forwarded-For")[0] \ No newline at end of file diff --git a/app/views/main.py b/app/views/main.py index c5d8ec8..7c95b65 100644 --- a/app/views/main.py +++ b/app/views/main.py @@ -2,7 +2,8 @@ import logging from flask import Blueprint, request, render_template -from app.functions.utils import exception_handler, load_yaml, send_webhook, get_client_ip, execute_command +from app.functions.utils import exception_handler, load_yaml, send_webhook, get_client_ip +from app.functions.netmiko import execute_command logger = logging.getLogger(__name__)