#!/usr/bin/python3
# coding=utf-8

"""
OneForAll多进程多协程异步子域爆破模块

:copyright: Copyright (c) 2019, Jing Ling. All rights reserved.
:license: GNU General Public License v3.0, see LICENSE for more details.
"""

import asyncio
import os
import queue
import secrets
import signal
import time

import aiomultiprocess
import exrex
import fire
import tqdm

import config
from common import database, resolve, utils
from common.module import Module
from config import logger


def init_worker():
    signal.signal(signal.SIGINT, signal.SIG_IGN)


def get_wordlist(name):
    return config.data_storage_path.joinpath(name)


def detect_wildcard(domain):
    """
    探测域名是否使用泛解析

    :param str domain: 域名
    :return: 如果没有使用泛解析返回False 使用返回泛解析的IP集合和ttl整型值
    """
    logger.log('INFOR', f'正在探测{domain}是否使用泛解析')
    token = secrets.token_hex(16)
    random_subdomain = f'{token}.{domain}'
    try:
        answers = resolve.dns_query_a(random_subdomain)
    except Exception as e:  # 如果查询随机域名A记录出错 说明不存在随机子域的A记录 即没有开启泛解析
        logger.log('DEBUG', e)
        logger.log('INFOR', f'{domain}没有使用泛解析')
        return False, None, None
    ttl = answers.ttl
    ips = {item.address for item in answers}
    logger.log('ALERT', f'{domain}使用了泛解析')
    logger.log('ALERT', f'{random_subdomain} 解析到IP: {ips} TTL: {ttl}')
    return True, ips, ttl


def wildcard_by_compare(ips, ttl, wildcard_ips, wildcard_ttl):
    """
    通过与泛解析返回的IP集合和判TTL值进行对比判断发现的子域是否是泛解析子域

    :param set ips: 子域A记录查询出的IP集合
    :param int ttl: 子域A记录查询出的TTL整型值
    :param set wildcard_ips: 泛解析的IP集合
    :param int wildcard_ttl: 泛解析的TTL整型值
    :return: 判断结果
    :rtype bool
    """
    # 参考：http://sh3ll.me/archives/201704041222.txt
    if not ips.issubset(wildcard_ips):
        return False
    if ttl != wildcard_ttl and ttl % 60 == 0 and wildcard_ttl % 60 == 0:
        return False
    return True


def wildcard_by_times(ips, ips_times):
    """
    对ips出现次数进行判断泛解析

    :param set ips: 发现子域的IP集合
    :param ips_times: 子域IP集合出现次数统计字典
    :return: 判断结果
    """
    times = ips_times.get(str(ips))
    if times > config.ips_appear_maximum:
        return True
    return False


def gen_fuzz_domains(domain, rule):
    domains = set()
    if '{fuzz}' not in domain:
        logger.log('FATAL', f'没有指定fuzz位置')
        return domains
    if not rule:
        logger.log('FATAL', f'没有指定fuzz规则')
        return domains
    fuzz_count = exrex.count(rule)
    if fuzz_count > 2000000:
        logger.log('FATAL', f'fuzz规则范围太大：{fuzz_count} > 2000000')
        return domains
    logger.log('INFOR', f'fuzz字典大小：{fuzz_count}')
    for i in range(3):
        random_domain = domain.replace('{fuzz}', exrex.getone(rule))
        logger.log('ALERT', f'请注意检查随机生成的{random_domain}是否正确')
    logger.log('ALERT', f'你有10秒检查时间退出使用`CTRL+C`')
    try:
        time.sleep(5)
    except KeyboardInterrupt:
        logger.log('INFOR', '爆破终止')
        exit(0)
    parts = domain.split('{fuzz}')
    for fuzz in exrex.generate(rule):
        fuzz_domain = parts[0] + fuzz + parts[1]
        domains.add(fuzz_domain)
    return domains


