mirror of
https://github.com/kernelkit/infix.git
synced 2026-07-26 10:43:02 +02:00
106 lines
3.3 KiB
Python
106 lines
3.3 KiB
Python
import subprocess
|
|
|
|
from dataclasses import dataclass
|
|
from . import env
|
|
|
|
|
|
@dataclass
|
|
class Location:
|
|
host: str
|
|
username: str
|
|
password: str
|
|
port: int = 22
|
|
|
|
|
|
import subprocess
|
|
import os
|
|
|
|
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):
|
|
self.name = name
|
|
self.location = location
|
|
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)
|