From 51e262762abd33cd4716908b43c435c6c7b940b5 Mon Sep 17 00:00:00 2001
From: Jerry Yan <792602257@qq.com>
Date: Wed, 8 Jun 2022 11:21:42 +0800
Subject: [PATCH] =?UTF-8?q?=E6=B5=8B=E8=AF=95=E5=9C=A8linux=E4=B8=8A?=
 =?UTF-8?q?=E5=8F=AF=E4=BB=A5=E6=AD=A3=E5=B8=B8gracefully=20exit=20?=
 =?UTF-8?q?=E5=90=8E=E6=9C=9FAFK=E4=BA=86=EF=BC=8C=E5=B7=A5=E4=BD=9C?=
 =?UTF-8?q?=E5=A4=AA=E5=BF=99=E4=BA=86?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 output/IOutput.py | 3 +++
 output/manager.py | 3 ++-
 proxy/manager.py  | 5 +++++
 3 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/output/IOutput.py b/output/IOutput.py
index 84b2601..d3dd3a1 100644
--- a/output/IOutput.py
+++ b/output/IOutput.py
@@ -43,6 +43,9 @@ class IOutput():
     def other_output(self, message_type: str, message_raw: bytes):
         ...
 
+    def debug_output(self, message_type: str, message_raw: str):
+        ...
+
     def error_output(self, message_type: str, message_raw: bytes, exception: Exception):
         ...
 
diff --git a/output/manager.py b/output/manager.py
index 79519c7..40a5c4c 100644
--- a/output/manager.py
+++ b/output/manager.py
@@ -122,7 +122,8 @@ class OutputManager():
             self.decode_payload(message)
 
     def terminate(self):
-        self._should_exit.set()
+        if self._should_exit:
+            self._should_exit.set()
         MESSAGE_QUEUE.put(None)
 
         for writer in self._writer:
diff --git a/proxy/manager.py b/proxy/manager.py
index 41d2617..d8f4b3a 100644
--- a/proxy/manager.py
+++ b/proxy/manager.py
@@ -18,6 +18,7 @@ _manager: "Optional[ProxyManager]" = None
 class ProxyManager:
     def __init__(self):
         self._mitm_instance = None
+        self._loop: "Optional[asyncio.AbstractEventLoop]" = None
         opts = Options(
             listen_host=config()['mitm']['host'],
             listen_port=config()['mitm']['port'],
@@ -34,6 +35,9 @@ class ProxyManager:
         self.terminate()
 
     def terminate(self):
+        if self._loop:
+            if self._loop.is_running():
+                self._loop.stop()
         if self._mitm_instance:
             self._mitm_instance.shutdown()
 
@@ -42,6 +46,7 @@ class ProxyManager:
 
     def _start(self):
         loop = asyncio.new_event_loop()
+        self._loop = loop
         asyncio.set_event_loop(loop)
         self._mitm_instance.run()