diff --git a/spug_api/apps/exec/executors.py b/spug_api/apps/exec/executors.py index de61db6..cd01e16 100644 --- a/spug_api/apps/exec/executors.py +++ b/spug_api/apps/exec/executors.py @@ -14,12 +14,19 @@ def exec_worker_handler(job): class Job: - def __init__(self, key, hostname, port, username, pkey, command, token=None): + def __init__(self, key, name, hostname, port, username, pkey, command, token=None): self.ssh = SSH(hostname, port, username, pkey) self.key = key self.command = command self.token = token self.rds_cli = None + self.env = dict( + SPUG_HOST_ID=str(self.key), + SPUG_HOST_NAME=name, + SPUG_HOST_HOSTNAME=hostname, + SPUG_SSH_PORT=str(port), + SPUG_SSH_USERNAME=username, + ) def _send(self, message, with_expire=False): if self.rds_cli is None: @@ -39,12 +46,12 @@ class Job: def run(self): if not self.token: with self.ssh: - return self.ssh.exec_command(self.command) + return self.ssh.exec_command(self.command, self.env) self.send('\x1b[36m### Executing ...\x1b[0m\r\n') code = -1 try: with self.ssh: - for code, out in self.ssh.exec_command_with_stream(self.command): + for code, out in self.ssh.exec_command_with_stream(self.command, self.env): self.send(out) except socket.timeout: code = 130 @@ -52,5 +59,6 @@ class Job: except Exception as e: code = 131 self.send(f'\r\n\x1b[31m### Exception {e}\x1b[0m') + raise e finally: self.send_status(code) diff --git a/spug_api/apps/exec/views.py b/spug_api/apps/exec/views.py index 6bd5020..bcd657e 100644 --- a/spug_api/apps/exec/views.py +++ b/spug_api/apps/exec/views.py @@ -57,6 +57,7 @@ def do_task(request): for host in Host.objects.filter(id__in=form.host_ids): data = dict( key=host.id, + name=host.name, token=token, hostname=host.hostname, port=host.port, diff --git a/spug_api/libs/ssh.py b/spug_api/libs/ssh.py index 99cc852..14e7b72 100644 --- a/spug_api/libs/ssh.py +++ b/spug_api/libs/ssh.py @@ -54,8 +54,10 @@ class SSH: if exit_code != 0: raise Exception(f'add public key error: {out}') - def exec_command_raw(self, command): + def exec_command_raw(self, command, environment=None): channel = self.client.get_transport().open_session() + if environment: + channel.update_environment(environment) channel.set_combine_stderr(True) channel.exec_command(command) code, output = channel.recv_exit_status(), channel.recv(-1) @@ -77,17 +79,19 @@ class SSH: out += line return exit_code, out - def _win_exec_command_with_stream(self, command): - chan = self.client.get_transport().open_session() - chan.set_combine_stderr(True) - chan.get_pty(width=102) - chan.exec_command(command) - stdout = chan.makefile("rb", -1) + def _win_exec_command_with_stream(self, command, environment=None): + channel = self.client.get_transport().open_session() + if environment: + channel.update_environment(environment) + channel.set_combine_stderr(True) + channel.get_pty(width=102) + channel.exec_command(command) + stdout = channel.makefile("rb", -1) out = stdout.readline() while out: - yield chan.exit_status, out.decode() + yield channel.exit_status, out.decode() out = stdout.readline() - yield chan.recv_exit_status(), out.decode() + yield channel.recv_exit_status(), out.decode() def exec_command_with_stream(self, command, environment=None): channel = self._get_channel()