From 234f7f3a2000e1cb1f8a489a56d38deb010ee8ec Mon Sep 17 00:00:00 2001
From: Stanley <stanley_shao@126.com>
Date: Tue, 2 Apr 2024 22:21:51 +0800
Subject: [PATCH] feat: cancel scan

---
 main.py    | 10 ++++++++++
 threads.py |  6 +++++-
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/main.py b/main.py
index aab0052..34ce383 100644
--- a/main.py
+++ b/main.py
@@ -300,6 +300,7 @@ def __test_ips(self, after_scan=True):
         if after_scan:
             self.progressBar.setValue(0)
             self.progressBar.setMaximum(len(ips))
+            self.ui.btnWait_Scan.setText('扫描')
         else:
             self.__init_progessBar(len(ips))
         thread = SpeedtestThread(self, ips, self.__add_result, self.__found_unavailable,
@@ -318,11 +319,16 @@ def __got_scan_result(self, ip):
 
     def __scan_finished(self):
         self.__set_buttons_enabled(True)
+        self.ui.btnWait_Scan.setText('扫描')
         self.__remove_progessBar()
         self.ui.statusbar.showMessage('扫描完成')
 
     @Slot()
     def on_btnWait_Scan_clicked(self):
+        if self.ui.btnWait_Scan.text() == '取消':
+            self.ui.btnWait_Scan.setEnabled(False)
+            self.sthread.cancel()
+            return
         dlg = dlgScan(self)
         if dlg.exec() == QDialog.Accepted:
             max_ips = dlg.ui.spinBox_MaxIP.value()
@@ -345,6 +351,10 @@ def on_btnWait_Scan_clicked(self):
             self.logLabel.setText(f'开始扫描，共 {total_addrs} 个 IP...')
             thread.start()
 
+            self.sthread = thread
+            self.ui.btnWait_Scan.setEnabled(True)
+            self.ui.btnWait_Scan.setText('取消')
+
     def dragEnterEvent(self, event):
         event.accept()
 
diff --git a/threads.py b/threads.py
index 6b2fdf1..df55b06 100644
--- a/threads.py
+++ b/threads.py
@@ -148,11 +148,15 @@ def __add_block(self):
             self.currentIndex += 1
         return True
 
+    def cancel(self):
+        self.running = False
+
     def run(self):
+        self.running = True
         self.pool = QThreadPool()
         self.pool.setMaxThreadCount(self.num_workers)
         self.counter = 0
-        while self.counter < self.max_ips and self.__add_block():
+        while self.running and self.counter < self.max_ips and self.__add_block():
             self.pool.waitForDone()
             self.progressUpdate.emit(self.num_added)
         self.pool.waitForDone()