from bbot.modules.templates.subdomain_enum import subdomain_enum


class dnsbrute(subdomain_enum):
    flags = ["subdomain-enum", "active", "aggressive"]
    watched_events = ["DNS_NAME"]
    produced_events = ["DNS_NAME"]
    meta = {
        "description": "Brute-force subdomains with massdns + static wordlist",
        "author": "@TheTechromancer",
        "created_date": "2024-04-24",
    }
    options = {
        "wordlist": "https://raw.githubusercontent.com/danielmiessler/SecLists/master/Discovery/DNS/subdomains-top1million-5000.txt",
        "max_depth": 5,
    }
    options_desc = {
        "wordlist": "Subdomain wordlist URL",
        "max_depth": "How many subdomains deep to brute force, i.e. 5.4.3.2.1.evilcorp.com",
    }
    deps_common = ["massdns"]
    reject_wildcards = "strict"
    dedup_strategy = "lowest_parent"
    _qsize = 10000

    async def setup(self):
        self.max_depth = max(1, self.config.get("max_depth", 5))
        self.subdomain_file = await self.helpers.wordlist(self.config.get("wordlist"))
        self.subdomain_list = set(self.helpers.read_file(self.subdomain_file))
        self.wordlist_size = len(self.subdomain_list)
        return await super().setup()

    async def filter_event(self, event):
        eligible, reason = await super().filter_event(event)
        query = self.make_query(event)

        # limit brute force depth
        subdomain_depth = self.helpers.subdomain_depth(query) + 1
        if subdomain_depth > self.max_depth:
            eligible = False
            reason = f"subdomain depth of *.{query} ({subdomain_depth}) > max_depth ({self.max_depth})"

        # don't brute-force things that look like autogenerated PTRs
        if self.helpers.dns.brute.has_excessive_digits(query):
            eligible = False
            reason = f'"{query}" looks like an autogenerated PTR'

        return eligible, reason

    async def handle_event(self, event):
        query = self.make_query(event)
        self.info(f"Brute-forcing {self.wordlist_size:,} subdomains for {query} (source: {event.data})")
        for hostname in await self.helpers.dns.brute(self, query, self.subdomain_list):
            await self.emit_event(
                hostname,
                "DNS_NAME",
                parent=event,
                context=f'{{module}} tried {self.wordlist_size:,} subdomains against "{query}" and found {{event.type}}: {{event.data}}',
            )
