mirror of
https://github.com/aaPanel/aaPanel.git
synced 2026-08-26 09:32:26 +02:00
261 lines
9.8 KiB
Python
261 lines
9.8 KiB
Python
import os
|
|
import re
|
|
import socket
|
|
import time
|
|
import uuid
|
|
|
|
import paramiko
|
|
|
|
|
|
class SSHManager:
|
|
"""基础ssh"""
|
|
|
|
def __init__(self, host: str, auth: dict, port: int = 22, user: str = 'root', timeout: int = 10,
|
|
keepalive: int = 30, retry: int = 2):
|
|
self.host = host
|
|
self.port = port
|
|
self.user = user
|
|
self.auth = auth
|
|
self.timeout = timeout
|
|
self.keepalive = keepalive
|
|
self.retry = retry
|
|
self.client = None
|
|
|
|
def __enter__(self):
|
|
self.connect()
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
self.close()
|
|
|
|
def connect(self, retry: int = None, retry_delay: float = 0.25):
|
|
if self.client and self.client.get_transport() and self.client.get_transport().is_active():
|
|
return
|
|
|
|
attempts = (self.retry if retry is None else retry) + 1
|
|
last_error = None
|
|
for i in range(attempts):
|
|
self.close()
|
|
client = paramiko.SSHClient()
|
|
client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
|
|
try:
|
|
if (
|
|
"key_file" in self.auth
|
|
and self.auth["key_file"]
|
|
and os.path.exists(self.auth["key_file"])
|
|
):
|
|
client.connect(
|
|
self.host,
|
|
port=self.port,
|
|
username=self.user,
|
|
key_filename=self.auth["key_file"],
|
|
timeout=self.timeout,
|
|
banner_timeout=self.timeout,
|
|
auth_timeout=self.timeout,
|
|
compress=True, # 启用压缩
|
|
)
|
|
else:
|
|
client.connect(
|
|
self.host,
|
|
port=self.port,
|
|
username=self.user,
|
|
password=self.auth.get("password"),
|
|
timeout=self.timeout,
|
|
banner_timeout=self.timeout,
|
|
auth_timeout=self.timeout,
|
|
compress=True, # 启用压缩
|
|
)
|
|
transport = client.get_transport()
|
|
transport.set_keepalive(self.keepalive)
|
|
sock = transport.sock
|
|
sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) # 禁用 Nagle 算法
|
|
sock.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 256 * 1024) # 增大发送缓冲区
|
|
sock.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 256 * 1024) # 增大接收缓冲区
|
|
self.client = client
|
|
return
|
|
except Exception as e:
|
|
last_error = e
|
|
try:
|
|
client.close()
|
|
except Exception:
|
|
pass
|
|
if i == attempts - 1:
|
|
break
|
|
time.sleep(retry_delay * (i + 1)) # 0.25, 0.5, 0.75...
|
|
raise Exception(
|
|
f"SSH Connection to {self.user}@{self.host}:{self.port} failed "
|
|
f"after {attempts} attempts: {last_error}"
|
|
)
|
|
|
|
def execute(self, command: str, timeout: int = 10, retry: int = 2, retry_delay: float = 0.3) -> tuple:
|
|
"""Returns (exit_status, stdout, stderr)"""
|
|
error = ""
|
|
retry = self.retry if self.retry else retry
|
|
for r in range(retry + 1):
|
|
self.connect()
|
|
try:
|
|
stdin, stdout, stderr = self.client.exec_command(command, timeout=timeout)
|
|
out = stdout.read().decode('utf-8', errors='ignore').strip()
|
|
err = stderr.read().decode('utf-8', errors='ignore').strip()
|
|
exit_status = stdout.channel.recv_exit_status()
|
|
return exit_status, out, err
|
|
except Exception as e:
|
|
error = str(e)
|
|
if r >= retry:
|
|
break
|
|
self.close()
|
|
time.sleep(retry_delay)
|
|
return -1, "", f"Command execution failed: {error}"
|
|
|
|
def execute_many(self, commands: list[str], timeout: int = 30, retry: int = 2, retry_delay: float = 0.3) -> list:
|
|
"""批量执行命令并发"""
|
|
results = []
|
|
if not commands:
|
|
return results
|
|
|
|
boundary = f"BND_{uuid.uuid4().hex}" # 唯一边界符
|
|
# 前批次唯一的临时目录
|
|
tmp_dir = f"/tmp/ssh_batch_{uuid.uuid4().hex}"
|
|
MAX_LEN = 100000
|
|
chunks = []
|
|
current_chunk = []
|
|
current_len = 0
|
|
for i, cmd in enumerate(commands):
|
|
# 转义
|
|
octal_cmd = "".join(f"\\{b:03o}" for b in cmd.encode('utf-8'))
|
|
# 执行
|
|
run_cmd = (
|
|
f"( ( eval \"$(printf '%b' '{octal_cmd}')\" ) "
|
|
f"> {tmp_dir}/o_{i} 2> {tmp_dir}/e_{i}; "
|
|
f"echo $? > {tmp_dir}/r_{i} ) & "
|
|
)
|
|
# 收集
|
|
collect_cmd = (
|
|
f"printf '\\n{boundary}_S_{i}\\n';"
|
|
f"cat {tmp_dir}/o_{i} 2>/dev/null;"
|
|
f"RET=$(cat {tmp_dir}/r_{i} 2>/dev/null || echo -1);"
|
|
f"printf '\\n{boundary}_E_{i}_%d\\n' \"$RET\";"
|
|
|
|
f"printf '\\n{boundary}_S_{i}\\n' >&2;"
|
|
f"cat {tmp_dir}/e_{i} 2>/dev/null >&2;"
|
|
f"printf '\\n{boundary}_E_{i}\\n' >&2;"
|
|
)
|
|
total_len = len(run_cmd) + len(collect_cmd)
|
|
# 分块
|
|
if current_len + total_len > MAX_LEN and current_chunk:
|
|
chunks.append(current_chunk)
|
|
current_chunk = []
|
|
current_len = 0
|
|
|
|
current_chunk.append((i, cmd, run_cmd, collect_cmd))
|
|
current_len += total_len
|
|
|
|
if current_chunk:
|
|
chunks.append(current_chunk)
|
|
|
|
for chunk in chunks:
|
|
# 1.建目录 2.并发 3.等待完成 4.顺序打印 5.清理垃圾
|
|
run_cmds_str = "".join([c[2] for c in chunk])
|
|
collect_cmds_str = "".join([c[3] for c in chunk])
|
|
full_command = f"mkdir -p {tmp_dir}; {run_cmds_str} wait; {collect_cmds_str} rm -rf {tmp_dir};"
|
|
exit_status, out, err = self.execute(
|
|
full_command,
|
|
timeout=timeout,
|
|
retry=retry,
|
|
retry_delay=retry_delay
|
|
)
|
|
global_fail = (exit_status == -1 and "Command execution failed:" in err)
|
|
for i, cmd_str, _, _ in chunk:
|
|
res_dict = {
|
|
"index": i,
|
|
"cmd": cmd_str,
|
|
"code": -1,
|
|
"stdout": "",
|
|
"stderr": ""
|
|
}
|
|
if global_fail:
|
|
res_dict["stderr"] = err
|
|
results.append(res_dict)
|
|
continue
|
|
|
|
out_pattern = re.compile(rf"{boundary}_S_{i}\s*(.*?)\s*{boundary}_E_{i}_(-?\d+)", re.DOTALL)
|
|
out_match = out_pattern.search(out)
|
|
err_pattern = re.compile(rf"{boundary}_S_{i}\s*(.*?)\s*{boundary}_E_{i}", re.DOTALL)
|
|
err_match = err_pattern.search(err)
|
|
if out_match:
|
|
res_dict["stdout"] = out_match.group(1).strip()
|
|
res_dict["code"] = int(out_match.group(2))
|
|
else:
|
|
res_dict["stderr"] += "\n[Batch Error] Failed to parse output boundary."
|
|
|
|
if err_match:
|
|
res_dict["stderr"] = (res_dict["stderr"] + "\n" + err_match.group(1).strip()).strip()
|
|
|
|
results.append(res_dict)
|
|
|
|
results.sort(key=lambda x: x["index"])
|
|
return results
|
|
|
|
def download(self, remote_path: str, local_path: str, progress_callback=None) -> bool:
|
|
self.connect()
|
|
sftp = self.client.open_sftp()
|
|
try:
|
|
try:
|
|
remote_stat = sftp.stat(remote_path)
|
|
remote_size = remote_stat.st_size
|
|
except IOError:
|
|
raise FileNotFoundError(f"Remote file {remote_path} not found")
|
|
|
|
local_size = os.path.getsize(local_path) if os.path.exists(local_path) else 0
|
|
if local_size >= remote_size:
|
|
if progress_callback:
|
|
progress_callback(remote_size, remote_size)
|
|
return True
|
|
|
|
with open(local_path, 'ab' if local_size > 0 else 'wb') as local_f:
|
|
with sftp.open(remote_path, 'rb') as remote_f:
|
|
if local_size > 0:
|
|
remote_f.seek(local_size)
|
|
|
|
# Paramiko预读机制
|
|
remote_f.prefetch(remote_size)
|
|
target_duration = 0.05
|
|
min_chunk = 16 * 1024
|
|
max_chunk = 2 * 1024 * 1024
|
|
chunk_size = 64 * 1024
|
|
|
|
transferred = 0
|
|
while True:
|
|
start_time = time.perf_counter()
|
|
|
|
data = remote_f.read(int(chunk_size))
|
|
if not data:
|
|
break
|
|
|
|
local_f.write(data)
|
|
transferred += len(data)
|
|
|
|
if progress_callback:
|
|
progress_callback(local_size + transferred, remote_size)
|
|
|
|
elapsed = int(time.perf_counter() - start_time)
|
|
|
|
if elapsed > 0:
|
|
speed = len(data) / elapsed
|
|
next_chunk_size = int(speed * target_duration)
|
|
chunk_size = max(min_chunk, min(next_chunk_size, max_chunk))
|
|
else:
|
|
chunk_size = min(chunk_size * 2, max_chunk)
|
|
|
|
return True
|
|
finally:
|
|
sftp.close()
|
|
|
|
def close(self):
|
|
if self.client:
|
|
try:
|
|
self.client.close()
|
|
except Exception:
|
|
pass
|
|
self.client = None
|