重构

oneforall/common/database.py

@@ -75,27 +75,30 @@ class Database(object):
         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)',
+                    f'insert into "{table_name}" ('
+                    f'id, url, subdomain, port, ips, status, reason, valid,'
+                    f'title, banner, module, source, elapsed, count)'
+                    f'values (:id, :url, :subdomain, :port, :ips, :status,'
+                    f':reason, :valid, :title, :banner, :module, :source,'
+                    f':elapsed, :count)',
                     results)
             except Exception as e:
                 logger.log('ERROR', e)
 
-    def copy_table(self, table_name):
+    def copy_table(self, table_name, bak_table_name):
         """
         复制表创建备份
 
         :param str table_name: 表名
+        :param str bak_table_name: 新表名
         """
         table_name = table_name.replace('.', '_')
-        new_table_name = table_name + '_bak'
-        logger.log('DEBUG', f'正在将{table_name}表复制到{new_table_name}新表')
+        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 "{new_table_name}"')
-            self.conn.query(
-                f'create table "{new_table_name}" as select * from "{table_name}"')
+            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)
 
@@ -112,9 +115,38 @@ class Database(object):
         except Exception as e:
             logger.log('ERROR', e)
 
+    def drop_table(self, table_name):
+        """
+        删除表
+
+        :param str table_name: 表名
+        """
+        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)
+
+    def rename_table(self, table_name, new_table_name):
+        """
+        复制表创建备份
+
+        :param str table_name: 表名
+        :param str new_table_name: 新表名
+        """
+        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)
+
     def deduplicate_subdomain(self, table_name):
         """
-        去重表中的子域并删除空值和无效值
+        去重表中的子域
 
         :param str table_name: 表名
         """
@@ -122,21 +154,23 @@ class Database(object):
         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)')
+                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)
 
     def remove_invalid(self, table_name):
         """
-            去除表中的空值或无效子域
+        去除表中的空值或无效子域
 
-            :param str 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 subdomain is null or valid == 0')
+                f'delete from "{table_name}" where '
+                f'subdomain is null or valid == 0')
         except Exception as e:
             logger.log('ERROR', e)
 

oneforall/common/module.py

@@ -227,7 +227,7 @@ class Module(object):
                       'ips': None,
                       'status': None,
                       'reason': None,
-                      'valid': 1,
+                      'valid': None,
                       'title': None,
                       'banner': None,
                       'module': self.module,
@@ -247,7 +247,7 @@ class Module(object):
                           'ips': ips,
                           'status': None,
                           'reason': None,
-                          'valid': 1,
+                          'valid': None,
                           'title': None,
                           'banner': None,
                           'module': self.module,

oneforall/common/request.py

@@ -32,7 +32,8 @@ def gen_new_datas(datas, ports):
     new_datas = []
     protocols = ['http://']
     for data in datas:
-        if data.get('valid'):  # 有效的子域才进行http请求探测
+        valid = data.get('valid')
+        if valid is None:  # 子域有效性未知的才进行http请求探测
             subdomain = data.get('subdomain')
             for port in ports:
                 for protocol in protocols:
@@ -83,6 +84,7 @@ def deal_results(datas, results):
             if resp.status == 400 or resp.status >= 500:
                 datas[index]['valid'] = 0
             else:
+                datas[index]['valid'] = 1
                 headers = resp.headers
                 banner = str({'Server': headers.get('Server'),
                               'Via': headers.get('Via'),

oneforall/common/resolve.py

@@ -78,7 +78,6 @@ def resolve_callback(future, index, datas):
                 datas[index]['ips'] = str(ips)
             else:
                 datas[index]['ips'] = 'No answers'
-                datas[index]['valid'] = 0
 
 
 async def bulk_query_a(datas):