优化内存

brute.py

@@ -90,9 +90,9 @@ def gen_subdomains(expression, path):
 
     :param  str  expression: generate subdomains's expression
     :param  str  path: path of wordlist
-    :return list subdomains: list of subdomains
+    :return set  subdomains: list of subdomains
     """
-    subdomains = list()
+    subdomains = set()
     with open(path, encoding='utf-8', errors='ignore') as fd:
         for line in fd:
             word = line.strip().lower()
@@ -105,14 +105,13 @@ def gen_subdomains(expression, path):
             if word.endswith('.'):
                 word = word[:-1]
             subdomain = expression.replace('*', word)
-            subdomains.append(subdomain)
+            subdomains.add(subdomain)
     size = len(subdomains)
     logger.log('DEBUG', f'The size of the dictionary generated by {path} is {size}')
     if size == 0:
         logger.log('ALERT', 'Please check the dictionary content!')
     else:
-        random_domain = random.choice(subdomains)
-        logger.log('ALERT', f'Please check whether {random_domain} is correct or not')
+        utils.check_random_subdomain(subdomains)
     return subdomains
 
 
@@ -123,12 +122,12 @@ def gen_fuzz_subdomains(expression, rule, fuzzlist):
     :param  str  expression: generate subdomains's expression
     :param  str  rule: regexp rule
     :param  str  fuzzlist: fuzz dictionary
-    :return list subdomains: list of subdomains
+    :return set  subdomains: list of subdomains
     """
-    subdomains = list()
+    subdomains = set()
     if fuzzlist:
         fuzz_domain = gen_subdomains(expression, fuzzlist)
-        subdomains = subdomains.extend(fuzz_domain)
+        subdomains.update(fuzz_domain)
     if rule:
         fuzz_count = exrex.count(rule)
         if fuzz_count > 10000000:
@@ -139,9 +138,8 @@ def gen_fuzz_subdomains(expression, rule, fuzzlist):
             if not fuzz_string.isalnum():
                 continue
             fuzz_domain = expression.replace('*', fuzz_string)
-            subdomains.append(fuzz_domain)
-        random_domain = random.choice(subdomains)
-        logger.log('ALERT', f'Please check whether {random_domain} is correct or not')
+            subdomains.add(fuzz_domain)
+        utils.check_random_subdomain(subdomains)
     logger.log('DEBUG', f'Dictionary size based on fuzz mode: {len(subdomains)}')
     return subdomains
 
@@ -152,7 +150,7 @@ def gen_word_subdomains(expression, path):
 
     :param  str  expression: generate subdomains's expression
     :param  str  path: path of wordlist
-    :return list subdomains: list of subdomains
+    :return set  subdomains: list of subdomains
     """
     subdomains = gen_subdomains(expression, path)
     logger.log('DEBUG', f'Dictionary based on word mode size: {len(subdomains)}')
@@ -543,11 +541,10 @@ class Brute(Module):
             wordlist = self.recursive_nextlist
         if self.word:
             word_subdomains = gen_word_subdomains(self.place, wordlist)
-            # set可以合并list
-            dict_set = dict_set.union(word_subdomains)
+            dict_set.update(word_subdomains)
         if self.fuzz:
             fuzz_subdomains = gen_fuzz_subdomains(self.place, self.rule, self.fuzzlist)
-            dict_set = dict_set.union(fuzz_subdomains)
+            dict_set.update(fuzz_subdomains)
         count = len(dict_set)
         logger.log('INFOR', f'Dictionary size: {count}')
         if count > 10000000:

common/utils.py

@@ -728,3 +728,13 @@ def match_subdomains(domain, html, distinct=True, fuzzy=True):
         return set(deal)
     else:
         return list(deal)
+
+
+def check_random_subdomain(subdomains):
+    if not subdomains:
+        logger.log('ALERT', f'The generated dictionary is empty')
+        return False
+    for subdomain in subdomains:
+        if subdomain:
+            logger.log('ALERT', f'Please check whether {subdomain} is correct or not')
+            return True