def gen_brute_domains(domain, path):
    domains = set()
    with open(path) as file:
        for line in file:
            brute_domain = line.strip() + '.' + domain
            domains.add(brute_domain)
    logger.log('INFOR', f'爆破字典大小：{len(domains)}')
    return domains


class AIOBrute(Module):
    """
    OneForAll多进程多协程异步子域爆破模块

    Example：
        python aiobrute.py --target example.com run
        python aiobrute.py --target ./domains.txt run
        python aiobrute.py --target example.com --processes 4 --coroutine 64 --wordlist data/subdomains.txt run
        python aiobrute.py --target example.com --recursive True --depth 2 --namelist data/next_subdomains.txt run
        python aiobrute.py --target www.{fuzz}.example.com --fuzz True --rule [a-z][0-9] run

    :param str target:       单个域名或者每行一个域名的文件路径
    :param int processes:    爆破的进程数(默认CPU核心数)
    :param int coroutine:    每个爆破进程下的协程数(默认16)
    :param str wordlist:     指定爆破所使用的字典路径(默认使用config.py配置)
    :param bool recursive:   是否使用递归爆破(默认False)
    :param int depth:        递归爆破的深度(默认2)
    :param str namelist:     指定递归爆破所使用的字典路径(默认使用config.py配置)
    :param bool fuzz:        是否使用fuzz模式进行爆破(默认False，开启必须指定fuzz正则规则)
    :param str rule:         fuzz模式使用的正则规则(默认使用config.py配置)
    """

    def __init__(self, target, processes=None, coroutine=64, wordlist=None,
                 recursive=False, depth=2, namelist=None, fuzz=False, rule=None):
        Module.__init__(self)
        self.domains = set()
        self.domain = str()
        self.module = 'Brute'
        self.source = 'AIOBrute'
        self.target = target
        self.processes = processes or config.brute_processes_num or os.cpu_count()
        self.coroutine = coroutine or config.brute_coroutine_num
        self.wordlist = wordlist or config.brute_wordlist_path or get_wordlist('subdomains.txt')
        self.recursive_brute = recursive or config.enable_recursive_brute
        self.recursive_depth = depth or config.brute_recursive_depth
        self.recursive_namelist = namelist or config.recursive_namelist_path or get_wordlist('next_subdomains.txt')
        self.fuzz = fuzz or config.enable_fuzz
        self.rule = rule or config.fuzz_rule
        self.nameservers = config.resolver_nameservers
        self.ips_times = dict()  # IP集合出现次数
        self.enable_wildcard = False  # 当前域名是否使用泛解析
        self.wildcard_ips = set()  # 泛解析IP集合
        self.wildcard_ttl = int()  # 泛解析TTL整型值

    def gen_tasks(self, domain):
        logger.log('INFOR', f'正在生成{domain}的字典')
        if self.domain != domain:  # 如果domain不是self.domain，而是self.domain的子域 生成递归爆破字典
            domains = gen_brute_domains(domain, self.recursive_namelist)
        elif self.fuzz and self.rule:  # 开启fuzz模式并指定了fuzz正则规则
            domains = gen_fuzz_domains(domain, self.rule)
        else:
            domains = gen_brute_domains(domain, self.wordlist)
        domains = list(domains)
        return utils.split_list(domains, 500)  # 分割任务组 500个子域为一组任务

    def deal_results(self, results):
        for result in results:
            if isinstance(result, Exception):
                # logger.log('DEBUG', f'爆破{subdomain}时出错 {str(answers)}')
                continue
            if isinstance(result, tuple):
                subdomain, answers = result
                ips = {record.host for record in answers}
                value = self.ips_times.setdefault(str(ips), 0)  # 取值 如果是首次出现的IP集合 出现次数先赋值0
                self.ips_times[str(ips)] = value + 1
                ttl = answers[0].ttl
                if self.enable_wildcard:
                    if wildcard_by_compare(ips, ttl, self.wildcard_ips, self.wildcard_ttl):
                        continue
                if wildcard_by_times(ips, self.ips_times):
                    continue
                logger.log('INFOR', f'发现{self.domain}的子域: {subdomain} 解析IP: {ips} TTL: {ttl}')
                self.subdomains.add(subdomain)
                self.records[subdomain] = str(ips)

    async def main(self, domain, rx_queue):
        if not self.fuzz:  # fuzz模式不探测域名是否使用泛解析
            self.enable_wildcard, self.wildcard_ips, self.wildcard_ttl = detect_wildcard(domain)
        tasks = self.gen_tasks(domain)
        logger.log('INFOR', f'正在爆破{domain}的域名')
        for task in tqdm.tqdm(tasks, desc='Progress', smoothing=1.0, ncols=True):
            async with aiomultiprocess.Pool(processes=self.processes, initializer=init_worker,
                                            childconcurrency=self.coroutine) as pool:
                try:
                    results = await pool.map(resolve.aiodns_query_a, task)
                except KeyboardInterrupt:
                    logger.log('ALERT', '爆破终止正在退出')
                    pool.terminate()  # 关闭pool，结束工作进程，不在处理未完成的任务。
                    self.save_json()
                    self.gen_result()
                    rx_queue.put(self.results)
                    return
                else:
                    self.deal_results(results)
            self.save_json()
            self.gen_result()
            rx_queue.put(self.results)

    def run(self, rx_queue=None):
        self.domains = utils.get_domains(self.target)
        while self.domains:
            self.domain = self.domains.pop()
            start = time.time()
            db_conn = database.connect_db()
            table_name = self.domain.replace('.', '_')
            database.create_table(db_conn, table_name)
            if not rx_queue:
                rx_queue = queue.Queue()
            logger.log('INFOR', f'开始执行{self.source}模块爆破域名{self.domain}')
            logger.log('INFOR', f'{self.source}模块使用{self.processes}个进程乘{self.coroutine}个协程')
            # logger.log('INFOR', f'{self.source}模块使用个进程乘{self.coroutine}个协程')
            if self.recursive_brute and not self.fuzz:  # fuzz模式不使用递归爆破
                logger.log('INFOR', f'开始递归爆破{self.domain}的第1层子域')
            loop = asyncio.get_event_loop()
            asyncio.set_event_loop(loop)
            loop.run_until_complete(self.main(self.domain, rx_queue))

            # 递归爆破下一层的子域
            if self.recursive_brute and not self.fuzz:  # fuzz模式不使用递归爆破
                for layer_num in range(1, self.recursive_depth):  # 之前已经做过1层子域爆破 当前实际递归层数是layer+1
                    logger.log('INFOR', f'开始递归爆破{self.domain}的第{layer_num + 1}层子域')
                    for subdomain in self.subdomains.copy():
                        if subdomain.count('.') - self.domain.count('.') == layer_num:  # 进行下一层子域爆破的限制条件
                            loop.run_until_complete(self.main(subdomain, rx_queue))

            while not rx_queue.empty():  # 队列不空就一直取数据存数据库
                database.save_db(db_conn, table_name, rx_queue.get())  # 将结果存入数据库中
            database.copy_table(db_conn, table_name)
            database.deduplicate_subdomain(db_conn, table_name)
            database.remove_invalid(db_conn, table_name)

            end = time.time()
            self.elapsed = round(end - start, 1)
            logger.log('INFOR', f'结束执行{self.source}模块爆破域名{self.domain}')
            logger.log('INFOR', f'{self.source}模块耗时{self.elapsed}秒发现{self.domain}的域名{len(self.subdomains)}个')
            logger.log('DEBUG', f'{self.source}模块发现{self.domain}的的域名 {self.subdomains}')


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

    :param str domain: 域名
    :param result: 结果集队列
    """
    brute = AIOBrute(domain)
    brute.run(result)


if __name__ == '__main__':
    fire.Fire(AIOBrute)
    # result_queue = queue.Queue()
    # do('example.com', result_queue)