from typing import Any, Optional
import requests
import dns.resolver as dns_resolver_module
from rich.progress import track


class DNSRecon:
    def dns_resolver(self, AIModels, target: str, akey: Optional[str], bkey: Optional[str], lkey, lendpoint, AI: str) -> Any:
        if target is None:
            raise ValueError("InvalidTarget: Target Not Provided")

        analyze = ''
        record_types = ['A', 'AAAA', 'NS', 'CNAME', 'MX', 'PTR', 'SOA', 'TXT']
        for record_type in track(record_types):
            try:
                answer = dns_resolver_module.resolve(target, record_type)
                for server in answer:
                    st = server.to_text()
                    analyze += f"\n{record_type} : {st}"
            except dns_resolver_module.NoAnswer:
                print('No record Found')
            except dns_resolver_module.NXDOMAIN:
                print('NXDOMAIN record NOT Found')
            except dns_resolver_module.LifetimeTimeout:
                print("Timed out, check your internet")
            except requests.exceptions.InvalidHeader:
                pass
            except KeyboardInterrupt:
                print("Bye")
                quit()

        response = ""
        match AI:
            case 'openai':
                if akey is None:
                    raise ValueError("KeyNotFound: Key Not Provided")
                akey = akey.replace('\n', '')
                response = AIModels.gpt_ai(akey, analyze)
            case 'bard':
                if bkey is None:
                    raise ValueError("KeyNotFound: Key Not Provided")
                bkey = bkey.replace('\n', '')
                response = AIModels.BardAI(bkey, analyze)
            case 'llama':
                response = AIModels.llama_AI(analyze, "local", lkey, lendpoint)
            case 'llama-api':
                response = AIModels.llama_AI(analyze, "runpod", lkey, lendpoint)
            except KeyboardInterrupt:
                print("Bye")
                quit()
        return str(response)