import asyncio
import logging
from ipaddress import ip_address

from theHarvester.discovery.constants import MissingKey
from theHarvester.lib.core import AsyncFetcher, Core
from theHarvester.lib.hostnames import normalize_scoped_hostname

logger = logging.getLogger(__name__)


class SearchPentestTools:
    def __init__(self, word) -> None:
        self.word = word.strip().lower().rstrip('.')
        self.key = Core.pentest_tools_key()
        if self.key is None:
            raise MissingKey('PentestTools')
        self.totalhosts: set[str] = set()
        self.totalips: set[str] = set()
        self.api = 'https://app.pentest-tools.com/api/v2'
        self.headers = {
            'Accept': 'application/json',
            'Authorization': f'Bearer {self.key}',
            'Content-Type': 'application/json',
        }
        self.proxy = False

    @staticmethod
    def _response_data(response: object) -> dict | None:
        if not isinstance(response, dict) or not isinstance(data := response.get('data'), dict):
            logger.info('Pentest-Tools returned a malformed response')
            return None
        return data

    async def poll(self, scan_id: int) -> None:
        for _attempt in range(10):
            await asyncio.sleep(3)
            status = self._response_data(
                await AsyncFetcher.fetch(
                    url=f'{self.api}/scans/{scan_id}',
                    headers=self.headers,
                    json=True,
                    proxy=self.proxy,
                )
            )
            if status is None:
                return
            status_name = status.get('status_name')
            if not isinstance(status_name, str):
                logger.info('Pentest-Tools returned a malformed status response')
                return
            if status_name in {'waiting', 'running'}:
                continue
            if status_name == 'finished':
                output = self._response_data(
                    await AsyncFetcher.fetch(
                        url=f'{self.api}/scans/{scan_id}/output',
                        headers=self.headers,
                        json=True,
                        proxy=self.proxy,
                    )
                )
                if output is not None:
                    await self.parse_json(output)
            else:
                logger.info('Pentest-Tools scan did not finish successfully')
            return
        logger.info('Pentest-Tools scan is still waiting after 10 status checks')

    async def parse_json(self, json_results) -> None:
        if json_results.get('output_type') != 'subdomain_list':
            return
        try:
            output_data = json_results['output_data']['subdomains']
        except (KeyError, TypeError):
            return
        if not isinstance(output_data, list):
            return
        for result in output_data:
            if not isinstance(result, dict):
                continue
            hostname = normalize_scoped_hostname(result.get('hostname'), self.word)
            if not hostname:
                continue
            if hostname != self.word:
                self.totalhosts.add(hostname)
            address_value = result.get('ip_address')
            if not isinstance(address_value, str):
                continue
            try:
                address = str(ip_address(address_value))
            except ValueError:
                continue
            self.totalips.add(address)

    async def get_hostnames(self) -> set[str]:
        return self.totalhosts

    async def get_ips(self) -> set[str]:
        return self.totalips

    async def do_search(self) -> None:
        # Pentest-Tools documents Subdomain Finder as tool 20:
        # https://pentest-tools.com/docs/api-reference/scans/start-a-scan
        subdomain_payload = {
            'tool_id': 20,
            'target_name': self.word,
            'tool_params': {
                'scan_type': 'light',
                'web_details': False,
                'unresolved_results': True,
            },
        }
        response = self._response_data(
            await AsyncFetcher.post_fetch(
                url=f'{self.api}/scans',
                headers=self.headers,
                json_body=subdomain_payload,
                json=True,
                proxy=self.proxy,
            )
        )
        if response is None:
            return
        scan_id = response.get('created_id')
        if not isinstance(scan_id, int) or isinstance(scan_id, bool):
            logger.info('Pentest-Tools returned a malformed start response')
            return
        await self.poll(scan_id)

    async def process(self, proxy: bool = False) -> None:
        self.proxy = proxy
        await self.do_search()  # Only need to do it once.
