添加IP是否为公网IP判断

oneforall/common/database.py

@@ -62,9 +62,10 @@ class Database(object):
                    f'new int,'
                    f'url text,'
                    f'subdomain text,'
-                   f'port int,'
                    f'level int,'
                    f'content text,'
+                   f'public int,'
+                   f'port int,'
                    f'status int,'
                    f'reason text,'
                    f'title text,'
@@ -92,12 +93,12 @@ class Database(object):
                 self.conn.bulk_query(
                     f'insert into "{table_name}" ('
                     f'id, type, valid, new, url, subdomain, port, level, content,'
-                    f'status, reason, title, banner, header, response, module,'
+                    f'public, status, reason, title, banner, header, response,'
-                    f'source, elapse, count)'
+                    f'module, source, elapse, count)'
                     f'values (:id, :type, :valid, :new, :url, :subdomain,'
-                    f':port, :level, :content, :status, :reason, :title, :banner,'
+                    f':port, :level, :content, :public, :status, :reason,'
-                    f':header, :response, :module, :source,:elapse, :count)',
+                    f':title, :banner, :header, :response, :module, :source,'
-                    results)
+                    f':elapse, :count)', results)
             except Exception as e:
                 logger.log('ERROR', e.args)
 
@@ -217,8 +218,9 @@ class Database(object):
         :param any valid: 有效性
         """
         table_name = table_name.replace('.', '_')
-        query = f'select id, type, valid, new, url, subdomain, port, level, content,' \
+        query = f'select id, type, valid, new, url, subdomain, level, ' \
-                f' status, reason, title, banner from "{table_name}"'
+                f'content, public, port, status, reason, title, banner ' \
+                f'from "{table_name}"'
         if valid:
             where = f' where valid = 1'
             query += where

oneforall/common/module.py

@@ -260,9 +260,10 @@ class Module(object):
                       'new': None,
                       'url': None,
                       'subdomain': None,
-                      'port': None,
                       'level': None,
                       'content': None,
+                      'public': None,
+                      'port': None,
                       'status': None,
                       'reason': None,
                       'title': None,
@@ -288,9 +289,10 @@ class Module(object):
                           'new': None,
                           'url': url,
                           'subdomain': subdomain,
-                          'port': None,
                           'level': level,
                           'content': content,
+                          'public': None,
+                          'port': None,
                           'status': None,
                           'reason': None,
                           'title': None,

oneforall/common/resolve.py

@@ -9,8 +9,11 @@ from dns.resolver import Resolver
 
 import config
 from config import logger
+from common import utils
 from common.database import Database
 
+socket.setdefaulttimeout(20)
+
 
 def dns_resolver():
     """
@@ -23,17 +26,17 @@ def dns_resolver():
     return resolver
 
 
-async def aiodns_query_a(hostname):
+async def aio_resolve_a(hostname, loop=None):
     """
-    异步查询A记录
+    异步解析A记录
 
     :param str hostname: 主机名
+    :param loop: 事件循环
     :return: 查询结果
     """
-    try:
+    if loop is None:
         loop = asyncio.get_event_loop()
-        socket.setdefaulttimeout(20)
+    try:
-        # answer = await loop.getaddrinfo(hostname, 80)
         answer = await loop.run_in_executor(None,
                                             socket.gethostbyname_ex,
                                             hostname)
@@ -53,9 +56,12 @@ def convert_results(result_list):
     result_dict = {}
     for result in result_list:
         hostname, answer = result
-        value_dict = {'content': None, 'reason': None, 'valid': None}
+        value_dict = {'content': None, 'reason': None,
+                      'public': None, 'valid': None}
         if isinstance(answer, tuple):
-            value_dict['content'] = ','.join(answer[2])
+            ip_list = answer[2]
+            value_dict['content'] = ','.join(ip_list)
+            value_dict['public'] = utils.check_ip_public(ip_list)
             result_dict[hostname] = value_dict
         elif isinstance(answer, Exception):
             value_dict['reason'] = str(answer.args)
@@ -113,9 +119,9 @@ def save_data(name, data):
     db.close()
 
 
