Merge branch 'dev'

scan/scanner.py

@@ -4,12 +4,13 @@
 import time
 
 from multiprocessing import Lock
+from collections import defaultdict
 
 CWD = os.path.dirname(__file__)
 sys.path.append(os.path.join(CWD, '..'))
 from scan.modules import modules
 from utils.net import get_all_ip, get_ip_seg_len
-from utils.base import multi_thread, process_bar, save_res
+from utils.base import multi_thread, process_bar, save_res, printf
 from utils.config import *
 
 
@@ -54,7 +55,6 @@ def __init__(self, args) -> None:
         self.done_lock = Lock()
         self.found_lock = Lock()
         self.file_lock = Lock()
-        self.bar_lock = Lock()
         self.start_time = time.time()
         self.bar = process_bar()
 
@@ -68,6 +68,7 @@ def _preprocess(self):
         # total ip
         with open(self.args.in_file, 'r') as f:
             total_ip = [l.strip() for l in f if not l.startswith('#') and l.strip()]
+            total_ip = list(set(total_ip))
         for ip in total_ip:
             self.total += get_ip_seg_len(ip) if '-' in ip or '/' in ip else 1
 
@@ -92,7 +93,38 @@ def _preprocess(self):
                         self.found += 1
 
     def __del__(self):
-        self.paused.close()
+        if os.path.exists(os.path.join(self.args.out_path, PAUSE)): self.paused.close()
+
+
+    def report(self):
+        """report the results"""
+        if not os.path.exists(os.path.join(self.args.out_path, RESULTS_ALL)):
+            return
+
+        with open(os.path.join(self.args.out_path, RESULTS_ALL), 'r') as f:
+            items = [l.strip().split(',') for l in f if l.strip()]
+
+        results = defaultdict(lambda: defaultdict(lambda: 0))
+        for i in items:
+            dev, vul = i[-2], i[-1]
+            results[dev][vul] += 1
+        results_sum = len(items)
+        results_max = max([val for vul in results.values() for val in vul.values()])
+
+        print('\n')
+        print('-' * 19, 'REPORT', '-' * 19)
+        for dev in results:
+            vuls = [(vul_name, vul_count) for vul_name, vul_count in results[dev].items()]
+            dev_sum = sum([i[1] for i in vuls])
+            printf(f"{dev} {dev_sum}", color='red', bold=True)
+            for vul_name, vul_count in vuls:
+                printf(f"{vul_name:>18} | ", end='')
+                block_num = int(vul_count / results_max * 25)
+                printf('▥' * block_num, end=' ')
+                printf(vul_count)
+        printf(f"{'sum: ' + str(results_sum):>46}", color='yellow', flash=True)
+        print('-' * 46)
+        print('\n')
 
 
     def scan_meta(self, ip, mod_name):
@@ -117,7 +149,6 @@ def scan_meta(self, ip, mod_name):
         except Exception as e:
             if DEBUG: print(e)
         finally:
-            with self.bar_lock: self.bar(self.total, self.done + 1, self.found, timer=True, start_time=self.start_time)
             return found
 
 
@@ -144,10 +175,7 @@ def scan(self, ip_term: str):
             found = False
             dev_type = modules['device_type'](ip)  # hikvision, dahua, cctv, dlink, unidentified
 
-            if dev_type == 'unidentified':
-                with self.done_lock: self.done += 1
-                continue
-            elif dev_type == 'hikvision' and self.mod_by_device('hikvision'):
+            if dev_type == 'hikvision' and self.mod_by_device('hikvision'):
                 if 'hik_weak' in self.modules: found |= self.scan_meta(ip, 'hik_weak')
                 if 'hb_weak' in self.modules: found |= self.scan_meta(ip, 'hb_weak')
                 if 'cve_2017_7921' in self.modules: found |= self.scan_meta(ip, 'cve_2017_7921')
@@ -159,9 +187,15 @@ def scan(self, ip_term: str):
                 if 'cctv_weak' in self.modules: found |= self.scan_meta(ip, 'cctv_weak')
             elif dev_type == 'dlink' and self.mod_by_device('dlink'):
                 if 'cve_2020_25078' in self.modules: found |= self.scan_meta(ip, 'cve_2020_25078')
-            if not found: save_res([ip, dev_type], os.path.join(self.args.out_path, RESULTS_FAILED))
 
-            with self.done_lock: self.done += 1
+            if not found and dev_type != 'unidentified':
+                save_res([ip, dev_type], os.path.join(self.args.out_path, RESULTS_FAILED))
+
+            with self.done_lock:
+                self.done += 1
+                # self.bar(self.total, self.done + 1, self.found, timer=True, start_time=self.start_time)
+                self.bar(self.total, self.done, self.found, timer=True, start_time=self.start_time)
+
         # write paused
         with self.file_lock:
             self.paused.write(ip_term + '\n')
@@ -183,4 +217,5 @@ def __call__(self):
                     self.modules[mod_name] = mod_func
 
         multi_thread(self.scan, self.ip_list, processes=self.args.th_num)
+        self.report()
         self._close()

utils/base.py

@@ -22,15 +22,15 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
-@run_time
+# @run_time
 def multi_process(func, items, processes=40):
     """multiprocess API"""
     with Pool(processes) as pool:
         res = pool.map_async(func, items).get()
     return res
 
 
-@run_time
+# @run_time
 def multi_thread(func, items, processes=40):
     """multiprocess API"""
     with ThreadPool(processes) as pool:
@@ -64,8 +64,8 @@ def output_formatter(info, color='green', bold=False, underline=False, flash=Fal
     return head + bold + underline + flash + color + str(info) + tail
 
 
-def printf(info, color='green', bold=False, underline=False, flash=False):
-    print(output_formatter(info, color=color, bold=bold, underline=underline, flash=flash))
+def printf(info, color='green', bold=False, underline=False, flash=False, *args, **kwargs):
+    print(output_formatter(info, color=color, bold=bold, underline=underline, flash=flash), *args, **kwargs)
 
 
 def time_formatter(t: float) -> str:
@@ -91,16 +91,16 @@ def wrapper(total, done, found=0, timer=False, start_time=0):
             time_pred = time_used * (total / done)
             time_used = output_formatter(time_formatter(time_used), color='cyan', bold=True)
             time_pred = output_formatter(time_formatter(time_pred), color='white', bold=True)
-            _time = f"time: {time_used}/{time_pred}"
+            _time = f"Time: {time_used}/{time_pred}"
 
         # count
         _total = output_formatter(total, color='blue', bold=True)
         _done = output_formatter(done, color='blue', bold=True)
         _percent = output_formatter(f"{round(done / total * 100, 1)}%", color='pink', bold=True)
         _found = 'Found ' + output_formatter(found, color='red', bold=True) if found else ''
-        count = f"{_done}/{_total} ({_percent}) {_found}"
+        count = f"{_done}/{_total}({_percent}) {_found}"
 
-        print(f"\r{icon} {count}  {_time:<55}", end='')
+        print(f"\r{icon} {count}  {_time:<55}", end='', flush=True)
     return wrapper