使用selenium + Chrome爬取某網站烏雲公開漏洞文章並保存為pdf文件

目的：使用selenium + Chrome爬取某網站指定類型的烏雲公開漏洞文章，即在win10終端輸入漏洞類型（如未授權），則爬取所有該類型的漏洞文章，並把每個分頁的數字作為文件夾名，來保存該分頁下面的所有的漏洞文章。

總結：本例只是能簡單的爬取某一類型漏洞的所有文章，但不能爬取多個類型漏洞的所有文章，有時可能會有一些小bug導致沒爬取完就崩潰，需要手工修改進而重新爬。其它問題解決看注釋。

1. 關於python代碼里面出現中文在windows里面的處理，還沒完全掌握。可參考Python for Windows 中文編碼問題匯總
2. 遇到的超時問題TimeoutException: Message: timeout 還沒解決
3. 本例代碼只是輸入漏洞的指定類型，才可以爬取其所有文章。但如想一塊爬取漏洞的類型如 未授權、sql等，導入paramunittest也沒搞定

#!/usr/bin/env python
# -*- coding:utf-8 -*-
from selenium import webdriver
from selenium.webdriver.common.keys import Keys
import unittest
import time
from lxml import etree
import urllib2
import pdfkit
import random
import os
import shutil
import urllib
# import chardet
import paramunittest

# @paramunittest.parametrized(
#     {'user':'編輯器漏洞','result':'true'},
    # {'user':'未授權','result':'true'},
# )

