from bs4 import BeautifulSoup
from bs4.element import Tag

from theHarvester.lib.core import AsyncFetcher, Core


class SearchRapidDns:
    def __init__(self, word) -> None:
        self.word = word
        self.total_results: list = []
        self.proxy = False

    async def do_search(self):
        try:
            headers = {'User-agent': Core.get_user_agent()}
            # TODO see if it's worth adding sameip searches
            # f'{self.hostname}/sameip/{self.word}?full=1#result'
            urls = [f'https://rapiddns.io/subdomain/{self.word}?full=1#result']
            responses = await AsyncFetcher.fetch_all(urls, headers=headers, proxy=self.proxy)
            if len(responses[0]) <= 1:
                return self.total_results
            soup = BeautifulSoup(responses[0], 'html.parser')
            table_el = soup.find('table')
            if not isinstance(table_el, Tag):
                return self.total_results
            tbody_el = table_el.find('tbody')
            if not isinstance(tbody_el, Tag):
                return self.total_results
            rows = tbody_el.find_all('tr')
            if rows:
                # Validation check
                for row in rows:
                    if not isinstance(row, Tag):
                        continue
                    cells = row.find_all('td')
                    if len(cells) > 0:
                        # sanity check
                        subdomain = str(cells[0].get_text())
                        if cells[-1].get_text() == 'CNAME':
                            self.total_results.append(f'{subdomain}')
                        else:
                            self.total_results.append(f'{subdomain}:{str(cells[1].get_text()).strip()}')
                self.total_results = list({domain for domain in self.total_results})
        except Exception as e:
            print(f'An exception has occurred: {e!s}')

    async def process(self, proxy: bool = False) -> None:
        self.proxy = proxy
        await self.do_search()

    async def get_hostnames(self):
        return self.total_results
