61 lines
No EOL
2.2 KiB
Python
61 lines
No EOL
2.2 KiB
Python
import subprocess
|
|
import os
|
|
import sys
|
|
|
|
class SSHClient:
|
|
def __init__(self, host, user="root", key_path=None, password=None, timeout=30):
|
|
self.host = host
|
|
self.user = user
|
|
self.key_path = os.path.expanduser(key_path) if key_path else None
|
|
self.password = password
|
|
self.timeout = timeout
|
|
|
|
def run(self, cmd, capture=True):
|
|
ssh_cmd = [
|
|
"ssh",
|
|
"-o", "StrictHostKeyChecking=no",
|
|
"-o", "UserKnownHostsFile=/dev/null",
|
|
"-o", f"ConnectTimeout={self.timeout}",
|
|
"-o", "BatchMode=yes" if not self.password else "BatchMode=no"
|
|
]
|
|
|
|
if self.key_path:
|
|
ssh_cmd += ["-i", self.key_path]
|
|
|
|
target = f"{self.user}@{self.host}"
|
|
|
|
if self.password:
|
|
# sshpass is required for password auth
|
|
full_cmd = ["sshpass", "-p", self.password] + ssh_cmd + [target, cmd]
|
|
else:
|
|
full_cmd = ssh_cmd + [target, cmd]
|
|
|
|
try:
|
|
result = subprocess.run(
|
|
full_cmd,
|
|
capture_output=capture,
|
|
text=True,
|
|
timeout=self.timeout + 10
|
|
)
|
|
return result
|
|
except subprocess.TimeoutExpired:
|
|
print(f"Error: SSH command timed out after {self.timeout}s on {self.host}", file=sys.stderr)
|
|
# Create a mock result for timeout
|
|
return subprocess.CompletedProcess(full_cmd, 1, "", "Timeout expired")
|
|
except Exception as e:
|
|
print(f"Error: SSH execution failed on {self.host}: {e}", file=sys.stderr)
|
|
return subprocess.CompletedProcess(full_cmd, 1, "", str(e))
|
|
|
|
def scp_to(self, local_path, remote_path):
|
|
scp_cmd = ["scp", "-o", "StrictHostKeyChecking=no", "-o", "UserKnownHostsFile=/dev/null"]
|
|
if self.key_path:
|
|
scp_cmd += ["-i", self.key_path]
|
|
|
|
target = f"{self.user}@{self.host}:{remote_path}"
|
|
|
|
if self.password:
|
|
full_cmd = ["sshpass", "-p", self.password] + scp_cmd + [local_path, target]
|
|
else:
|
|
full_cmd = scp_cmd + [local_path, target]
|
|
|
|
return subprocess.run(full_cmd, capture_output=True, timeout=self.timeout + 60) |