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

167 lines
5.2 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
from fabric import Connection
from app.exceptions import NetworkError
from invoke.exceptions import UnexpectedExit
import paramiko
import time
import socket
from plugins.base.plugin_base import BasePlugin
class SSHFeatures(BasePlugin):
def __init__(self):
super().__init__()
self.c = None
def connect(self):
""" Connect to a machine """
if self.c:
return self.c
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)
self.c = 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(uhp, 3)
self.c = 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:
self.vprint(f"Connection: {self.c}", 3)
return self.c
self.vprint("SSH network error", 0)
raise NetworkError
def remote_run(self, cmd, disown=False):
""" Connects to the machine and runs a command there
@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)
result = None
retry = 2
while retry > 0:
try:
result = self.c.run(cmd, disown=disown)
print(result)
# paramiko.ssh_exception.SSHException in the next line is needed for windows openssh
except (paramiko.ssh_exception.NoValidConnectionsError, UnexpectedExit, paramiko.ssh_exception.SSHException):
if retry <= 0:
raise NetworkError
else:
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 ""
def put(self, src, dst):
""" 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:
try:
res = self.c.put(src, dst)
except (paramiko.ssh_exception.SSHException, socket.timeout, UnexpectedExit):
self.vprint(f"PUT Failed to connect, will retry {retries} times. Timeout: {timeout}", 3)
retries -= 1
timeout += 10
time.sleep(retry_sleep)
self.disconnect()
self.connect()
except FileNotFoundError as e:
self.vprint(f"File not found: {e}", 0)
break
else:
return res
self.vprint("SSH network error on PUT command", 0)
raise NetworkError
def get(self, src, dst):
""" Get a file to a machine
@param src: source dir
@param dst: destination
"""
self.connect()
if os.path.isdir(dst):
dst = os.path.join(dst, os.path.basename(src))
retry = 2
while retry > 0:
try:
res = self.c.get(src, dst)
except (paramiko.ssh_exception.NoValidConnectionsError, UnexpectedExit):
if retry <= 0:
raise NetworkError
else:
self.disconnect()
self.connect()
retry -= 1
self.vprint("Got some SSH errors. Retrying", 2)
except FileNotFoundError as e:
self.vprint(e, 0)
break
else:
break
return res
def disconnect(self):
""" Disconnect from a machine """
if self.c:
self.c.close()
self.c = None