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,7 +11,12 @@ 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)
+
+    @staticmethod
+    def get_connection(db_path):
         """
         获取数据库对象
 
@@ -30,17 +35,16 @@ def connect_db(db_path=None):
         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: 要创建的表名
         """
+        table_name = table_name.replace('.', '_')
         logger.log('DEBUG', f'正在创建{table_name}表')
         try:
-        db_conn.query(f'create table if not exists "{table_name}" ('
+            self.conn.query(f'create table if not exists "{table_name}" ('
                             f'id integer primary key,'
                             f'url text,'
                             f'subdomain text,'
@@ -58,112 +62,111 @@ def create_table(db_conn, table_name):
         except Exception as e:
             logger.log('ERROR', e)
 
-
-def save_db(db_conn, table_name, results, module_name=None):
+    def save_db(self, 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}的子域结果存入数据库')
+        table_name = table_name.replace('.', '_')
         if results:
             try:
-            db_conn.bulk_query(f'insert into "{table_name}" (id, url, subdomain, port, ips, status,'
+                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)
+                    f':title, :banner, :module, :source, :elapsed, :count)',
+                    results)
             except Exception as e:
                 logger.log('ERROR', e)
 
-
-def copy_table(db_conn, table_name):
+    def copy_table(self, table_name):
         """
         复制表创建备份
 
-    :param db_conn: 数据库连接
         :param str table_name: 表名
         """
+        table_name = table_name.replace('.', '_')
         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}"')
+            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(db_conn, table_name):
+    def clear_table(self, table_name):
         """
         清空表中数据
 
-    :param db_conn: 数据库连接
         :param str table_name: 表名
         """
+        table_name = table_name.replace('.', '_')
         logger.log('DEBUG', f'正在清空{table_name}表中的数据')
         try:
-        db_conn.query(f'delete from "{table_name}"')
+            self.conn.query(f'delete from "{table_name}"')
         except Exception as e:
             logger.log('ERROR', e)
 
-
-def deduplicate_subdomain(db_conn, table_name):
+    def deduplicate_subdomain(self, table_name):
         """
         去重表中的子域并删除空值和无效值
 
-    :param db_conn: 数据库连接
         :param str table_name: 表名
         """
+        table_name = table_name.replace('.', '_')
         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)')
+            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 remove_invalid(db_conn, table_name):
+    def remove_invalid(self, table_name):
         """
             去除表中的空值或无效子域
 
-        :param db_conn: 数据库连接
             :param str table_name: 表名
             """
+        table_name = table_name.replace('.', '_')
         logger.log('DEBUG', f'正在去除{table_name}表中的无效子域')
         try:
-        db_conn.query(f'delete from "{table_name}" where subdomain is null or valid == 0')
+            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(db_conn, table_name):
+    def get_data(self, table_name):
         """
         获取表中的所有数据
 
-    :param db_conn: 数据库连接
         :param str table_name: 表名
         """
+        table_name = table_name.replace('.', '_')
         logger.log('DEBUG', f'获取{table_name}表中的所有数据')
         try:
-        rows = db_conn.query(f'select * from "{table_name}"')
+            rows = self.conn.query(f'select * from "{table_name}"')
         except Exception as e:
             logger.log('ERROR', e)
         else:
             return rows
 
-
-def get_subdomain(db_conn, table_name, valid):
+    def get_subdomain(self, table_name, valid):
         """
         获取表中的子域数据
 
-    :param db_conn: 数据库连接
         :param str table_name: 表名
         :param int valid: 是否有效
         """
+        table_name = table_name.replace('.', '_')
         logger.log('DEBUG', f'获取{table_name}表中的所有数据')
         try:
-        rows = db_conn.query(f'select * from "{table_name}" where valid = {valid}')
+            rows = self.conn.query(
+                f'select * from "{table_name}" where valid = {valid}')
         except Exception as e:
             logger.log('ERROR', e)
         else:

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')