From 9fc7571422ffbca2772c60c38e2112295bd65df1 Mon Sep 17 00:00:00 2001
From: vapao <live1989@foxmail.com>
Date: Sat, 8 Oct 2022 15:33:55 +0800
Subject: [PATCH] =?UTF-8?q?A=20=E6=96=B0=E5=A2=9E=E7=BB=88=E6=AD=A2?=
 =?UTF-8?q?=E4=BB=BB=E5=8A=A1=E6=89=A7=E8=A1=8C=E5=8A=9F=E8=83=BD?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 spug_api/apps/exec/executors.py |  9 ++++-----
 spug_api/apps/exec/urls.py      |  3 ++-
 spug_api/apps/exec/views.py     | 17 ++++++++++++++++-
 spug_api/libs/ssh.py            | 15 +++++++++++++--
 4 files changed, 35 insertions(+), 9 deletions(-)

diff --git a/spug_api/apps/exec/executors.py b/spug_api/apps/exec/executors.py
index 5950f49..a422757 100644
--- a/spug_api/apps/exec/executors.py
+++ b/spug_api/apps/exec/executors.py
@@ -16,13 +16,13 @@ def exec_worker_handler(job):
 
 
 class Job:
-    def __init__(self, key, name, hostname, port, username, pkey, command, interpreter, params=None, token=None,
-                 term=None):
+    def __init__(self, token, key, name, hostname, port, username, pkey, command, interpreter, params=None, term=None):
         self.ssh = SSH(hostname, port, username, pkey, term=term)
         self.key = key
         self.command = self._handle_command(command, interpreter)
         self.token = token
         self.rds = get_redis_connection()
+        self.rds_key = f'PID:{self.token}:{self.key}'
         self.env = dict(
             SPUG_HOST_ID=str(self.key),
             SPUG_HOST_NAME=name,
@@ -50,14 +50,12 @@ class Job:
         self._send({'key': self.key, 'status': code})
 
     def run(self):
-        if not self.token:
-            with self.ssh:
-                return self.ssh.exec_command(self.command, self.env)
         flag = time.time()
         self.send('\r\n\x1b[36m### Executing ...\x1b[0m\r\n')
         code = -1
         try:
             with self.ssh:
+                self.rds.set(self.rds_key, self.ssh.get_pid(), 7200)
                 for code, out in self.ssh.exec_command_with_stream(self.command, self.env):
                     self.send(out)
             human_time = human_seconds_time(time.time() - flag)
@@ -70,4 +68,5 @@ class Job:
             self.send(f'\r\n\x1b[31m### Exception {e}\x1b[0m')
             raise e
         finally:
+            self.rds.delete(self.rds_key)
             self.send_status(code)
diff --git a/spug_api/apps/exec/urls.py b/spug_api/apps/exec/urls.py
index ffce12f..9ad212b 100644
--- a/spug_api/apps/exec/urls.py
+++ b/spug_api/apps/exec/urls.py
@@ -3,11 +3,12 @@
 # Released under the AGPL-3.0 License.
 from django.conf.urls import url
 
-from apps.exec.views import *
+from apps.exec.views import TemplateView, TaskView, handle_terminate
 from apps.exec.transfer import TransferView
 
 urlpatterns = [
     url(r'template/$', TemplateView.as_view()),
     url(r'do/$', TaskView.as_view()),
     url(r'transfer/$', TransferView.as_view()),
+    url(r'terminate/$', handle_terminate),
 ]
diff --git a/spug_api/apps/exec/views.py b/spug_api/apps/exec/views.py
index a6db18c..a6545c7 100644
--- a/spug_api/apps/exec/views.py
+++ b/spug_api/apps/exec/views.py
@@ -120,4 +120,19 @@ class TaskView(View):
         return json_response(error=error)
 
 
-
+@auth('exec.task.do')
+def handle_terminate(request):
+    form, error = JsonParser(
+        Argument('token', help='参数错误'),
+        Argument('host_id', type=int, help='参数错误')
+    ).parse(request.body)
+    if error is None:
+        host = Host.objects.get(pk=form.host_id)
+        rds = get_redis_connection()
+        rds_key = f'PID:{form.token}:{host.id}'
+        pid = rds.get(rds_key)
+        if pid:
+            with host.get_ssh() as ssh:
+                ssh.exec_command_raw(f'kill -9 {pid.decode()}')
+                rds.delete(rds_key)
+    return json_response(error=error)
diff --git a/spug_api/libs/ssh.py b/spug_api/libs/ssh.py
index 7f4d600..16558bc 100644
--- a/spug_api/libs/ssh.py
+++ b/spug_api/libs/ssh.py
@@ -58,6 +58,7 @@ class SSH:
         self.sftp = None
         self.exec_file = None
         self.term = term or {}
+        self.pid = None
         self.eof = 'Spug EOF 2108111926'
         self.default_env = default_env
         self.regex = re.compile(r'Spug EOF 2108111926 (-?\d+)[\r\n]?')
@@ -174,6 +175,12 @@ class SSH:
         sftp = self._get_sftp()
         sftp.remove(path)
 
+    def get_pid(self):
+        if self.pid:
+            return self.pid
+        self._get_channel()
+        return self.pid
+
     def _get_channel(self):
         if self.channel:
             return self.channel
@@ -183,13 +190,17 @@ class SSH:
         command = '[ -n "$BASH_VERSION" ] && set +o history\n'
         command += '[ -n "$ZSH_VERSION" ] && set +o zle && set -o no_nomatch\n'
         command += 'export PS1= && stty -echo\n'
-        command = self._handle_command(command, self.default_env)
+        command += f'echo {self.eof} $$\n'
         self.channel.sendall(command)
         out = ''
         while True:
             if self.channel.recv_ready():
                 out += self._decode(self.channel.recv(8196))
-                if self.regex.search(out):
+                match = self.regex.search(out)
+                if match:
+                    self.pid = int(match.group(1))
+                    if self.pid <= 1:
+                        raise Exception('Failed to get process pid')
                     self.stdout = self.channel.makefile('r')
                     break
             elif counter >= 100: