diff --git a/openct/__main__.py b/openct/__main__.py index 1079460..a291682 100644 --- a/openct/__main__.py +++ b/openct/__main__.py @@ -20,15 +20,7 @@ logging.info("Backing up device %s", host) pbar.update(100 / len(devices)) - device = RouterOSDevice( - host, - config.identity.username, - config.settings.connection_timeout, - config.identity.key_file, - config.dirs.backup_dir, - ) + device = RouterOSDevice(host, config) if device.is_available(): device.fetch_backup() - else: - logging.error("Could not connect to device %s", host) diff --git a/openct/connections/connections.py b/openct/connections/connections.py index b7d2d16..42e02c2 100644 --- a/openct/connections/connections.py +++ b/openct/connections/connections.py @@ -4,8 +4,15 @@ import logging from fabric import Connection as FabricConnection -from paramiko.ssh_exception import SSHException, NoValidConnectionsError -from invoke import UnexpectedExit +from paramiko.ssh_exception import ( + SSHException, + NoValidConnectionsError, + BadHostKeyException, + AuthenticationException, +) +from invoke import UnexpectedExit, Failure + +from openct.setup import Config class DeviceConnection(Protocol): @@ -21,19 +28,12 @@ def fetch_backup(self) -> None: class SshConnection(DeviceConnection): """DeviceConnection using Fabric""" - def __init__( - self, - ip_address: str, - username: str, - connection_timeout: int, - key_file: str, - backup_dir: str, - ) -> None: + def __init__(self, ip_address: str, config: Config) -> None: self.ip_address = ip_address - self.username = username - self.connection_timeout = connection_timeout - self.key_file = key_file - self.backup_dir = backup_dir + self.username = config.identity.username + self.connection_timeout = config.settings.connection_timeout + self.key_file = config.identity.key_file + self.backup_dir = config.dirs.backup_dir def test_connection(self) -> bool: with FabricConnection( @@ -45,8 +45,15 @@ def test_connection(self) -> bool: try: connection.open() return True - except (TimeoutError, SSHException, NoValidConnectionsError): + except ( + TimeoutError, + BadHostKeyException, + AuthenticationException, + SSHException, + NoValidConnectionsError, + ) as error: logging.error("Could not connect to device %s", self.ip_address) + logging.info(error) return False def fetch_backup(self) -> None: @@ -57,12 +64,23 @@ def fetch_backup(self) -> None: connect_kwargs={"key_filename": self.key_file}, ) as connection: try: - connection.run("/export file=backup", hide=True, warn=False) + connection.run( + "/export file=backup", + hide=True, + warn=False, + timeout=self.connection_timeout, + ) connection.get( "backup.rsc", f"{self.backup_dir}/backup_{self.ip_address}.rsc" ) - connection.run("file/remove backup.rsc", hide=True, warn=False) - except UnexpectedExit: + connection.run( + "file/remove backup.rsc", + hide=True, + warn=False, + timeout=self.connection_timeout, + ) + except (UnexpectedExit, Failure) as error: logging.error( "Error while fetching backup from device %s", self.ip_address ) + logging.info(error) diff --git a/openct/devices/devices.py b/openct/devices/devices.py index e4fb34c..626d79a 100644 --- a/openct/devices/devices.py +++ b/openct/devices/devices.py @@ -3,6 +3,7 @@ from abc import ABC, abstractmethod from openct.connections import DeviceConnection, SshConnection +from openct.setup import Config class Device(ABC): @@ -22,19 +23,8 @@ def is_available(self) -> bool: class RouterOSDevice(Device): """RouterOS device""" - def __init__( - self, - ip_address: str, - username: str, - connection_timeout: int, - key_file: str, - backup_dir: str, - ) -> None: - super().__init__( - SshConnection( - ip_address, username, connection_timeout, key_file, backup_dir - ) - ) + def __init__(self, ip_address: str, config: Config) -> None: + super().__init__(SshConnection(ip_address, config)) def fetch_backup(self) -> str: return self.connection.fetch_backup() diff --git a/pyproject.toml b/pyproject.toml index 6eb30de..ba2cbf3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "openct" -version = "0.1.3" +version = "0.1.4" description = "Configuration backup and analysis tools for devices running pfSense and RouterOS" authors = ["Weehooey "] license = "GNU GPL v3.0"