database.py修改为类

oneforall/aiobrute.py

@@ -20,8 +20,9 @@ import fire
 import tqdm
 
 import config
-from common import database, resolve, utils
+from common import resolve, utils
 from common.module import Module
+from common.database import Database
 from config import logger
 
 
@@ -248,9 +249,8 @@ class AIOBrute(Module):
         while self.domains:
             self.domain = self.domains.pop()
             start = time.time()
-            db_conn = database.connect_db()
-            table_name = self.domain.replace('.', '_')
-            database.create_table(db_conn, table_name)
+            db = Database()
+            db.create_table(self.domain)
             if not rx_queue:
                 rx_queue = queue.Queue()
             logger.log('INFOR', f'开始执行{self.source}模块爆破域名{self.domain}')
@@ -278,10 +278,10 @@ class AIOBrute(Module):
             while not rx_queue.empty():
                 source, results = rx_queue.get()
                 # 将结果存入数据库中
-                database.save_db(db_conn, table_name, results, source)
-            database.copy_table(db_conn, table_name)
-            database.deduplicate_subdomain(db_conn, table_name)
-            database.remove_invalid(db_conn, table_name)
+                db.save_db(self.domain, results, source)
+            db.copy_table(self.domain)
+            db.deduplicate_subdomain(self.domain)
+            db.remove_invalid(self.domain)
 
             end = time.time()
             self.elapsed = round(end - start, 1)

oneforall/collect.py

@@ -3,12 +3,11 @@
 被动收集类
 """
 import time
-import queue
 import threading
 import importlib
 import config
 import dbexport
-from common import database
+from common.database import Database
 from config import logger
 
 
@@ -54,7 +53,7 @@ class Collect(object):
             import_object = importlib.import_module('.'+name, package)
             self.collect_func.append(getattr(import_object, 'do'))
 
-    def run(self, rx_queue=None):
+    def run(self):
         """
         类运行入口
         """
@@ -77,19 +76,18 @@ class Collect(object):
         for thread in threads:
             thread.join()
 
-        db_conn = database.connect_db()
-        table_name = self.domain.replace('.', '_')
-        database.create_table(db_conn, table_name)
-        database.copy_table(db_conn, table_name)
-        database.deduplicate_subdomain(db_conn, table_name)
-        database.remove_invalid(db_conn, table_name)
-        db_conn.close()
+        db = Database()
+        db.create_table(self.domain)
+        db.copy_table(self.domain)
+        db.deduplicate_subdomain(self.domain)
+        db.remove_invalid(self.domain)
+        # conn.close()
         # 数据库导出
         if self.export:
             if not self.path:
                 name = f'{self.domain}.{self.format}'
                 self.path = config.result_save_path.joinpath(name)
-            dbexport.export(table_name, path=self.path, format=self.format)
+            dbexport.export(self.domain, path=self.path, format=self.format)
         end = time.time()
         self.elapsed = round(end - start, 1)
 

oneforall/common/database.py

@@ -11,160 +11,163 @@ from records import Connection
 from config import logger
 
 
-def connect_db(db_path=None):
-    """
-    获取数据库对象
+class Database(object):
+    def __init__(self, db_path=None):
+        self.conn = self.get_connection(db_path)
 
-    :param db_path: 数据库连接或路径
-    :return: SQLite数据库
-    """
-    logger.log('DEBUG', f'正在获取数据库连接')
-    if isinstance(db_path, Connection):
-        return db_path
-    protocol = 'sqlite:///'
-    if not db_path:  # 数据库路径为空连接默认数据库
-        db_path = f'{protocol}{config.result_save_path}/result.sqlite3'
-    else:
-        db_path = protocol + db_path
-    db = records.Database(db_path)  # 不存在数据库时会新建一个数据库
-    logger.log('DEBUG', f'使用数据库: {db_path}')
-    return db.get_connection()
+    @staticmethod
+    def get_connection(db_path):
+        """
+        获取数据库对象
 
