实现新子域标记

oneforall/common/database.py

@@ -37,7 +37,7 @@ class Database(object):
 
     def create_table(self, table_name):
         """
-        初始化数据库
+        创建表结构
 
         :param str table_name: 要创建的表名
         """
@@ -90,6 +90,26 @@ class Database(object):
             except Exception as e:
                 logger.log('ERROR', e)
 
+    def exist_table(self, table_name):
+        """
+        判断是否存在某表
+
+        :param str table_name: 表名
+        """
+        table_name = table_name.replace('.', '_')
+        logger.log('DEBUG', f'正在查询是否存在{table_name}表')
+        try:
+            result = self.conn.query(f'select count() from sqlite_master '
+                                     f'where type = "table" and '
+                                     f'name = "{table_name}"')
+        except Exception as e:
+            logger.log('ERROR', e)
+        else:
+            if len(result) != 0:
+                return True
+            else:
+                return False
+
     def copy_table(self, table_name, bak_table_name):
         """
         复制表创建备份
@@ -135,7 +155,7 @@ class Database(object):
 
     def rename_table(self, table_name, new_table_name):
         """
-        复制表创建备份
+        重命名表名
 
         :param str table_name: 表名
         :param str new_table_name: 新表名

oneforall/common/utils.py

@@ -207,3 +207,29 @@ def check_response(method, resp):
         else:
             logger.log('ALERT', msg)
     return False
+
+
+def mark_subdomain(old_data, new_data):
+    """
+    标记新增子域并返回新的数据集
+
+    :param old_data: 之前数据集
+    :param new_data: 现在数据集
+    :return: 已标记的新的数据集
+    """
+    # 第一次收集子域的情况
+    if not old_data:
+        for index, item in enumerate(new_data):
+            item['new'] = 1
+            new_data[index] = item
+        return new_data
+    # 非第一次收集子域的情况
+    old_subdomains = {item.get('subdomain') for item in old_data}
+    for index, item in enumerate(new_data):
+        subdomain = item.get('subdomain')
+        if subdomain in old_subdomains:
+            item['new'] = 0
+        else:
+            item['new'] = 1
+        new_data[index] = item
+    return new_data

oneforall/oneforall.py

@@ -77,7 +77,7 @@ class OneForAll(object):
         self.port = port
         self.domains = set()
         self.domain = str()
-        self.datas = list()
+        self.data = list()
         self.brute = brute
         self.verify = verify
         self.takeover = takeover
@@ -90,7 +90,8 @@ class OneForAll(object):
             self.brute = config.enable_brute_module
         if self.verify is None:
             self.verify = config.enable_verify_subdomain
-        rename_table = self.domain + '_last'
+        old_table = self.domain + '_last'
+        new_table = self.domain + '_now'
         collect = Collect(self.domain, export=False)
         collect.run()
         if self.brute:
@@ -102,49 +103,61 @@ class OneForAll(object):
         db.copy_table(self.domain, self.domain+'_ori')
         db.remove_invalid(self.domain)
         db.deduplicate_subdomain(self.domain)
+
+        old_data = []
+        # 非第一次收集子域的情况时数据库预处理
+        if db.exist_table(new_table):
+            db.drop_table(old_table)  # 如果存在上次收集结果表就先删除
+            db.rename_table(new_table, old_table)  # 新表重命名为旧表
+            old_data = db.get_data(old_table).as_dict()
+
         # 不验证子域的情况
         if not self.verify:
             # 数据库导出
             self.valid = None
-            dbexport.export(self.domain, valid=self.valid, format=self.format,
-                            show=self.show)
-            db.drop_table(rename_table)
-            db.rename_table(self.domain, rename_table)
+            dbexport.export(self.domain, valid=self.valid,
+                            format=self.format, show=self.show)
+            db.drop_table(new_table)
+            db.rename_table(self.domain, new_table)
             return
         # 开始验证子域工作
-        self.datas = db.get_data(self.domain).as_dict()
+        self.data = db.get_data(self.domain).as_dict()
+
+        # 标记新发现子域
+        self.data = utils.mark_subdomain(old_data, self.data)
+
         loop = asyncio.get_event_loop()
         asyncio.set_event_loop(loop)
 
         # 解析域名地址
-        task = resolve.bulk_query_a(self.datas)
-        self.datas = loop.run_until_complete(task)
+        task = resolve.bulk_query_a(self.data)
+        self.data = loop.run_until_complete(task)
 
         # 保存解析结果
         resolve_table = self.domain + '_res'
         db.drop_table(resolve_table)
         db.create_table(resolve_table)
-        db.save_db(resolve_table, self.datas, 'resolve')
+        db.save_db(resolve_table, self.data, 'resolve')
 
         # 请求域名地址
-        task = request.bulk_get_request(self.datas, self.port)
-        self.datas = loop.run_until_complete(task)
+        task = request.bulk_get_request(self.data, self.port)
+        self.data = loop.run_until_complete(task)
         # 在关闭事件循环前加入一小段延迟让底层连接得到关闭的缓冲时间
         loop.run_until_complete(asyncio.sleep(0.25))
 
         db.clear_table(self.domain)
-        db.save_db(self.domain, self.datas)
+        db.save_db(self.domain, self.data)
 
         # 数据库导出
-        dbexport.export(self.domain, valid=self.valid, format=self.format,
-                        show=self.show)
-        db.drop_table(rename_table)
-        db.rename_table(self.domain, rename_table)
+        dbexport.export(self.domain, valid=self.valid,
+                        format=self.format, show=self.show)
+        db.drop_table(new_table)
+        db.rename_table(self.domain, new_table)
         db.close()
         # 子域接管检查
 
         if self.takeover:
-            subdomains = set(map(lambda x: x.get('subdomain'), self.datas))
+            subdomains = set(map(lambda x: x.get('subdomain'), self.data))
             takeover = Takeover(subdomains)
             takeover.run()