#!/usr/bin/env python3

"""
检查域名证书收集子域名
"""
import socket
import ssl

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


class CheckCert(Module):
    """
    通过检查域名的 SSL/TLS 证书来收集子域名信息
    """
    def __init__(self, domain):
        """
        初始化 CheckCert 类

        :param domain: 目标域名
        """
        Module.__init__(self)  # 调用父类 Module 的初始化方法
        self.domain = self.register(domain)  # 注册域名并保存
        self.port = 443  # SSL/TLS 默认端口
        self.module = 'Check'  # 模块名称
        self.source = 'CertInfo'  # 数据源标识

    def check(self):
        """
        获取域名的 SSL/TLS 证书并匹配证书中的子域名
        """
        try:
            # 创建 SSL 上下文
            ctx = ssl.create_default_context()
            # 创建一个 SSL 套接字并连接到指定域名的 443 端口
            sock = ctx.wrap_socket(socket.socket(), server_hostname=self.domain)
            sock.connect((self.domain, self.port))  # 连接到域名和端口
            # 获取对等证书
            cert_dict = sock.getpeercert()
        except Exception as e:
            # 捕获异常并记录日志
            logger.log('DEBUG', e.args)
            return

        # 从证书信息中匹配子域名
        subdomains = utils.match_subdomain(self.domain, str(cert_dict))
        # 将匹配到的子域名添加到子域名集合中
        self.subdomains = self.subdomains.union(subdomains)

    def run(self):
        """
        类执行入口
        """
        self.begin()  # 开始执行
        self.check()  # 执行证书检查
        self.finish()  # 完成执行

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


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

    :param domain: 域名
    """
    check = CheckCert(domain)  # 创建 CheckCert 实例
    check.run()  # 执行检查


if __name__ == '__main__':
    do('example.com')  # 对 example.com 执行检查