+        :param db_path: 数据库连接或路径
+        :return: SQLite数据库
+        """
+        logger.log('DEBUG', f'正在获取数据库连接')
+        if isinstance(db_path, Connection):
+            return db_path
+        protocol = 'sqlite:///'
+        if not db_path:  # 数据库路径为空连接默认数据库
+            db_path = f'{protocol}{config.result_save_path}/result.sqlite3'
+        else:
+            db_path = protocol + db_path
+        db = records.Database(db_path)  # 不存在数据库时会新建一个数据库
+        logger.log('DEBUG', f'使用数据库: {db_path}')
+        return db.get_connection()
 
-def create_table(db_conn, table_name):
-    """
-    初始化数据库
+    def create_table(self, table_name):
+        """
+        初始化数据库
 
-    :param db_conn: 数据库连接
-    :param str table_name: 要创建的表名
-    """
-    logger.log('DEBUG', f'正在创建{table_name}表')
-    try:
-        db_conn.query(f'create table if not exists "{table_name}" ('
-                      f'id integer primary key,'
-                      f'url text,'
-                      f'subdomain text,'
-                      f'port int,'
-                      f'ips text,'
-                      f'status int,'
-                      f'reason text,'
-                      f'valid int,'
-                      f'title text,'
-                      f'banner text,'
-                      f'module text,'
-                      f'source text,'
-                      f'elapsed float,'
-                      f'count int)')
-    except Exception as e:
-        logger.log('ERROR', e)
-
-
-def save_db(db_conn, table_name, results, module_name=None):
-    """
-    将各模块结果存入数据库
-
-    :param db_conn: 数据库连接
-    :param str table_name: 表名
-    :param list results: 结果列表
-    :param str module_name: 模块名
-    """
-    logger.log('DEBUG', f'正在将{module_name}模块发现{table_name}的子域结果存入数据库')
-    if results:
+        :param str table_name: 要创建的表名
+        """
+        table_name = table_name.replace('.', '_')
+        logger.log('DEBUG', f'正在创建{table_name}表')
         try:
-            db_conn.bulk_query(f'insert into "{table_name}" (id, url, subdomain, port, ips, status,'
-                               f'reason, valid, title, banner, module, source, elapsed, count)'
-                               f'values (:id, :url, :subdomain, :port, :ips, :status, :reason, :valid,'
-                               f':title, :banner, :module, :source, :elapsed, :count)', results)
+            self.conn.query(f'create table if not exists "{table_name}" ('
+                            f'id integer primary key,'
+                            f'url text,'
+                            f'subdomain text,'
+                            f'port int,'
+                            f'ips text,'
+                            f'status int,'
+                            f'reason text,'
+                            f'valid int,'
+                            f'title text,'
+                            f'banner text,'
+                            f'module text,'
+                            f'source text,'
+                            f'elapsed float,'
+                            f'count int)')
         except Exception as e:
             logger.log('ERROR', e)
 
+    def save_db(self, table_name, results, module_name=None):
+        """
+        将各模块结果存入数据库
 
-def copy_table(db_conn, table_name):
-    """
-    复制表创建备份
+        :param str table_name: 表名
+        :param list results: 结果列表
+        :param str module_name: 模块名
+        """
+        logger.log('DEBUG', f'正在将{module_name}模块发现{table_name}的子域结果存入数据库')
+        table_name = table_name.replace('.', '_')
+        if results:
+            try:
+                self.conn.bulk_query(
+                    f'insert into "{table_name}" (id, url, subdomain, port, ips, status,'
+                    f'reason, valid, title, banner, module, source, elapsed, count)'
+                    f'values (:id, :url, :subdomain, :port, :ips, :status, :reason, :valid,'
+                    f':title, :banner, :module, :source, :elapsed, :count)',
+                    results)
+            except Exception as e:
+                logger.log('ERROR', e)
 
-    :param db_conn: 数据库连接
-    :param str table_name: 表名
-    """
-    new_table_name = table_name + '_bak'
-    logger.log('DEBUG', f'正在将{table_name}表复制到{new_table_name}新表')
-    try:
-        db_conn.query(f'drop table if exists "{new_table_name}"')
-        db_conn.query(f'create table "{new_table_name}" as select * from "{table_name}"')
-    except Exception as e:
-        logger.log('ERROR', e)
+    def copy_table(self, table_name):
+        """
+        复制表创建备份
 
