mirror of
https://github.com/kernelkit/infix.git
synced 2026-07-28 03:33:01 +02:00
Tobias Waldekranz
0298dc3c1a
By default, ensure that the target device accepts TCP connections on
the specified port before returning the device object. This aligns the
SSH implementation with {NET,REST}CONF.
This solves the issue where tests that use SSH commands early in their
execution would sometimes fail as the remote server had not started up
yet.
112 lines
3.4 KiB
Python
112 lines
3.4 KiB
Python
import subprocess
|
|
|
|
from dataclasses import dataclass
|
|
from . import env, netutil, util
|
|
|
|
|
|
@dataclass
|
|
class Location:
|
|
host: str
|
|
username: str
|
|
password: str
|
|
port: int = 22
|
|
|
|
|
|
import subprocess
|
|
import os
|
|
|
|
def ssh_syn(addr, port=22):
|
|
return netutil.tcp_port_is_open(addr, port)
|
|
|
|
def fetch_file(remote_user, remote_address, remote_file, local_file, key_file, check=False, remove=False):
|
|
"""
|
|
Fetches a file over SSH using scp and the provided private key.
|
|
|
|
:param remote_user: The user on the remote machine.
|
|
:param remote_address: The address of the remote machine.
|
|
:param remote_file: The file to fetch from the remote machine.
|
|
:param local_file: The local path where the file will be stored.
|
|
:param key_file: The path to the private SSH key.
|
|
:param check: check the return code of the command.
|
|
:param remove: remove the fetched local file after copying.
|
|
"""
|
|
try:
|
|
result = subprocess.run(
|
|
f"scp -q -o StrictHostKeyChecking=no -i {key_file} {remote_user}@[{remote_address}]:{remote_file} {local_file}",
|
|
shell=True
|
|
)
|
|
|
|
if check:
|
|
if result.returncode != 0:
|
|
raise RuntimeError("Failed to copy file from remote host")
|
|
|
|
if not os.path.exists(local_file):
|
|
raise RuntimeError(f"File {local_file} does not exist after copy")
|
|
|
|
if os.path.getsize(local_file) == 0:
|
|
raise RuntimeError(f"File {local_file} is empty after copy")
|
|
|
|
except Exception as e:
|
|
print(f"Error during file transfer: {e}")
|
|
raise
|
|
|
|
finally:
|
|
if os.path.exists(key_file):
|
|
try:
|
|
os.remove(key_file)
|
|
except OSError as e:
|
|
print(f"Error removing key file {key_file}: {e}")
|
|
|
|
if remove:
|
|
try:
|
|
if os.path.exists(local_file):
|
|
os.remove(local_file)
|
|
except OSError as e:
|
|
print(f"Error removing fetched file {local_file}: {e}")
|
|
|
|
class Device(object):
|
|
def __init__(self, name: str, location: Location, wait: bool=True):
|
|
self.name = name
|
|
self.location = location
|
|
if wait:
|
|
util.until(lambda: ssh_syn(location.host, location.port))
|
|
|
|
def __str__(self):
|
|
nm = f"{self.name}"
|
|
if env.ENV.ltop:
|
|
nm += f"({env.ENV.ltop.xlate(self.name)})"
|
|
return nm + " [SSH]"
|
|
|
|
def _mangle_subprocess_args(self, args, kwargs):
|
|
if not args:
|
|
return None
|
|
|
|
args = list(args)
|
|
if type(args[0]) is str:
|
|
if kwargs.get("shell"):
|
|
args[0] = ["/bin/sh", "-c", args[0]]
|
|
kwargs["shell"] = False
|
|
else:
|
|
args[0] = [args[0]]
|
|
|
|
args[0] = ["ssh",
|
|
"-oStrictHostKeyChecking no",
|
|
"-oUserKnownHostsFile /dev/null",
|
|
"-oLogLevel QUIET",
|
|
f"-l{self.location.username}",
|
|
self.location.host] + args[0]
|
|
|
|
if self.location.password:
|
|
args[0] = ["sshpass", f"-p{self.location.password}"] + args[0]
|
|
|
|
return args, kwargs
|
|
|
|
def run(self, *args, **kwargs):
|
|
args, kwargs = self._mangle_subprocess_args(args, kwargs)
|
|
return subprocess.run(*args, **kwargs)
|
|
|
|
def runsh(self, script, *args, **kwargs):
|
|
return self.run("/bin/sh", text=True, input=script,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.STDOUT, *args, **kwargs)
|