You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
PurpleDome/plugins/base/ssh_features.py

193 lines
6.5 KiB
Python

#!/usr/bin/env python3
""" A class you can use to add SSH features to you plugin. Useful for vm_controller/machinery classes """
import os.path
3 years ago
import socket
import time
import paramiko
from fabric import Connection
from invoke.exceptions import UnexpectedExit
3 years ago
from app.exceptions import NetworkError
from plugins.base.plugin_base import BasePlugin
class SSHFeatures(BasePlugin):
3 years ago
""" A Mixin class to add SSH features to all kind of VM machinery """
def __init__(self):
3 years ago
self.config = None
super().__init__()
3 years ago
self.connection = None
3 years ago
def get_ip(self):
""" Get the IP of a machine, must be overwritten in the machinery class """
raise NotImplementedError
def connect(self):
""" Connect to a machine """
3 years ago
if self.connection is not None:
return self.connection
retries = 10
retry_sleep = 10
timeout = 30
while retries:
try:
if self.config.os() == "linux":
uhp = self.get_ip()
self.vprint(f"Connecting to {uhp}", 3)
3 years ago
self.connection = Connection(uhp, connect_timeout=timeout)
if self.config.os() == "windows":
args = {}
# args = {"key_filename": self.config.ssh_keyfile() or self.v.keyfile(vm_name=self.config.vmname())}
if self.config.ssh_keyfile():
args["key_filename"] = self.config.ssh_keyfile()
if self.config.ssh_password():
args["password"] = self.config.ssh_password()
self.vprint(args, 3)
uhp = self.get_ip()
self.vprint(f"IP to connect to: {uhp}", 3)
3 years ago
self.connection = Connection(uhp, connect_timeout=timeout, user=self.config.ssh_user(), connect_kwargs=args)
except (paramiko.ssh_exception.SSHException, socket.timeout):
self.vprint(f"Failed to connect, will retry {retries} times. Timeout: {timeout}", 0)
retries -= 1
timeout += 10
time.sleep(retry_sleep)
else:
3 years ago
self.vprint(f"Connection: {self.connection}", 3)
return self.connection
self.vprint("SSH network error", 0)
raise NetworkError
3 years ago
def remote_run(self, cmd: str, disown: bool = False):
""" Connects to the machine and runs a command there
@param cmd: The command to execute
@param disown: Send the connection into background
"""
if cmd is None:
return ""
self.connect()
cmd = cmd.strip()
self.vprint("Running SSH remote run: " + cmd, 3)
self.vprint("Disown: " + str(disown), 3)
3 years ago
# self.vprint("Connection: " + self.connection, 1)
result = None
retry = 2
while retry > 0:
do_retry = False
try:
3 years ago
result = self.connection.run(cmd, disown=disown)
print(result)
# paramiko.ssh_exception.SSHException in the next line is needed for windows openssh
3 years ago
except (paramiko.ssh_exception.NoValidConnectionsError, UnexpectedExit, paramiko.ssh_exception.SSHException) as error:
if retry <= 0:
3 years ago
raise NetworkError from error
do_retry = True
except paramiko.ssh_exception.NoValidConnectionsError as error:
self.vprint(f"No valid connection. Errors: {error.errors}", 1)
do_retry = True
if do_retry:
3 years ago
self.disconnect()
self.connect()
retry -= 1
self.vprint("Got some SSH errors. Retrying", 2)
else:
break
if result and result.stderr:
self.vprint("Debug: Stderr: " + str(result.stderr.strip()), 0)
if result:
return result.stdout.strip()
return ""
3 years ago
def put(self, src: str, dst: str):
""" Send a file to a machine
@param src: source dir
@param dst: destination
"""
self.connect()
self.vprint(f"PUT {src} -> {dst}", 3)
res = ""
retries = 10
retry_sleep = 10
timeout = 30
while retries > 0:
do_retry = False
try:
3 years ago
res = self.connection.put(src, dst)
except (paramiko.ssh_exception.SSHException, socket.timeout, UnexpectedExit):
self.vprint("PUT Failed to connect", 1)
do_retry = True
except paramiko.ssh_exception.NoValidConnectionsError as error:
self.vprint(f"No valid connection. Errors: {error.errors}", 1)
do_retry = True
except FileNotFoundError as error:
self.vprint(f"File not found: {error}", 0)
break
if do_retry:
self.vprint(f"Will retry {retries} times. Timeout: {timeout}", 3)
retries -= 1
timeout += 10
time.sleep(retry_sleep)
self.disconnect()
self.connect()
else:
return res
self.vprint("SSH network error on PUT command", 0)
raise NetworkError
3 years ago
def get(self, src: str, dst: str):
""" Get a file to a machine
@param src: source dir
@param dst: destination
"""
self.connect()
res = None
if os.path.isdir(dst):
dst = os.path.join(dst, os.path.basename(src))
retry = 2
while retry > 0:
do_retry = False
try:
3 years ago
res = self.connection.get(src, dst)
3 years ago
except (paramiko.ssh_exception.NoValidConnectionsError, UnexpectedExit) as error:
if retry <= 0:
3 years ago
raise NetworkError from error
do_retry = True
except paramiko.ssh_exception.NoValidConnectionsError as error:
self.vprint(f"No valid connection. Errors: {error.errors}", 1)
do_retry = True
except FileNotFoundError as error:
self.vprint(error, 0)
break
if do_retry:
3 years ago
self.disconnect()
self.connect()
retry -= 1
self.vprint("Got some SSH errors. Retrying", 2)
else:
break
return res
def disconnect(self):
""" Disconnect from a machine """
3 years ago
if self.connection:
self.connection.close()
self.connection = None