-
-def clear_table(db_conn, table_name):
-    """
-    清空表中数据
-
-    :param db_conn: 数据库连接
-    :param str table_name: 表名
-    """
-    logger.log('DEBUG', f'正在清空{table_name}表中的数据')
-    try:
-        db_conn.query(f'delete from "{table_name}"')
-    except Exception as e:
-        logger.log('ERROR', e)
-
-
-def deduplicate_subdomain(db_conn, table_name):
-    """
-    去重表中的子域并删除空值和无效值
-
-    :param db_conn: 数据库连接
-    :param str table_name: 表名
-    """
-    logger.log('DEBUG', f'正在去重{table_name}表中的子域')
-    try:
-        db_conn.query(f'delete from "{table_name}" where id not in (select min(id) from "{table_name}" group by subdomain)')
-    except Exception as e:
-        logger.log('ERROR', e)
-
-
-def remove_invalid(db_conn, table_name):
-    """
-        去除表中的空值或无效子域
-
-        :param db_conn: 数据库连接
         :param str table_name: 表名
         """
-    logger.log('DEBUG', f'正在去除{table_name}表中的无效子域')
-    try:
-        db_conn.query(f'delete from "{table_name}" where subdomain is null or valid == 0')
-    except Exception as e:
-        logger.log('ERROR', e)
+        table_name = table_name.replace('.', '_')
+        new_table_name = table_name + '_bak'
+        logger.log('DEBUG', f'正在将{table_name}表复制到{new_table_name}新表')
+        try:
+            self.conn.query(f'drop table if exists "{new_table_name}"')
+            self.conn.query(
+                f'create table "{new_table_name}" as select * from "{table_name}"')
+        except Exception as e:
+            logger.log('ERROR', e)
 
+    def clear_table(self, table_name):
+        """
+        清空表中数据
 
-def get_data(db_conn, table_name):
-    """
-    获取表中的所有数据
+        :param str table_name: 表名
+        """
+        table_name = table_name.replace('.', '_')
+        logger.log('DEBUG', f'正在清空{table_name}表中的数据')
+        try:
+            self.conn.query(f'delete from "{table_name}"')
+        except Exception as e:
+            logger.log('ERROR', e)
 
-    :param db_conn: 数据库连接
-    :param str table_name: 表名
-    """
-    logger.log('DEBUG', f'获取{table_name}表中的所有数据')
-    try:
-        rows = db_conn.query(f'select * from "{table_name}"')
-    except Exception as e:
-        logger.log('ERROR', e)
-    else:
-        return rows
+    def deduplicate_subdomain(self, table_name):
+        """
+        去重表中的子域并删除空值和无效值
 
+        :param str table_name: 表名
+        """
+        table_name = table_name.replace('.', '_')
+        logger.log('DEBUG', f'正在去重{table_name}表中的子域')
+        try:
+            self.conn.query(
+                f'delete from "{table_name}" where id not in (select min(id) from "{table_name}" group by subdomain)')
+        except Exception as e:
+            logger.log('ERROR', e)
 
-def get_subdomain(db_conn, table_name, valid):
-    """
-    获取表中的子域数据
+    def remove_invalid(self, table_name):
+        """
+            去除表中的空值或无效子域
 
-    :param db_conn: 数据库连接
-    :param str table_name: 表名
-    :param int valid: 是否有效
-    """
-    logger.log('DEBUG', f'获取{table_name}表中的所有数据')
-    try:
-        rows = db_conn.query(f'select * from "{table_name}" where valid = {valid}')
-    except Exception as e:
-        logger.log('ERROR', e)
-    else:
-        return rows
+            :param str table_name: 表名
+            """
+        table_name = table_name.replace('.', '_')
+        logger.log('DEBUG', f'正在去除{table_name}表中的无效子域')
+        try:
+            self.conn.query(
+                f'delete from "{table_name}" where subdomain is null or valid == 0')
+        except Exception as e:
+            logger.log('ERROR', e)
+
+    def get_data(self, table_name):
+        """
+        获取表中的所有数据
+
+        :param str table_name: 表名
+        """
+        table_name = table_name.replace('.', '_')
+        logger.log('DEBUG', f'获取{table_name}表中的所有数据')
+        try:
+            rows = self.conn.query(f'select * from "{table_name}"')
+        except Exception as e:
+            logger.log('ERROR', e)
+        else:
+            return rows
+
+    def get_subdomain(self, table_name, valid):
+        """
+        获取表中的子域数据
+
+        :param str table_name: 表名
+        :param int valid: 是否有效
+        """
+        table_name = table_name.replace('.', '_')
+        logger.log('DEBUG', f'获取{table_name}表中的所有数据')
+        try:
+            rows = self.conn.query(
+                f'select * from "{table_name}" where valid = {valid}')
+        except Exception as e:
+            logger.log('ERROR', e)
+        else:
+            return rows

