优化srv模块

modules/srv.py

@@ -2,7 +2,6 @@
 通过枚举域名常见的SRV记录并做查询来发现子域
 """
 
-import json
 import queue
 import threading
 
@@ -15,38 +14,30 @@ class BruteSRV(Module):
     def __init__(self, domain):
         Module.__init__(self)
         self.domain = domain
-        self.module = 'dnsquery'
+        self.module = 'BruteSRV'
         self.source = "BruteSRV"
         self.qtype = 'SRV'  # 利用的DNS记录的SRV记录查询子域
-        self.thread_num = 10
-        self.names_que = queue.Queue()
-        self.answers_que = queue.Queue()
+        self.thread_num = 20
+        self.names_queue = queue.Queue()
+        self.answers_queue = queue.Queue()
 
-    def gen_names(self):
+    def fill_queue(self):
         path = data_storage_dir.joinpath('srv_prefixes.json')
-        with open(path, encoding='utf-8', errors='ignore') as file:
-            prefixes = json.load(file)
-        names = map(lambda prefix: prefix + self.domain, prefixes)
-
-        for name in names:
-            self.names_que.put(name)
-
-    def brute(self):
-        """
-        枚举域名的SRV记录
-        """
-        self.gen_names()
+        prefixes = utils.load_json(path)
+        for prefix in prefixes:
+            self.names_queue.put(prefix + self.domain)
 
+    def do_brute(self):
         for num in range(self.thread_num):
-            thread = BruteThread(self.names_que, self.answers_que)
+            thread = BruteThread(self.names_queue, self.answers_queue)
             thread.name = f'BruteThread-{num}'
             thread.daemon = True
             thread.start()
+        self.names_queue.join()
 
-        self.names_que.join()
-
-        while not self.answers_que.empty():
-            answer = self.answers_que.get()
+    def deal_answers(self):
+        while not self.answers_queue.empty():
+            answer = self.answers_queue.get()
             if answer is None:
                 continue
             for item in answer:
@@ -59,7 +50,9 @@ class BruteSRV(Module):
         类执行入口
         """
         self.begin()
-        self.brute()
+        self.fill_queue()
+        self.do_brute()
+        self.deal_answers()
         self.finish()
         self.save_json()
         self.gen_result()
@@ -67,29 +60,19 @@ class BruteSRV(Module):
 
 
 class BruteThread(threading.Thread):
-    def __init__(self, names_que, answers_que):
+    def __init__(self, names_queue, answers_queue):
         threading.Thread.__init__(self)
-        self.names_que = names_que
-        self.answers_que = answers_que
+        self.names_queue = names_queue
+        self.answers_queue = answers_queue
 
     def run(self):
         while True:
-            name = self.names_que.get()
+            name = self.names_queue.get()
             answer = utils.dns_query(name, 'SRV')
-            self.answers_que.put(answer)
-            self.names_que.task_done()
-
-
-def run(domain):
-    """
-    类统一调用入口
-
-    :param str domain: 域名
-    """
-    brute = BruteSRV(domain)
-    brute.run()
+            self.answers_queue.put(answer)
+            self.names_queue.task_done()
 
 
 if __name__ == '__main__':
-    run('zonetransfer.me')
-    # run('example.com')
+    brute = BruteSRV('zonetransfer.me')
+    brute.run()