mirror of
https://github.com/aaPanel/aaPanel.git
synced 2026-08-17 21:25:47 +02:00
809 lines
30 KiB
Python
809 lines
30 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import stat
|
|
from typing import Optional, Tuple, Callable, Union, Dict, Any, Iterator
|
|
import time
|
|
import io
|
|
import math
|
|
|
|
import paramiko
|
|
from dataclasses import dataclass
|
|
from .rate_limiter import RateLimiter, ProgressTracker
|
|
|
|
|
|
@dataclass
|
|
class CommandResult:
|
|
exit_code: int
|
|
stdout: str
|
|
stderr: str
|
|
|
|
|
|
class SSHExecutor:
|
|
"""
|
|
High-level SSH executor wrapping Paramiko for command execution and SFTP upload.
|
|
|
|
Usage:
|
|
with SSHExecutor(host, user, password=...) as ssh:
|
|
code, out, err = ssh.run("uname -a")
|
|
ssh.upload("./local.txt", "/tmp/remote.txt")
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
host: str,
|
|
username: str,
|
|
port: int = 22,
|
|
password: Optional[str] = None,
|
|
key_file: Optional[str] = None,
|
|
passphrase: Optional[str] = None,
|
|
key_data: Optional[str] = None,
|
|
timeout: Optional[int] = None,
|
|
strict_host_key_checking: bool = False,
|
|
allow_agent: bool = False,
|
|
look_for_keys: bool = False,
|
|
threading_mod: bool = False, # 线程模式，默认为False，当线程模式，会在每次获取sftp客户端时重新获取，
|
|
) -> None:
|
|
self.host = host
|
|
self.port = port
|
|
self.username = username
|
|
self.password = password
|
|
self.key_file = key_file
|
|
self.passphrase = passphrase
|
|
self.key_data = key_data
|
|
self.timeout = timeout or 20
|
|
self.strict_host_key_checking = strict_host_key_checking
|
|
self.allow_agent = allow_agent
|
|
self.look_for_keys = look_for_keys
|
|
self._client: Optional[paramiko.SSHClient] = None
|
|
self._sftp: Optional[paramiko.SFTPClient] = None
|
|
self._threading_mod = threading_mod
|
|
|
|
def open(self) -> None:
|
|
if self._client is not None:
|
|
return
|
|
client = paramiko.SSHClient()
|
|
if self.strict_host_key_checking:
|
|
client.set_missing_host_key_policy(paramiko.RejectPolicy())
|
|
else:
|
|
client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
|
|
|
|
connect_kwargs: Dict[str, Any] = {
|
|
"hostname": self.host,
|
|
"port": self.port,
|
|
"username": self.username,
|
|
"look_for_keys": self.look_for_keys,
|
|
"allow_agent": self.allow_agent,
|
|
}
|
|
if self.timeout is not None:
|
|
connect_kwargs.update({
|
|
"timeout": self.timeout,
|
|
"banner_timeout": self.timeout,
|
|
"auth_timeout": self.timeout,
|
|
})
|
|
if self.password:
|
|
connect_kwargs["password"] = self.password
|
|
|
|
if self.key_file or self.key_data:
|
|
pkey = self._load_private_key(self.key_file, self.key_data, self.passphrase)
|
|
connect_kwargs["pkey"] = pkey
|
|
|
|
try:
|
|
client.connect(**connect_kwargs)
|
|
except Exception as e:
|
|
client.close()
|
|
raise RuntimeError(f"SSH connection failed: {e}")
|
|
|
|
self._client = client
|
|
|
|
def close(self) -> None:
|
|
if self._sftp is not None:
|
|
try:
|
|
self._sftp.close()
|
|
finally:
|
|
self._sftp = None
|
|
if self._client is not None:
|
|
try:
|
|
self._client.close()
|
|
finally:
|
|
self._client = None
|
|
|
|
def __enter__(self) -> "SSHExecutor":
|
|
self.open()
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc, tb) -> None:
|
|
self.close()
|
|
|
|
def run(self, command: str, timeout: Optional[int] = None) -> Tuple[int, str, str]:
|
|
client = self._require_client()
|
|
try:
|
|
effective_timeout = timeout if timeout is not None else self.timeout
|
|
stdin, stdout, stderr = client.exec_command(command, timeout=effective_timeout)
|
|
exit_status = stdout.channel.recv_exit_status()
|
|
out = stdout.read().decode("utf-8", errors="replace")
|
|
err = stderr.read().decode("utf-8", errors="replace")
|
|
return exit_status, out, err
|
|
except Exception as e:
|
|
raise RuntimeError(f"Command execution failed: {e}")
|
|
|
|
def upload(
|
|
self,
|
|
local_path: str,
|
|
remote_path: str,
|
|
rate_limit: Optional[int] = None,
|
|
progress_callback: Optional[Callable[[int, int], None]] = None,
|
|
resume: bool = False,
|
|
rate_algorithm: str = "token_bucket",
|
|
burst_capacity: Optional[int] = None,
|
|
rate_limiter: Optional[RateLimiter] = None,
|
|
) -> None:
|
|
"""
|
|
Upload a file via SFTP with optional rate limiting and resume support.
|
|
|
|
Args:
|
|
local_path: Local file path
|
|
remote_path: Remote destination path
|
|
rate_limit: Rate limit in bytes per second (None = no limit)
|
|
progress_callback: Callback(transferred_bytes, total_bytes) for progress updates
|
|
resume: Whether to resume upload if remote file exists and is smaller.
|
|
WARNING: Only checks file size, no content verification. Use with caution.
|
|
rate_algorithm: Rate limiting algorithm ("token_bucket" or "leaky_bucket")
|
|
burst_capacity: For token bucket, maximum burst capacity in bytes
|
|
rate_limiter: External RateLimiter instance for unified control across multiple transfers
|
|
"""
|
|
expanded_local = os.path.expanduser(local_path)
|
|
if not os.path.isfile(expanded_local):
|
|
raise FileNotFoundError(f"Local file not found: {expanded_local}")
|
|
|
|
local_size = os.path.getsize(expanded_local)
|
|
resume_offset = 0
|
|
|
|
if resume:
|
|
sftp = None
|
|
try:
|
|
sftp = self.get_sftp()
|
|
remote_stat = sftp.stat(remote_path)
|
|
if remote_stat.st_size < local_size:
|
|
resume_offset = remote_stat.st_size
|
|
if progress_callback:
|
|
progress_callback(resume_offset, local_size)
|
|
except FileNotFoundError:
|
|
pass
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
if self._threading_mod and sftp:
|
|
sftp.close()
|
|
|
|
sftp = None
|
|
try:
|
|
sftp = self.get_sftp()
|
|
# Use external rate limiter if provided, otherwise create new one
|
|
if rate_limiter is None:
|
|
rate_limiter = RateLimiter(rate_limit, rate_algorithm, burst_capacity)
|
|
progress_tracker = ProgressTracker(progress_callback)
|
|
|
|
# Use chunked transfer for better control
|
|
self._upload_chunked(sftp, expanded_local, remote_path, resume_offset, rate_limiter, progress_tracker)
|
|
|
|
except Exception as e:
|
|
raise RuntimeError(f"SFTP upload failed: {e}")
|
|
finally:
|
|
if self._threading_mod and sftp:
|
|
sftp.close()
|
|
|
|
def _ensure_remote_dir(self, sftp, path):
|
|
try:
|
|
sftp.stat(path)
|
|
except FileNotFoundError:
|
|
self._create_remote_dir_recursive(sftp, path)
|
|
|
|
@staticmethod
|
|
def _create_remote_dir_recursive(sftp, path):
|
|
dirs = [d for d in path.split('/') if d]
|
|
current = ''
|
|
for d in dirs:
|
|
current += '/' + d
|
|
try:
|
|
sftp.stat(current)
|
|
except FileNotFoundError:
|
|
sftp.mkdir(current)
|
|
|
|
def _upload_chunked(
|
|
self,
|
|
sftp: paramiko.SFTPClient,
|
|
local_path: str,
|
|
remote_path: str,
|
|
resume_offset: int,
|
|
rate_limiter: Optional[RateLimiter] = None,
|
|
progress_tracker: Optional[ProgressTracker] = None,
|
|
) -> None:
|
|
"""Upload file in chunks with optional rate limiting and progress tracking."""
|
|
local_size = os.path.getsize(local_path)
|
|
chunk_size = 32768 # 32KB chunks
|
|
transferred = resume_offset
|
|
|
|
# Initialize components if provided
|
|
if rate_limiter:
|
|
rate_limiter.start()
|
|
if progress_tracker:
|
|
progress_tracker.start()
|
|
|
|
with open(local_path, "rb") as local_file:
|
|
if resume_offset > 0:
|
|
local_file.seek(resume_offset)
|
|
|
|
# 检查目录是否存在，不存在则创建
|
|
self._ensure_remote_dir(sftp, os.path.dirname(remote_path))
|
|
|
|
with sftp.file(remote_path, "ab" if resume_offset > 0 else "wb") as remote_file:
|
|
while transferred < local_size:
|
|
chunk = local_file.read(chunk_size)
|
|
if not chunk:
|
|
break
|
|
|
|
# Apply rate limiting before transfer
|
|
if rate_limiter:
|
|
rate_limiter.wait_if_needed(len(chunk))
|
|
|
|
# Perform the actual transfer
|
|
remote_file.write(chunk)
|
|
transferred += len(chunk)
|
|
|
|
# Update progress after transfer
|
|
if progress_tracker:
|
|
progress_tracker.update(transferred, local_size)
|
|
|
|
# Final progress update
|
|
if progress_tracker:
|
|
progress_tracker.finish(local_size)
|
|
|
|
def download(
|
|
self,
|
|
remote_path: str,
|
|
local_path: str,
|
|
rate_limit: Optional[int] = None,
|
|
progress_callback: Optional[Callable[[int, int], None]] = None,
|
|
resume: bool = False,
|
|
rate_algorithm: str = "token_bucket",
|
|
burst_capacity: Optional[int] = None,
|
|
rate_limiter: Optional[RateLimiter] = None,
|
|
) -> None:
|
|
"""
|
|
Download a file via SFTP with optional rate limiting and resume support.
|
|
|
|
Args:
|
|
remote_path: Remote file path
|
|
local_path: Local destination path
|
|
rate_limit: Rate limit in bytes per second (None = no limit)
|
|
progress_callback: Callback(transferred_bytes, total_bytes) for progress updates
|
|
resume: Whether to resume download if local file exists and is smaller.
|
|
WARNING: Only checks file size, no content verification. Use with caution.
|
|
rate_algorithm: Rate limiting algorithm ("token_bucket" or "leaky_bucket")
|
|
burst_capacity: For token bucket, maximum burst capacity in bytes
|
|
rate_limiter: External RateLimiter instance for unified control across multiple transfers
|
|
"""
|
|
expanded_local = os.path.expanduser(local_path)
|
|
resume_offset = 0
|
|
|
|
sftp = None
|
|
if resume and os.path.exists(expanded_local):
|
|
local_size = os.path.getsize(expanded_local)
|
|
try:
|
|
sftp = self.get_sftp()
|
|
remote_stat = sftp.stat(remote_path)
|
|
if local_size < remote_stat.st_size:
|
|
resume_offset = local_size
|
|
if progress_callback:
|
|
progress_callback(resume_offset, remote_stat.st_size)
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
if self._threading_mod and sftp:
|
|
sftp.close()
|
|
|
|
try:
|
|
sftp = self.get_sftp()
|
|
# Use external rate limiter if provided, otherwise create new one
|
|
if rate_limiter is None:
|
|
rate_limiter = RateLimiter(rate_limit, rate_algorithm, burst_capacity)
|
|
progress_tracker = ProgressTracker(progress_callback)
|
|
|
|
# Use chunked transfer for better control
|
|
self._download_chunked(sftp, remote_path, expanded_local, resume_offset, rate_limiter, progress_tracker)
|
|
|
|
except Exception as e:
|
|
raise RuntimeError(f"SFTP download failed: {e}")
|
|
finally:
|
|
if self._threading_mod and sftp:
|
|
sftp.close()
|
|
|
|
@staticmethod
|
|
def _download_chunked(
|
|
sftp,
|
|
remote_path: str,
|
|
local_path: str,
|
|
resume_offset: int,
|
|
rate_limiter: Optional[RateLimiter] = None,
|
|
progress_tracker: Optional[ProgressTracker] = None,
|
|
) -> None:
|
|
"""Download file in chunks with optional rate limiting and progress tracking."""
|
|
remote_size = sftp.stat(remote_path).st_size
|
|
chunk_size = 32768 # 32KB chunks
|
|
transferred = resume_offset
|
|
|
|
# Initialize components if provided
|
|
if rate_limiter:
|
|
rate_limiter.start()
|
|
if progress_tracker:
|
|
progress_tracker.start()
|
|
|
|
mode = "ab" if resume_offset > 0 else "wb"
|
|
with open(local_path, mode) as local_file:
|
|
with sftp.file(remote_path, "rb") as remote_file:
|
|
if resume_offset > 0:
|
|
remote_file.seek(resume_offset)
|
|
|
|
while transferred < remote_size:
|
|
# Apply rate limiting before transfer
|
|
if rate_limiter:
|
|
rate_limiter.wait_if_needed(chunk_size)
|
|
|
|
# Perform the actual transfer
|
|
chunk = remote_file.read(chunk_size)
|
|
if not chunk:
|
|
break
|
|
|
|
local_file.write(chunk)
|
|
transferred += len(chunk)
|
|
|
|
# Update progress after transfer
|
|
if progress_tracker:
|
|
progress_tracker.update(transferred, remote_size)
|
|
|
|
# Final progress update
|
|
if progress_tracker:
|
|
progress_tracker.finish(remote_size)
|
|
|
|
def _require_client(self) -> paramiko.SSHClient:
|
|
if self._client is None:
|
|
raise RuntimeError("SSH client is not connected. Call open() or use a with-context.")
|
|
return self._client
|
|
|
|
def get_sftp(self) -> paramiko.SFTPClient:
|
|
if self._threading_mod:
|
|
th_sftp = self._require_client().open_sftp()
|
|
return th_sftp
|
|
if self._sftp is None:
|
|
self._sftp = self._require_client().open_sftp()
|
|
return self._sftp
|
|
|
|
@staticmethod
|
|
def _load_private_key(
|
|
key_file: Optional[str],
|
|
key_data: Optional[str],
|
|
passphrase: Optional[str],
|
|
) -> paramiko.PKey:
|
|
"""Load a private key by normalizing to key_data and parsing it.
|
|
|
|
Priority is mutually exclusive by design: key_file > key_data.
|
|
Supported types: RSA, DSS, ECDSA, Ed25519.
|
|
"""
|
|
if not key_data and key_file:
|
|
path = os.path.expanduser(key_file)
|
|
try:
|
|
with open(path, "r", encoding="utf-8") as f:
|
|
key_data = f.read()
|
|
except Exception as e:
|
|
raise RuntimeError(f"Failed to read private key file: {e}")
|
|
|
|
if not key_data:
|
|
raise RuntimeError("No private key provided")
|
|
|
|
stream = io.StringIO(key_data)
|
|
last_error: Optional[Exception] = None
|
|
key_classes = [paramiko.RSAKey, paramiko.ECDSAKey, paramiko.Ed25519Key]
|
|
if hasattr(paramiko, "DSSKey"): # 兼容无DSSKey功能的paramiko版本
|
|
key_classes.append(paramiko.DSSKey)
|
|
for key_cls in key_classes:
|
|
try:
|
|
stream.seek(0)
|
|
return key_cls.from_private_key(stream, password=passphrase)
|
|
except paramiko.PasswordRequiredException:
|
|
raise RuntimeError("Private key is encrypted; provide passphrase.")
|
|
except Exception as e:
|
|
last_error = e
|
|
raise RuntimeError(f"Failed to load private key from data: {last_error}")
|
|
|
|
def run_streaming(
|
|
self,
|
|
command: str,
|
|
on_stdout: Optional[Callable[[bytes], None]] = None,
|
|
on_stderr: Optional[Callable[[bytes], None]] = None,
|
|
timeout: Optional[int] = None,
|
|
read_chunk_size: int = 32768,
|
|
poll_interval_sec: float = 0.05,
|
|
) -> int:
|
|
"""
|
|
Execute a remote command and stream output chunks to callbacks to minimize memory usage.
|
|
|
|
Returns the process exit status when the command completes.
|
|
"""
|
|
client = self._require_client()
|
|
transport = client.get_transport()
|
|
if transport is None:
|
|
raise RuntimeError("SSH transport is not available")
|
|
|
|
effective_timeout = timeout if timeout is not None else self.timeout
|
|
chan = transport.open_session(timeout=effective_timeout)
|
|
chan.exec_command(command)
|
|
|
|
start_time = time.time()
|
|
|
|
try:
|
|
while True:
|
|
if chan.recv_ready():
|
|
data = chan.recv(read_chunk_size)
|
|
if data and on_stdout is not None:
|
|
on_stdout(data)
|
|
if chan.recv_stderr_ready():
|
|
data = chan.recv_stderr(read_chunk_size)
|
|
if data and on_stderr is not None:
|
|
on_stderr(data)
|
|
|
|
if chan.exit_status_ready() and not chan.recv_ready() and not chan.recv_stderr_ready():
|
|
break
|
|
|
|
if effective_timeout is not None and (time.time() - start_time) > effective_timeout:
|
|
chan.close()
|
|
raise TimeoutError("Command execution timed out")
|
|
|
|
time.sleep(poll_interval_sec)
|
|
|
|
exit_code = chan.recv_exit_status()
|
|
return exit_code
|
|
finally:
|
|
try:
|
|
chan.close()
|
|
except Exception:
|
|
pass
|
|
|
|
def execute_script_streaming(
|
|
self,
|
|
script_content: str,
|
|
script_type: str = "shell",
|
|
remote_dir: str = "/tmp",
|
|
script_name: Optional[str] = None,
|
|
timeout: Optional[int] = None,
|
|
cleanup: bool = True,
|
|
env_vars: Optional[Dict[str, str]] = None,
|
|
on_stdout: Optional[Callable[[bytes], None]] = None,
|
|
on_stderr: Optional[Callable[[bytes], None]] = None,
|
|
) -> int:
|
|
"""
|
|
Execute a bash script with streaming output.
|
|
|
|
Args:
|
|
script_content: The bash script content to execute
|
|
remote_dir: Remote directory to place the script (default: /tmp)
|
|
script_name: Name for the script file (auto-generated if None)
|
|
timeout: Command execution timeout in seconds
|
|
cleanup: Whether to delete the script file after execution
|
|
env_vars: Environment variables to set before script execution
|
|
on_stdout: Callback to receive stdout chunks (bytes)
|
|
on_stderr: Callback to receive stderr chunks (bytes)
|
|
|
|
Returns:
|
|
int: Exit code of the script execution
|
|
|
|
Raises:
|
|
RuntimeError: If script upload or execution fails
|
|
"""
|
|
remote_script_path = self._prepare_script(script_content, remote_dir, script_name)
|
|
|
|
try:
|
|
command = self._build_command(remote_script_path, script_type, env_vars)
|
|
return self.run_streaming(
|
|
command,
|
|
on_stdout=on_stdout,
|
|
on_stderr=on_stderr,
|
|
timeout=timeout,
|
|
)
|
|
finally:
|
|
if cleanup:
|
|
self._cleanup_script(remote_script_path)
|
|
|
|
def execute_script_collect(
|
|
self,
|
|
script_content: str,
|
|
script_type: str = "shell",
|
|
remote_dir: str = "/tmp",
|
|
script_name: Optional[str] = None,
|
|
timeout: Optional[int] = None,
|
|
cleanup: bool = True,
|
|
env_vars: Optional[Dict[str, str]] = None,
|
|
) -> CommandResult:
|
|
"""
|
|
Execute a bash script and collect all output.
|
|
|
|
Args:
|
|
script_content: The bash script content to execute
|
|
remote_dir: Remote directory to place the script (default: /tmp)
|
|
script_name: Name for the script file (auto-generated if None)
|
|
timeout: Command execution timeout in seconds
|
|
cleanup: Whether to delete the script file after execution
|
|
env_vars: Environment variables to set before script execution
|
|
|
|
Returns:
|
|
CommandResult: The execution result with exit_code, stdout, stderr
|
|
|
|
Raises:
|
|
RuntimeError: If script upload or execution fails
|
|
"""
|
|
remote_script_path = self._prepare_script(script_content, remote_dir, script_name)
|
|
|
|
try:
|
|
command = self._build_command(remote_script_path, script_type, env_vars)
|
|
code, out, err = self.run(command, timeout=timeout)
|
|
return CommandResult(exit_code=code, stdout=out, stderr=err)
|
|
finally:
|
|
if cleanup:
|
|
self._cleanup_script(remote_script_path)
|
|
|
|
def execute_local_script_streaming(
|
|
self,
|
|
local_script_path: str,
|
|
script_type: str = "shell",
|
|
remote_dir: str = "/tmp",
|
|
script_name: Optional[str] = None,
|
|
timeout: Optional[int] = None,
|
|
cleanup: bool = True,
|
|
env_vars: Optional[Dict[str, str]] = None,
|
|
on_stdout: Optional[Callable[[bytes], None]] = None,
|
|
on_stderr: Optional[Callable[[bytes], None]] = None,
|
|
) -> int:
|
|
"""
|
|
Execute a local bash script with streaming output.
|
|
|
|
Args:
|
|
local_script_path: Path to the local script file
|
|
remote_dir: Remote directory to place the script (default: /tmp)
|
|
script_name: Name for the script file (uses basename if None)
|
|
timeout: Command execution timeout in seconds
|
|
cleanup: Whether to delete the script file after execution
|
|
env_vars: Environment variables to set before script execution
|
|
on_stdout: Callback to receive stdout chunks (bytes)
|
|
on_stderr: Callback to receive stderr chunks (bytes)
|
|
|
|
Returns:
|
|
int: Exit code of the script execution
|
|
|
|
Raises:
|
|
FileNotFoundError: If local script file not found
|
|
RuntimeError: If script upload or execution fails
|
|
"""
|
|
if not os.path.isfile(local_script_path):
|
|
raise FileNotFoundError(f"Local script not found: {local_script_path}")
|
|
|
|
if not script_name:
|
|
script_name = os.path.basename(local_script_path)
|
|
|
|
remote_script_path = f"{remote_dir.rstrip('/')}/{script_name}"
|
|
|
|
# Upload the local script file via SFTP with LF normalization
|
|
sftp = self.get_sftp()
|
|
with open(local_script_path, "r", encoding="utf-8", newline="") as f:
|
|
content = f.read()
|
|
content_lf = content.replace("\r\n", "\n").replace("\r", "\n")
|
|
with sftp.file(remote_script_path, "w") as remote_file:
|
|
remote_file.write(content_lf.encode("utf-8"))
|
|
|
|
try:
|
|
command = self._build_command(remote_script_path, script_type, env_vars)
|
|
return self.run_streaming(
|
|
command,
|
|
on_stdout=on_stdout,
|
|
on_stderr=on_stderr,
|
|
timeout=timeout,
|
|
)
|
|
finally:
|
|
if cleanup:
|
|
self._cleanup_script(remote_script_path)
|
|
|
|
def execute_local_script_collect(
|
|
self,
|
|
local_script_path: str,
|
|
script_type: str = "shell",
|
|
remote_dir: str = "/tmp",
|
|
script_name: Optional[str] = None,
|
|
timeout: Optional[int] = None,
|
|
cleanup: bool = True,
|
|
env_vars: Optional[Dict[str, str]] = None,
|
|
) -> CommandResult:
|
|
"""
|
|
Execute a local bash script and collect all output.
|
|
|
|
Args:
|
|
local_script_path: Path to the local script file
|
|
remote_dir: Remote directory to place the script (default: /tmp)
|
|
script_name: Name for the script file (uses basename if None)
|
|
timeout: Command execution timeout in seconds
|
|
cleanup: Whether to delete the script file after execution
|
|
env_vars: Environment variables to set before script execution
|
|
|
|
Returns:
|
|
CommandResult: The execution result with exit_code, stdout, stderr
|
|
|
|
Raises:
|
|
FileNotFoundError: If local script file not found
|
|
RuntimeError: If script upload or execution fails
|
|
"""
|
|
if not os.path.isfile(local_script_path):
|
|
raise FileNotFoundError(f"Local script not found: {local_script_path}")
|
|
|
|
if not script_name:
|
|
script_name = os.path.basename(local_script_path)
|
|
|
|
remote_script_path = f"{remote_dir.rstrip('/')}/{script_name}"
|
|
|
|
# Upload the local script file via SFTP with LF normalization
|
|
sftp = self.get_sftp()
|
|
with open(local_script_path, "r", encoding="utf-8", newline="") as f:
|
|
content = f.read()
|
|
content_lf = content.replace("\r\n", "\n").replace("\r", "\n")
|
|
with sftp.file(remote_script_path, "w") as remote_file:
|
|
remote_file.write(content_lf.encode("utf-8"))
|
|
|
|
try:
|
|
command = self._build_command(remote_script_path, script_type, env_vars)
|
|
code, out, err = self.run(command, timeout=timeout)
|
|
return CommandResult(exit_code=code, stdout=out, stderr=err)
|
|
finally:
|
|
if cleanup:
|
|
self._cleanup_script(remote_script_path)
|
|
|
|
def _prepare_script(self, script_content: str, remote_dir: str, script_name: Optional[str]) -> str:
|
|
"""Prepare script by uploading content (LF normalized)."""
|
|
if not script_name:
|
|
import uuid
|
|
script_name = f"script_{uuid.uuid4().hex[:8]}"
|
|
|
|
remote_script_path = f"{remote_dir.rstrip('/')}/{script_name}"
|
|
|
|
sftp = None
|
|
try:
|
|
# Upload script content directly via SFTP (normalize to LF)
|
|
sftp = self.get_sftp()
|
|
content_lf = script_content.replace("\r\n", "\n").replace("\r", "\n")
|
|
with sftp.file(remote_script_path, "w") as remote_file:
|
|
remote_file.write(content_lf.encode("utf-8"))
|
|
except:
|
|
pass
|
|
finally:
|
|
if sftp and self._threading_mod:
|
|
sftp.close()
|
|
|
|
return remote_script_path
|
|
|
|
@staticmethod
|
|
def _build_command(
|
|
remote_script_path: str,
|
|
script_type: str = "shell",
|
|
env_vars: Optional[Dict[str, str]] = None) -> str:
|
|
"""Build the command string with environment variables."""
|
|
env_string = ""
|
|
if env_vars:
|
|
env_pairs = [f"{k}='{v}'" for k, v in env_vars.items()]
|
|
env_string = " ".join(env_pairs) + " "
|
|
|
|
if script_type == "shell":
|
|
return f"{env_string}bash {remote_script_path}"
|
|
elif script_type == "python":
|
|
get_py_bin = "pyBin=$(which python3 2> /dev/null || which python 2> /dev/null || echo 'python')"
|
|
py_info = "echo ""; echo \"Current Python environment:${pyBin} $(${pyBin} -c 'import sys,platform;print(sys.version.split()[0],platform.platform())')\""
|
|
cmd = "%s;${pyBin} %s; ret=$?; [ $ret -eq 0 ] && exit $ret; %s;exit $ret;" % (
|
|
get_py_bin, remote_script_path, py_info
|
|
)
|
|
return cmd
|
|
else:
|
|
raise ValueError("Invalid script type")
|
|
|
|
def _cleanup_script(self, remote_script_path: str) -> None:
|
|
"""Clean up the remote script file via SFTP without invoking shell."""
|
|
sftp = None
|
|
try:
|
|
sftp = self.get_sftp()
|
|
# Ensure path exists before removal
|
|
try:
|
|
sftp.stat(remote_script_path)
|
|
except FileNotFoundError:
|
|
return
|
|
sftp.remove(remote_script_path)
|
|
except Exception:
|
|
# Swallow cleanup errors
|
|
pass
|
|
finally:
|
|
if sftp and self._threading_mod:
|
|
sftp.close()
|
|
|
|
def path_exists(self, path: str) -> Tuple[bool, str]:
|
|
"""
|
|
Check if a path exists on the remote server.
|
|
|
|
Args:
|
|
path: Path to check
|
|
|
|
Returns:
|
|
Tuple[bool, str]: A tuple containing a boolean indicating whether the path exists and an error message
|
|
"""
|
|
sftp = None
|
|
try:
|
|
sftp = self.get_sftp()
|
|
try:
|
|
sftp.stat(path)
|
|
return True, ""
|
|
except FileNotFoundError:
|
|
return False, ""
|
|
except Exception as e:
|
|
return False, str(e)
|
|
finally:
|
|
if sftp and self._threading_mod:
|
|
sftp.close()
|
|
|
|
|
|
def create_dir(self, path: str):
|
|
"""
|
|
Create a directory on the remote server.
|
|
|
|
Args:
|
|
path: Path to create
|
|
|
|
Returns:
|
|
Tuple[bool, str]: A tuple containing a boolean indicating whether the directory was created successfully and an error message
|
|
"""
|
|
sftp = None
|
|
try:
|
|
sftp = self.get_sftp()
|
|
self._ensure_remote_dir(sftp, path)
|
|
return True, ""
|
|
except Exception as e:
|
|
return False, str(e)
|
|
finally:
|
|
if sftp and self._threading_mod:
|
|
return sftp.close()
|
|
|
|
def path_info(self, path: str) -> Dict:
|
|
"""
|
|
Get information about a path on the remote server.
|
|
|
|
Args:
|
|
path: Path to get information about
|
|
|
|
Returns:
|
|
Dict: A dictionary containing information about the path, including path, isdir, size, mtime, mode, uid, gid, and exists
|
|
"""
|
|
sftp = None
|
|
not_found = {"path": path,"isdir": False,"size": 0,"mtime": 0,"mode": 0,"uid": 0,"gid": 0, "exists": False}
|
|
try:
|
|
sftp = self.get_sftp()
|
|
info = sftp.stat(path)
|
|
return {
|
|
"path": path,
|
|
"isdir": stat.S_ISDIR(info.st_mode),
|
|
"size": info.st_size,
|
|
"mtime": info.st_mtime,
|
|
"mode": info.st_mode,
|
|
"uid": info.st_uid,
|
|
"gid": info.st_gid,
|
|
"exists": True
|
|
}
|
|
except FileNotFoundError:
|
|
return not_found
|
|
except:
|
|
return not_found
|
|
finally:
|
|
if sftp and self._threading_mod:
|
|
sftp.close()
|
|
|