Update Vul_db.py

Vul_db.py

@@ -1,58 +1,81 @@
-import redis
+import redis
-import pymysql
+import pymysql
-import json,re
+import json,re
-
+
-class Vuldboption:
+from func_timeout import func_set_timeout
-    def table_exists(self, con, table_name):
+import func_timeout
-        # 判断数据表是否已经创建
+
-        sql = 'show tables;'
+class Vuldboption:
-        con.execute(sql)
+    def table_exists(self, con, table_name):
-        tables = [con.fetchall()]
+        # 判断数据表是否已经创建
-        table_list = re.findall('(\'.*?\')', str(tables))
+        sql = 'show tables;'
-        table_list = [re.sub("'", '', each) for each in table_list]  # 遍历并获得数据库表
+        con.execute(sql)
-        if table_name in table_list:
+        tables = [con.fetchall()]
-            return 1  # 创建了返回1
+        table_list = re.findall('(\'.*?\')', str(tables))
-        else:
+        table_list = [re.sub("'", '', each) for each in table_list]  # 遍历并获得数据库表
-            return 0  # 不创建返回0
+        if table_name in table_list:
-
+            return 1  # 创建了返回1
-    def process_item(self,table_name,redis_item_name):
+        else:
-
+            return 0  # 不创建返回0
-        pool = redis.ConnectionPool(host='127.0.0.1', port=6379, db=0, password=None)
+
-        redis_cli = redis.Redis(connection_pool=pool)
+    @func_set_timeout(1)
-
+    def process_item(self, table_name, redis_item_name, sql1, sql2):
-        conn = pymysql.connect(user='root',  # 用户名
+        pool = redis.ConnectionPool(host='127.0.0.1', port=6379, db=0, password=None)
-                password='root',  # 密码
+        redis_cli = redis.Redis(connection_pool=pool)
-                db = 'vulcrawl',  # 数据库名
+
-                host='127.0.0.1',  # 地址
+        conn = pymysql.connect(user='root',  # 用户名
-                port=3306,
+                password='root',  # 密码
-                charset='utf8')
+                db = 'vulcrawl',  # 数据库名
-        #table_name = 'db_cnnvd'  # 数据库表
+                host='127.0.0.1',  # 地址
-        # 没有对应数据库表则创建
+                port=3306,
-        if (self.table_exists(conn.cursor(), table_name) != 1):
+                charset='utf8')
-            sql = "create table db_cnnvd(time VARCHAR (30),title VARCHAR (1000),Numbering VARCHAR (40),url VARCHAR (200))"
+        page_start = 0
-            conn.cursor().execute(sql)  # 不存在则创建数据库表
+        step = 2
-
+        page_end = page_start + step
-        offset = 0
+        page_num = 0
-        while True:
+        #table_name = 'db_cnnvd'  # 数据库表
-            # 将数据从redis里pop出来
+        # 没有对应数据库表则创建
-            source, data = redis_cli.blpop(redis_item_name)
+        if (self.table_exists(conn.cursor(), table_name) != 1):
-            # 将redis字符串转为格式化的字典
+            #sql = "create table db_cnnvd(time VARCHAR (30),title VARCHAR (1000),Numbering VARCHAR (40),url VARCHAR (200))"
-            item = json.loads(data)
+            conn.cursor().execute(sql1)  # 不存在则创建数据库表
-            # 创建数据库游标
+
-            cursor = conn.cursor()
+        offset = 0
-
+        while True:
-            sql = "insert into db_cnnvd(time,title,Numbering,url)VALUES ('%s','%s','%s','%s')" % (item['time'], item['title'], item['Numbering'], item['url'])
+            # 将数据从redis里pop出来
-            cursor.execute(sql)  # 执行数据插入
+            source, data = redis_cli.blpop(redis_item_name)
-            conn.commit()  # 提交记录
+            # 将redis字符串转为格式化的字典
-
+            if redis_cli.blpop(redis_item_name):
-            cursor.close()  # 关闭游标
+                page_start = page_end + 1
-
+                page_end = page_start + step
-            offset += 1
+                page_num += 1
-            print("正在保存第：" + str(offset) + "条记录")
+            # 如果获取到的数据为空,则跳出循环
-
+            else:
-
+                break
-if __name__ == '__main__':
+            item = json.loads(data)
-    vul_cnnvd = Vuldboption()
+            # 创建数据库游标
-    vul_cnnvd.process_item('db_cnnvd','vulcrawl:items')
+            cursor = conn.cursor()
-    vul_woodyun = Vuldboption()
+            # "insert into db_cnnvd(time,title,Numbering,url)VALUES ('%s','%s','%s','%s')"
-    vul_woodyun.process_item('db_woodyun','vulcrawl2:items')
+            sql = sql2 % (item['time'], item['title'], item['Numbering'], item['url'])
+            cursor.execute(sql)  # 执行数据插入
+            conn.commit()  # 提交记录
+
+            cursor.close()  # 关闭游标
+
+            offset += 1
+            print("正在保存第：" + str(offset) + "条记录")
+
+
+
+if __name__ == '__main__':
+    creat_cnnvd_table = "create table db_cnnvd(time VARCHAR (30),title VARCHAR (1000),Numbering VARCHAR (40),url VARCHAR (200))"
+    insert_cnnvd_data = "insert into db_cnnvd(time,title,Numbering,url)VALUES ('%s','%s','%s','%s')"
+    creat_woodyun_table = "create table db_woodyun(time VARCHAR (100),title VARCHAR (1000),Numbering VARCHAR (40),url VARCHAR (200))"
+    insert_woodyun_data = "insert into db_woodyun(time,title,Numbering,url)VALUES ('%s','%s','%s','%s')"
+    try:
+        vul_cnnvd = Vuldboption()
+        vul_cnnvd.process_item('db_cnnvd', 'vulcrawl:items', creat_cnnvd_table, insert_cnnvd_data)
+    except func_timeout.exceptions.FunctionTimedOut:
+        vul_woodyun = Vuldboption()
+        vul_woodyun.process_item('db_woodyun', 'vulcrawl2:items', creat_woodyun_table, insert_woodyun_data)
+    else:
+        pass