-def resolve_progress(pr_queue, total):
+def resolve_progress_func(pr_queue, total):
     """
-    解析进度
+    解析进度函数
 
     :param pr_queue: 进度队列
     :param int total: 待解析的子域个数
@@ -134,17 +140,18 @@ def resolve_progress(pr_queue, total):
     bar.close()
 
 
-async def aio_query(pr_queue, hostname):
+async def do_resolve(pr_queue, hostname):
     """
-    异步查询主机名的A记录
+    异步解析主机名的A记录
 
     :param pr_queue: 进度队列
     :param str hostname: 主机名
     :return: 查询结果
     """
-    results = await aiodns_query_a(hostname)
+    loop = asyncio.get_event_loop()
+    result = await aio_resolve_a(hostname, loop)
     pr_queue.put(1)
-    return results
+    return result
 
 
 async def aio_resolve(subdomain_list, process_num, coroutine_num):
@@ -159,31 +166,28 @@ async def aio_resolve(subdomain_list, process_num, coroutine_num):
     m = Manager()
     pr_queue = m.Queue()
     loop = asyncio.get_event_loop()
-    loop.run_in_executor(None, resolve_progress, pr_queue, len(subdomain_list))
+    loop.run_in_executor(None, resolve_progress_func,
-    wrapped_query = functools.partial(aio_query, pr_queue)
+                         pr_queue, len(subdomain_list))
+    wrapped_resolve_func = functools.partial(do_resolve, pr_queue)
     async with aiomp.Pool(processes=process_num,
                           childconcurrency=coroutine_num) as pool:
-        results = await pool.map(wrapped_query, subdomain_list)
+        result_list = await pool.map(wrapped_resolve_func, subdomain_list)
-        return results
+        return result_list
 
 
-async def bulk_resolve(data_list):
+async def run_aio_resolve(subdomain_list):
     """
-    批量解析A记录并返回解析结果
+    异步解析子域A记录
 
-    :param list data_list: 待解析的数据列表
+    :param list subdomain_list: 待解析的子域列表
     :return: 解析得到的结果列表
     """
-    logger.log('INFOR', '正在异步查询子域的A记录')
-    # semaphore = asyncio.Semaphore(config.limit_resolve_conn)
-    query_subdomains = filter_subdomain(data_list)
     process_num = config.brute_process_num
     coroutine_num = config.brute_coroutine_num
-    results = await aio_resolve(query_subdomains, process_num, coroutine_num)
+    logger.log('INFOR', '正在异步查询子域的A记录')
-    results_dict = convert_results(results)
+    result_list = await aio_resolve(subdomain_list, process_num, coroutine_num)
-    data_list = update_data(data_list, results_dict)
     logger.log('INFOR', '完成异步查询子域的A记录')
-    return data_list
+    return result_list
 
 
 def run_resolve(data):
@@ -196,7 +200,9 @@ def run_resolve(data):
     """
     loop = asyncio.get_event_loop()
     asyncio.set_event_loop(loop)
-    resolve_coroutine = bulk_resolve(data)
+    need_resolve_subdomains = filter_subdomain(data)
-    # 在关闭事件循环前加入一小段延迟让底层连接得到关闭的缓冲时间
+    resolve_coroutine = run_aio_resolve(need_resolve_subdomains)
-    loop.run_until_complete(asyncio.sleep(0.25))
+    results_list = loop.run_until_complete(resolve_coroutine)
-    return loop.run_until_complete(resolve_coroutine)
+    results_dict = convert_results(results_list)
+    resolved_data = update_data(data, results_dict)
+    return resolved_data

oneforall/common/utils.py

@@ -2,8 +2,8 @@ import re
 import sys
 import time
 import random
-import ipaddress
 import platform
+from ipaddress import IPv4Address, ip_address
 
 import config
 from pathlib import Path
@@ -53,7 +53,7 @@ def gen_random_ip():
     生成随机的点分十进制的IP字符串
     """
     while True:
-        ip = ipaddress.IPv4Address(random.randint(0, 2 ** 32 - 1))
+        ip = IPv4Address(random.randint(0, 2 ** 32 - 1))
         if ip.is_global:
             return ip.exploded
 
@@ -389,3 +389,11 @@ def get_sample_banner(headers):
         temp_list.append(power)
     banner = ','.join(temp_list)
     return banner
+
+
+def check_ip_public(ip_list):
+    for ip_str in ip_list:
+        ip = ip_address(ip_str)
+        if not ip.is_global:
+            return 0
+    return 1