"""
检查内容安全策略收集子域名
"""
import requests

from client.subdomain.oneforall.common import utils
from client.subdomain.oneforall.common.module import Module
from client.subdomain.oneforall.config import logger


class CheckCSP(Module):
    """
    检查内容安全策略 (CSP) 收集子域名
    """
    def __init__(self, domain, header=None):
        """
        初始化 CheckCSP 类

        :param domain: 目标域名
        :param header: (可选) 响应头，若提供则使用此值
        """
        Module.__init__(self)  # 调用父类 Module 的初始化方法
        self.domain = self.register(domain)  # 注册域名
        self.module = 'Check'  # 模块名称
        self.source = 'ContentSecurityPolicy'  # 数据源标识
        self.csp_header = header  # 存储 CSP 响应头

    def grab_header(self):
        """
        抓取目标域名的响应头，查找 Content-Security-Policy 字段

        :return: 响应头字典
        """
        csp_header = dict()  # 初始化为空字典
        urls = [f'http://{self.domain}', f'https://{self.domain}',
                f'http://www.{self.domain}', f'https://www.{self.domain}']
        
        # 循环通过不同的 URL 进行请求，获取头部信息
        for url in urls:
            self.header = self.get_header()
            self.proxy = self.get_proxy(self.source)
            response = self.get(url, check=False)
            if response:
                csp_header = response.headers  # 获取响应头
                break
        return csp_header

    def check(self):
        """
        通过解析 Content-Security-Policy 响应头，匹配子域名
        """
        if not self.csp_header:
            self.csp_header = self.grab_header()  # 获取响应头
        csp = self.csp_header.get('Content-Security-Policy')  # 提取 CSP 字段
        if not self.csp_header:
            logger.log('DEBUG', f'获取 {self.domain} 域的请求头失败')
            return
        if not csp:
            logger.log('DEBUG', f'{self.domain} 域的响应头中没有内容安全策略字段')
            return
        # 匹配 CSP 中的子域名信息
        self.subdomains = utils.match_subdomain(self.domain, csp)

    def run(self):
        """
        执行检查过程并保存结果
        """
        self.begin()  # 开始执行
        self.check()  # 执行 CSP 检查
        self.finish()  # 完成执行

        # 保存检查结果
        self.save_json()
        self.gen_result()
        self.save_db()


def do(domain, header=None):  # 统一入口，方便多线程调用
    """
    类统一调用入口

    :param domain: 域名
    :param header: 响应头（可选）
    """
    check = CheckCSP(domain, header)  # 创建 CheckCSP 实例
    check.run()  # 执行检查


if __name__ == '__main__':
    # 示例：获取 'google-analytics.com' 域的响应头并执行检查
    resp = requests.get('https://content-security-policy.com/')
    do('google-analytics.com', resp.headers)