oneforall/common/module.py

@@ -13,7 +13,7 @@ import config
 from config import logger
 from . import utils
 from .domain import Domain
-from common import database
+from common.database import Database
 
 
 lock = threading.Lock()
@@ -259,10 +259,9 @@ class Module(object):
 
     def save_db(self):
         lock.acquire()
-        db_conn = database.connect_db()
-        table_name = self.domain.replace('.', '_')
-        database.create_table(db_conn, table_name)
+        db = Database()
+        db.create_table(self.domain)
         source, results = self.results
         # 将结果存入数据库中
-        database.save_db(db_conn, table_name, results, source)
+        db.save_db(self.domain, results, source)
         lock.release()

oneforall/dbexport.py

@@ -9,7 +9,7 @@ OneForAll数据库导出模块
 """
 
 import fire
-from common import database
+from common.database import Database
 from config import logger
 
 
@@ -33,13 +33,13 @@ def export(table, db=None, valid=None, path=None, format='xlsx', output=False):
     :param str path:    导出路径(默认None)
     :param bool output: 是否将导出数据输出到终端(默认False)
     """
-    db_conn = database.connect_db(db)
+    database = Database(db)
     if valid is None:
-        rows = database.get_data(db_conn, table)
+        rows = database.get_data(table)
     elif isinstance(valid, int):
-        rows = database.get_subdomain(db_conn, table, valid)
+        rows = database.get_subdomain(table, valid)
     else:
-        rows = database.get_data(db_conn, table)  # 意外情况导出全部子域
+        rows = database.get_data(table)  # 意外情况导出全部子域
     if output:
         print(rows.dataset)
     if not path:

oneforall/oneforall.py

@@ -16,7 +16,8 @@ from datetime import datetime
 from config import logger
 from collect import Collect
 from aiobrute import AIOBrute
-from common import utils, database, resolve, request
+from common import utils, resolve, request
+from  common.database import Database
 
 
 banner = """\033[01;33m
@@ -80,9 +81,8 @@ class OneForAll(object):
                     # 由于爆破会有大量dns解析请求 并发常常会导致其他任务中的网络请求超时
                     brute = AIOBrute(self.domain)
                     brute.run()
-                table_name = self.domain.replace('.', '_')
-                db_conn = database.connect_db()
-                self.datas = database.get_data(db_conn, table_name).as_dict()
+                db = Database()
+                self.datas = db.get_data(self.domain).as_dict()
                 loop = asyncio.get_event_loop()
                 asyncio.set_event_loop(loop)
                 self.datas = loop.run_until_complete(resolve.bulk_query_a(self.datas))
@@ -90,13 +90,12 @@ class OneForAll(object):
                 # 在关闭事件循环前加入一小段延迟让底层连接得到关闭的缓冲时间
                 loop.run_until_complete(asyncio.sleep(0.25))
                 loop.close()
-                database.clear_table(db_conn, table_name)
-                database.save_db(db_conn, table_name, self.datas)
+                db.clear_table(self.domain)
+                db.save_db(self.domain, self.datas)
                 # 数据库导出
                 if not self.path:
                     self.path = config.result_save_path.joinpath(f'{self.domain}.{self.format}')
-                dbexport.export(table_name, db_conn, self.valid, self.path, self.format, self.output)
-                db_conn.close()
+                dbexport.export(self.domain, db.conn, self.valid, self.path, self.format, self.output)
         else:
             logger.log('FATAL', f'获取域名失败')
         logger.log('INFOR', f'结束运行OneForAll')