fix issues

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)

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,

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()