# coding: utf8
import urlAudit
import global_list
import os
import sys
from selenium import webdriver
import time

class XSS_detect():
    def __init__(self):
        self.payload = []
        filename = os.path.join(sys.path[0], "Data", "xss_payload.txt")
        f = open(filename)
        for i in f:
            self.payload.append(i.strip())

        self.url = 'https://www.baidu.com/s?ie=utf-8&f=8&rsv_bp=1&rsv_idx=1&tn=80035161_1_dg&wd=webdriver.switch_to_alert().accept()&fenlei=256&oq=driver.switch_to_alert().accept()&rsv_pq=f8652eeb000392bf&rsv_t=f105iz6ovkH27OEZtKtM67BC259VLVXVveYXt7tz7E1qSAx0IEcu7L0rqD%2FjKtJMcIVJYQ&rqlang=cn&rsv_dl=tb&rsv_enter=1&rsv_btype=t&inputT=3331&rsv_sug3=115&rsv_sug2=0&rsv_sug4=3331'
        self.verificationErrors = []
        self.accept_next_alert = True
        self.driver = webdriver.Firefox()
        self.username = 'admin'
        self.password = 'password'

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

    def detect(self):
        test = urlAudit.UrlAudit()
        testurl = test.param(self.url)
        if testurl:
            rooturl, query, query_length = test.split(self.url.rstrip('\n'))
            for i in range(query_length):
            # print 'i=',i
                self.add_payload(rooturl, query, i)
            self.spider_reflected(self.url)
            self.spider_stored(self.url)

        else:
            self.spider_reflected(self.url)
            self.spider_stored(self.url)

        self.output_report()

    def add_payload(self, rooturl, query, i):
        for xsscase in self.payload:
            query2 = query[:]
            query2[i] = query2[i] + xsscase
            # print 'quety2[%d]===' % (i,query2[i])
            querystr = '&'.join(query2)
            #print "querystr====",querystr
            urlrequest = '%s?%s' % (rooturl, querystr)
            #print 'payloadUrl:',urlrequest

            self.load(urlrequest)

            res = self.reflected_catch_alert(query[i], xsscase)

            if res:
                break

    def reflected_catch_alert(self, injection_point, xsscase):
        try:
            #self.driver.switch_to_alert().accept()
            global_list.REFlECTED_XSS_INJECTION_dict[injection_point] = xsscase
            global_list.NUM_REFlECTED_XSS += 1
            return True
        except:
            return False

    def stored_catch_alert(self, injection_point, xsscase):
        try:

            #self.driver.switch_to_alert().accept()
            global_list.STROED_XSS_INJECTION_dict[injection_point] = xsscase
            global_list.NUM_STROED_XSS += 1
            return True
        except:

            return False


    def spider_reflected(self, url):
        self.load(url)
        # print u"#分析反射型xss"
        try:
            form = self.driver.find_element_by_xpath("//form")
            inputs = form.find_elements_by_xpath("//input[@type='text']")
            inputs_len = len(inputs)
            # print "inputs_len:",inputs_len
            # 分别测试input输入框
            for i in range(0, inputs_len):
                self.load(url)
                form = self.driver.find_element_by_xpath("//form")
                inputs = form.find_elements_by_xpath("//input[@type='text']")
                input_name = inputs[i].get_attribute('name')
                injection_point = "//form//input[@name='"
                injection_point += input_name
                injection_point += "']"

                for xsscase in self.payload:
                    self.load(url)
                    form = self.driver.find_element_by_xpath("//form")
                    inputs = form.find_elements_by_xpath("//input[@type='text']")
                    inputs[i].clear()
                    inputs[i].send_keys(xsscase)
                    form.find_element_by_xpath("//input[@type='submit']").click()
                    res = self.reflected_catch_alert(injection_point, xsscase)
                    if res:
                        break
        except:
            pass

    def spider_stored(self, url):
        try:
            for xsscase in self.payload:
                self.load(url)
                form = self.driver.find_element_by_xpath("//form")
                textarea = form.find_element_by_xpath('//textarea')
                textarea.clear()
                textarea.send_keys(xsscase)
                injection_point = '//form/textarea/'
                form.find_element_by_xpath("//input[@type='submit']").click()
                res = self.stored_catch_alert(injection_point, xsscase)
                if res:
                    break
        except:
            pass

    def output_report(self):
        print u"\n检测结论："
        print "-----------------------------------------------------------------------------------"
        print u"测试URL：\n", self.url
        print "-----------------------------------------------------------------------------------"
        if global_list.NUM_REFlECTED_XSS == 0:
            print u"该网页没有检测到反射型xss漏洞。"
        else:
            print u"该网页检测到反射型xss漏洞:"
            for k, v in global_list.REFlECTED_XSS_INJECTION_dict.items():
                print u"注入点：", k
                print u"payload用例：", v
                print "-----------------------------------------------------------------------------------"

        if global_list.NUM_STROED_XSS == 0:
            print u"该网页没有检测到存储型xss漏洞。"
        else:
            print u"该网页检测到存储型xss漏洞:"

        for k, v in global_list.STROED_XSS_INJECTION_dict.items():
            print u"注入点：", k
            print u"payload用例：", v

