降低代码重复率

oneforall/common/database.py

@@ -35,6 +35,14 @@ class Database(object):
         logger.log('DEBUG', f'使用数据库: {db_path}')
         return db.get_connection()
 
+    def query(self, sql):
+        try:
+            results = self.conn.query(sql)
+        except Exception as e:
+            logger.log('ERROR', e.args)
+        else:
+            return results
+
     def create_table(self, table_name):
         """
         创建表结构
@@ -43,27 +51,24 @@ class Database(object):
         """
         table_name = table_name.replace('.', '_')
         logger.log('DEBUG', f'正在创建{table_name}表')
-        try:
-            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'new int,'
-                            f'title text,'
-                            f'banner text,'
-                            f'header text,'
-                            f'response text,'
-                            f'module text,'
-                            f'source text,'
-                            f'elapsed float,'
-                            f'count int)')
-        except Exception as e:
-            logger.log('ERROR', e.args)
+        self.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'new int,'
+                   f'title text,'
+                   f'banner text,'
+                   f'header text,'
+                   f'response text,'
+                   f'module text,'
+                   f'source text,'
+                   f'elapsed float,'
+                   f'count int)')
 
     def save_db(self, table_name, results, module_name=None):
         """
@@ -74,7 +79,7 @@ class Database(object):
         :param str module_name: 模块名
         """
         logger.log('DEBUG', f'正在将{module_name}模块发现{table_name}的子域'
-                   '结果存入数据库')
+                            '结果存入数据库')
         table_name = table_name.replace('.', '_')
         if results:
             try:
@@ -98,17 +103,13 @@ class Database(object):
         """
         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.args)
+        results = self.query(f'select count() from sqlite_master '
+                             f'where type = "table" and '
+                             f'name = "{table_name}"')
+        if len(results) != 0:
+            return True
         else:
-            if len(result) != 0:
-                return True
-            else:
-                return False
+            return False
 
     def copy_table(self, table_name, bak_table_name):
         """
@@ -120,12 +121,9 @@ class Database(object):
         table_name = table_name.replace('.', '_')
         bak_table_name = bak_table_name.replace('.', '_')
         logger.log('DEBUG', f'正在将{table_name}表复制到{bak_table_name}新表')
-        try:
-            self.conn.query(f'drop table if exists "{bak_table_name}"')
-            self.conn.query(f'create table "{bak_table_name}" '
-                            f'as select * from "{table_name}"')
-        except Exception as e:
-            logger.log('ERROR', e.args)
+        self.query(f'drop table if exists "{bak_table_name}"')
+        self.query(f'create table "{bak_table_name}" '
+                   f'as select * from "{table_name}"')
 
     def clear_table(self, table_name):
         """
@@ -135,10 +133,7 @@ class Database(object):
         """
         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.args)
+        self.query(f'delete from "{table_name}"')
 
     def drop_table(self, table_name):
         """
@@ -148,10 +143,7 @@ class Database(object):
         """
         table_name = table_name.replace('.', '_')
         logger.log('DEBUG', f'正在删除{table_name}表')
-        try:
-            self.conn.query(f'drop table if exists "{table_name}"')
-        except Exception as e:
-            logger.log('ERROR', e.args)
+        self.query(f'drop table if exists "{table_name}"')
 
     def rename_table(self, table_name, new_table_name):
         """
@@ -163,11 +155,8 @@ class Database(object):
         table_name = table_name.replace('.', '_')
         new_table_name = new_table_name.replace('.', '_')
         logger.log('DEBUG', f'正在将{table_name}表重命名为{table_name}表')
-        try:
-            self.conn.query(f'alter table "{table_name}" '
-                            f'rename to "{new_table_name}"')
-        except Exception as e:
-            logger.log('ERROR', e.args)
+        self.query(f'alter table "{table_name}" '
+                   f'rename to "{new_table_name}"')
 
     def deduplicate_subdomain(self, table_name):
         """
@@ -177,12 +166,9 @@ class Database(object):
         """
         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) '
-                f'from "{table_name}" group by subdomain)')
-        except Exception as e:
-            logger.log('ERROR', e.args)
+        self.query(f'delete from "{table_name}" where '
+                   f'id not in (select min(id) '
+                   f'from "{table_name}" group by subdomain)')
 
     def remove_invalid(self, table_name):
         """
@@ -192,12 +178,8 @@ class Database(object):
         """
         table_name = table_name.replace('.', '_')
         logger.log('DEBUG', f'正在去除{table_name}表中的无效子域')
-        try:
-            self.conn.query(
-                f'delete from "{table_name}" where '
-                f'subdomain is null or valid == 0')
-        except Exception as e:
-            logger.log('ERROR', e.args)
+        self.query(f'delete from "{table_name}" where '
+                   f'subdomain is null or valid == 0')
 
     def get_data(self, table_name):
         """
@@ -207,12 +189,7 @@ class Database(object):
         """
         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.args)
-        else:
-            return rows
+        return self.query(f'select * from "{table_name}"')
 
     def export_data(self, table_name, valid):
         """
@@ -228,12 +205,7 @@ class Database(object):
             where = f' where valid = {valid}'
             query += where
         logger.log('DEBUG', f'获取{table_name}表中的所有数据')
-        try:
-            rows = self.conn.query(query)
-        except Exception as e:
-            logger.log('ERROR', e.args)
-        else:
-            return rows
+        return self.query(query)
 
     def close(self):
         self.conn.close()