From 9bfcc9020efc93ccdddf82c43f3ff9140d6ab9c4 Mon Sep 17 00:00:00 2001
From: ChronosX88 <chronosx88@gmail.com>
Date: Mon, 21 Jan 2019 19:22:17 +0400
Subject: [PATCH] Corrected output of current running threads and made some
 small changes.

---
 CoreModel.py     |  8 +++-----
 MainPresenter.py | 30 ++++++++++++++++--------------
 main.py          |  4 ++--
 3 files changed, 21 insertions(+), 21 deletions(-)

diff --git a/CoreModel.py b/CoreModel.py
index 29bfa65..feac6a1 100644
--- a/CoreModel.py
+++ b/CoreModel.py
@@ -2,13 +2,11 @@ import socket
 
 
 class CoreModel:
-    def __init__(self):
+    def __init__(self, timeout):
         self.defSocket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        self.defSocket.settimeout(int(timeout))
 
-    def scanIP(self, host, ports, timeout):
-        if timeout == '':
-            timeout = '3'
-        self.defSocket.settimeout(3)
+    def scanIP(self, host, ports):
         openPorts = []
         for i in ports[0]:
             result = self.defSocket.connect_ex((host, int(i)))
diff --git a/MainPresenter.py b/MainPresenter.py
index 3a94e44..3003e97 100644
--- a/MainPresenter.py
+++ b/MainPresenter.py
@@ -2,13 +2,13 @@ import CoreModel
 import Parser
 import threading
 import queue
+import datetime
 from PyQt5.Qt import *
 from netaddr import IPNetwork
 
 
 class MainPresenter:
     def __init__(self, ui):
-        self.coreModel = CoreModel.CoreModel()
         self.ui = ui
         self.threads = []
         self.isScanEnabled = False
@@ -17,7 +17,6 @@ class MainPresenter:
     def startScan(self, ipRanges, portsStr, threadNumber, timeout):
         if timeout == '':
             timeout = '3'
-        self.isScanEnabled = True
         cidrIPRanges = Parser.getCIDRFromRanges(ipRanges)
         ports = Parser.getPortsFromString(portsStr)
         ips = []
@@ -27,29 +26,29 @@ class MainPresenter:
         for ip in ips:
             self.queue.put(ip)
         for i in range(int(threadNumber)):
-            self.threads.append(ScanThread(self.queue, ports, timeout))
+            self.threads.append(ScanThread(self.queue, ports, timeout, self))
             self.setCurrentThreadsLabel(len(self.threads))
         for thread in self.threads:
             thread.signal.connect(self.setLogText)
-            thread.finished.connect(self.onThreadExit)
+            thread.exit_signal.connect(self.on_thread_exit)
             thread.start()
 
-    def onThreadExit(self):
+    def on_thread_exit(self, thread):
+        self.threads.pop(self.threads.index(thread[0]))
+        self.setCurrentThreadsLabel(len(self.threads))
         if len(self.threads) == 0:
-            self.setCurrentThreadsLabel(0)
+            self.isScanEnabled = False
             self.ui.startButton.setText("Start")
-        else:
-            self.setCurrentThreadsLabel(len(self.threads) - 1)
 
     def stopScan(self):
         self.isScanEnabled = False
         for thread in self.threads:
-            thread.exit(0)
-        self.threads = []
+            thread.exit_signal.emit(self.threads.index(thread))
+            thread.exit()
         self.queue = queue.Queue()
 
     def setLogText(self, string):
-        self.ui.dataText.append(str(string) + '\n')
+        self.ui.dataText.append("[" + str(datetime.datetime.now()) + "] " + str(string))
 
     def setCurrentThreadsLabel(self, threadNumber):
         self.ui.currentThreadsLabel.setText(str(threadNumber))
@@ -58,21 +57,24 @@ class MainPresenter:
 class ScanThread(QThread):
 
     signal = pyqtSignal(str)
+    exit_signal = pyqtSignal(list)  # This signal for correct pop'ing from thread list
 
-    def __init__(self, scanQueue, ports, timeout, parent=None):
+    def __init__(self, scanQueue, ports, timeout, presenter, parent=None):
         QThread.__init__(self, parent)
         self.scanQueue = scanQueue
-        self.coreModel = CoreModel.CoreModel()
+        self.coreModel = CoreModel.CoreModel(timeout)
         self.ports = ports
         self._stop_event = threading.Event()
         self.timeout = timeout
+        self.presenter = presenter
 
     def run(self):
         while True:
             if self.scanQueue.empty():
+                self.exit_signal.emit([self])
                 self.exit(1)
             hostObject = self.scanQueue.get()
-            open_ports = self.coreModel.scanIP(str(hostObject), self.ports, self.timeout)
+            open_ports = self.coreModel.scanIP(str(hostObject), self.ports)
             signalStr = ''.join(str(x) for x in open_ports)
             if(signalStr == ''):
                 self.signal.emit(str(hostObject) + ' has no open ports!')
diff --git a/main.py b/main.py
index 6c8c244..0015859 100644
--- a/main.py
+++ b/main.py
@@ -20,12 +20,12 @@ class MyWin(QtWidgets.QMainWindow):
 
     def startButtonClicked(self):
         if self.isScanActive == False:
-            self.isScanActive = True
+            self.presenter.isScanEnabled = True
             self.ui.startButton.setText("Stop")
             self.presenter.startScan(self.ui.ipLine.text(), self.ui.portsLine.text(), self.ui.threadsLine.text(),
                                      self.ui.timeoutLine.text())
         else:
-            self.isScanActive = False
+            self.presenter.isScanEnabled = False
             self.ui.startButton.setText("Start")
             self.presenter.stopScan()