U 优化ssh类库

pull/410/head
vapao 2021-11-08 14:26:04 +08:00
parent a6621d30c6
commit c90eed1c8a
2 changed files with 27 additions and 20 deletions

View File

@ -18,6 +18,7 @@ class SSH:
self.channel = None self.channel = None
self.sftp = None self.sftp = None
self.eof = 'Spug EOF 2108111926' self.eof = 'Spug EOF 2108111926'
self.already_init = False
self.default_env = self._make_env_command(default_env) self.default_env = self._make_env_command(default_env)
self.regex = re.compile(r'Spug EOF 2108111926 (-?\d+)[\r\n]?') self.regex = re.compile(r'Spug EOF 2108111926 (-?\d+)[\r\n]?')
self.arguments = { self.arguments = {
@ -63,11 +64,7 @@ class SSH:
channel.set_combine_stderr(True) channel.set_combine_stderr(True)
channel.exec_command(command) channel.exec_command(command)
code, output = channel.recv_exit_status(), channel.recv(-1) code, output = channel.recv_exit_status(), channel.recv(-1)
try: return code, self._decode(output)
output = output.decode()
except UnicodeDecodeError:
output = output.decode(encoding='GBK')
return code, output
def exec_command(self, command, environment=None): def exec_command(self, command, environment=None):
channel = self._get_channel() channel = self._get_channel()
@ -94,9 +91,9 @@ class SSH:
stdout = channel.makefile("rb", -1) stdout = channel.makefile("rb", -1)
out = stdout.readline() out = stdout.readline()
while out: while out:
yield channel.exit_status, out.decode() yield channel.exit_status, self._decode(out)
out = stdout.readline() out = stdout.readline()
yield channel.recv_exit_status(), out.decode() yield channel.recv_exit_status(), self._decode(out)
def exec_command_with_stream(self, command, environment=None): def exec_command_with_stream(self, command, environment=None):
channel = self._get_channel() channel = self._get_channel()
@ -104,7 +101,7 @@ class SSH:
channel.send(command) channel.send(command)
exit_code, line = -1, '' exit_code, line = -1, ''
while True: while True:
line = channel.recv(8196).decode() line = self._decode(channel.recv(8196))
if not line: if not line:
break break
match = self.regex.search(line) match = self.regex.search(line)
@ -127,6 +124,10 @@ class SSH:
sftp = self._get_sftp() sftp = self._get_sftp()
return sftp.listdir_attr(path) return sftp.listdir_attr(path)
def sftp_stat(self, path):
sftp = self._get_sftp()
return sftp.stat(path)
def remove_file(self, path): def remove_file(self, path):
sftp = self._get_sftp() sftp = self._get_sftp()
sftp.remove(path) sftp.remove(path)
@ -137,14 +138,14 @@ class SSH:
counter = 0 counter = 0
self.channel = self.client.invoke_shell() self.channel = self.client.invoke_shell()
command = 'export PS1= && stty -echo; unsetopt zle; set -e\n' command = 'set +o zle\nset -o no_nomatch\nexport PS1= && stty -echo\n'
if self.default_env: if self.default_env:
command += f'{self.default_env}\n' command += f'{self.default_env}\n'
command += f'echo {self.eof} $?\n' command += f'echo {self.eof} $?\n'
self.channel.send(command.encode()) self.channel.send(command.encode())
while True: while True:
if self.channel.recv_ready(): if self.channel.recv_ready():
line = self.channel.recv(8196).decode() line = self._decode(self.channel.recv(8196))
if self.regex.search(line): if self.regex.search(line):
self.stdout = self.channel.makefile('r') self.stdout = self.channel.makefile('r')
break break
@ -163,11 +164,6 @@ class SSH:
self.sftp = self.client.open_sftp() self.sftp = self.client.open_sftp()
return self.sftp return self.sftp
def _break(self):
time.sleep(5)
command = f'\x03 echo {self.eof} -1\n'
self.channel.send(command.encode())
def _make_env_command(self, environment): def _make_env_command(self, environment):
if not environment: if not environment:
return None return None
@ -181,18 +177,29 @@ class SSH:
return f'export {str_envs}' return f'export {str_envs}'
def _handle_command(self, command, environment): def _handle_command(self, command, environment):
new_command = f'trap \'echo {self.eof} $?; rm -f $SPUG_EXEC_FILE\' EXIT\n' new_command = commands = ''
if not self.already_init:
commands = 'export SPUG_EXEC_FILE=$(mktemp)\n'
commands += 'trap \'rm -f $SPUG_EXEC_FILE\' EXIT\n'
self.already_init = True
env_command = self._make_env_command(environment) env_command = self._make_env_command(environment)
if env_command: if env_command:
new_command += f'{env_command}\n' new_command += f'{env_command}\n'
new_command += command new_command += command
new_command += f'\necho {self.eof} $?\n'
b64_command = base64.standard_b64encode(new_command.encode()) b64_command = base64.standard_b64encode(new_command.encode())
commands = 'export SPUG_EXEC_FILE=$(mktemp)\n'
commands += f'echo {b64_command.decode()} | base64 -di > $SPUG_EXEC_FILE\n' commands += f'echo {b64_command.decode()} | base64 -di > $SPUG_EXEC_FILE\n'
commands += 'bash $SPUG_EXEC_FILE\n' commands += 'source $SPUG_EXEC_FILE\n'
return commands return commands
def _decode(self, content):
try:
content = content.decode()
except UnicodeDecodeError:
content = content.decode(encoding='GBK', errors='ignore')
return content
def __enter__(self): def __enter__(self):
self.get_client() self.get_client()
transport = self.client.get_transport() transport = self.client.get_transport()

View File

@ -2,7 +2,7 @@ apscheduler==3.7.0
Django==2.2.* Django==2.2.*
channels==2.3.1 channels==2.3.1
channels_redis==2.4.1 channels_redis==2.4.1
paramiko==2.7.2 paramiko==2.8.0
django-redis==4.10.0 django-redis==4.10.0
requests==2.22.0 requests==2.22.0
GitPython==3.0.8 GitPython==3.0.8