class wooyunSelenium(unittest.TestCase):
# class TestDemo(unittest.TestCase):
    '''
    利用unittest模塊的測試類來解決js分頁，模擬點擊下一頁。
    
    '''
    
    def setUp(self):
    # def setParameters(self, user, result):
        ''' 
        初始化方法（固定寫法）,為測試准備環境。

        '''
        # self.user = user
        # self.result = result

        # 創建谷歌瀏覽器對象。
        self.driver = webdriver.Chrome()

        # 創建無頭瀏覽器對象，實測速度比谷歌瀏覽器的要快點，但不是很多，沒有瀏覽器窗口界面。
        # self.driver = webdriver.PhantomJS()

        # input輸入的提示語如果是字符串，要加雙引號，如果是數字，直接輸入(python2)，如果是raw_input函數，則不管輸入的是什么，全部返回的類型為字符串類型。
        # vname = input('請輸入要查詢的漏洞類型：')  "編輯器漏洞"

        # raw_input輸入的字符不用加雙引號，始終返回字符串類型
        self.vname = raw_input('請輸入要查詢的漏洞類型：')

        # 輸出類型為str類型,檢測該字符串的編碼格式
        # print chardet.detect(vname)

        # 對字符串中文進行url編碼使用urllib.quote函數，對字典里面的中文進行url編碼使用urlencode函數
        # print type(urllib.quote('編輯器漏洞'))
        # print urllib.quote('編輯器漏洞')

        # 實測發現windows系統的代碼里面出現中文，默認編碼格式為gbk，要進行解碼處理，才不會出現中文亂碼。
        # print self.vname
        os.mkdir(self.vname.decode('gbk'))
        os.chdir(self.vname.decode('gbk'))

        # 該處發現先解碼后編碼為utf-8，才可以解決下面成功獲取響應內容。
        self.vname = urllib.quote(self.vname.decode('gbk').encode('utf-8'))

        self.page = 1

        # 對self.vname進行url編碼，由於是字符串，所以采用urllib.quote()函數來編碼字符串成為url編碼格式，如果為字典，則使用urlencode函數來處理。
        # urllib.uquote(self.vname)
        
       
        # 請求網頁,如果該行代碼放在testMooc方法，會導致抓取的數據有重復，而且漏抓。注意vname要轉換為urlcode編碼格式，否則會報錯 UnicodeDecodeError: 'utf8' codec can't decode byte 0xb1 in position 41: invalid start byte
        self.driver.get("http://wooyun.jozxing.cc/search?keywords=" + self.vname + "&&content_search_by=by_bugs&&search_by_html=False&&page=" + str(self.page))
        # self.driver.get('http://wooyun.jozxing.cc/search?keywords=%E7%BC%96%E8%BE%91%E5%99%A8%E6%BC%8F%E6%B4%9E&content_search_by=by_bugs')

    def testWooyun(self):
    # def testcase(self):
        '''
        具體的測試用例方法（方法名開頭固定以test開頭）

        '''

        # 加上u用來處理中文亂碼，但發現不能使用PhantomJS，老是爆錯，而谷歌瀏覽器則不會。
        # os.mkdir(u'代碼執行')
        # os.chdir(u'命令執行')

        # 外循環控制頁數，因為經過觀察，某類型漏洞文章的最多分頁是沒有達到1000的，所以用該數以保證爬取到該類型漏洞的所有分頁
        for i in range(self.page,1000):

            # 創建一個文件夾來存放該頁所有文章，文件夾名字為分頁數字
            os.mkdir(str(i))

            # 讓網頁完全加載，避免網頁沒加載完導致獲取的數據丟失
            time.sleep(3)
            # 獲取網頁源碼
            html = self.driver.page_source
            # 把源碼解析為html dom文檔
            content = etree.HTML(html)
            # 使用xpath去匹配所有的課程鏈接

            links = content.xpath('//td[2]/a/@href')

            n = 0
            # 遍歷該分頁里面的文章鏈接
            for each in links:
                each = 'http://wooyun.jozxing.cc/' + each
                print each
                try:
                    # self.driver2 = webdriver.Chrome()報錯，發現頁面空白。至於是漏洞文章頁面還是分頁頁面，還沒確定，
                    # 處代碼表明是漏洞文章頁面請求為空。大約請求第234個，會出現一個空白頁面，一直等待好像死機一樣。解決方法把谷歌瀏覽器驅動程序放入python安裝目錄下的script目錄
                    self.driver2 = webdriver.Chrome()
                    # self.driver2 = webdriver.PhantomJS()
                    # 加上下面2行代碼，即設置超時時間，實測發現有效去除 不能讀取vr下面的某個文件 報錯。Unable to read VR Path Registry from C:\Users\hp\AppData\Local\openvr\openvrpaths.vrpath
                    self.driver2.set_page_load_timeout(10)
                    self.driver2.set_script_timeout(10)
                    # self.driver2.implicitly_wait(10)
                    self.driver2.get(each)
                    html2 = self.driver2.page_source
                    content2 = etree.HTML(html2)
                    # 獲取文章章標題
                    # title = content2.xpath("//h3[@class='wybug_title']/text()")[0]
                    # 處理獲取不到頁面時的異常（頁面空白），再請求訪問一次（可設置多次甚至循環），獲取不到頁面的原因可能是網站的反爬機制造成的。
                    title = content2.xpath("//h3[@class='wybug_title']/text()")[0]
                except:
                    # 關閉當前打開漏洞文章的瀏覽器窗口，發現使用close有時沒作用，則使用quit直接退出當前漏洞文章瀏覽器窗口
                    self.driver2.quit()
                    # 重新獲取該漏洞頁面
                    self.driver2 = webdriver.Chrome()
                    # self.driver2.implicitly_wait(10)
                    self.driver2.set_page_load_timeout(10)
                    self.driver2.set_script_timeout(10)
                    # self.driver2 = webdriver.PhantomJS()

                    # 怎么老是報超時錯誤：TimeoutException: Message: timeout
                    self.driver2.get(each)
                    html2 = self.driver2.page_source
                    content2 = etree.HTML(html2)
                    title = content2.xpath("//h3[@class='wybug_title']/text()")[0]

                # 設置保存的文件名，由於windows環境對文件名命名有'/'、'\'、'?'、'|'、'<'、'>'、'"'、'*'有限制，所以要有如下過濾
                filename = title[6:].strip().replace('"','').replace('/','_').replace('\\','_').replace('<','').replace('>','').replace('(','').replace(')','').replace('[','').replace(']','').replace('\\','').replace('%','').replace(';','').replace('*','').replace('?','').replace(':','').replace('|','').replace(u'®','').replace('.','').replace('&','')
                # file = filename + ".pdf"
                n += 1
                # 初始文件名
                file1 = str(n) + '.pdf'
                # 保存文件名
                file2 = filename + '.pdf'
                try:
                    path_wk = r'C:\Program Files\wkhtmltopdf\bin\wkhtmltopdf.exe'
                    config = pdfkit.configuration(wkhtmltopdf = path_wk)
                    pdfkit.from_url(each, file1, configuration=config)
                except:
                    path_wk = r'C:\Program Files\wkhtmltopdf\bin\wkhtmltopdf.exe'
                    config = pdfkit.configuration(wkhtmltopdf = path_wk)
                    pdfkit.from_url(each, file1, configuration=config)
                
                
                self.driver2.quit()
                # m變量值用來區分同名文件
                m = 1
                # 由於文件名有重復，所以使用遞歸函數來處理，文件重名的，文件名末尾加遞增數字保存。如a.pdf,a2.pdf
                self.modify_filename(file1,file2,filename,m)
                # time.sleep(random.randint(1,3))
            # 把當前分頁里面的所有的文章文件移動到對應的分頁文件里面
            for d in os.listdir('.'):
                if d.split('.')[-1] == 'pdf':
                    shutil.move(d,str(i))
            

            # 退出循環條件，從網頁源碼里面沒找到某個類名，則返回為-1，進而作為當點擊下一頁時，如果沒法點擊時，則此時的返回值不為-1
            if (html.find('next disabled')) != -1:
                break
            
            # 模擬瀏覽器手動點擊下一頁。  phantomjs好像不支持翻頁，但有時候測試發現也支持。搞不懂。
            self.driver.find_element_by_xpath('//li[@class="next"]/a').click()
            # self.assertTrue(self.result == "true")

    def modify_filename(self,file1,file2,filename,m):
       '''
       更改文件名函數
       如有多個同名文件，自動在文件名末尾加上數字，從2開始。
       方法遞歸

       '''

       if os.path.exists(file2):
            m += 1
            file2 = filename + str(m) + '.pdf'
            self.modify_filename(file1,file2,filename,m)
       else:
            os.rename(file1,file2)
            return


    def tearDown(self):
        ''' 退出方法（固定寫法），清理測試環境，以備下次測試使用'''
        
        self.driver.quit()

if __name__ == '__main__':
    unittest.main()
    # unittest.main(verbosity=2)