"""
DNS resolution provider for DNSRecon.
Discovers domain relationships through DNS record analysis.
"""

import socket
import dns.resolver
import dns.reversename
from typing import List, Dict, Any, Tuple, Optional
from .base_provider import BaseProvider
from core.graph_manager import RelationshipType, NodeType


class DNSProvider(BaseProvider):
    """
    Provider for standard DNS resolution and reverse DNS lookups.
    Discovers domain-to-IP and IP-to-domain relationships through DNS records.
    """
    
    def __init__(self):
        """Initialize DNS provider with appropriate rate limiting."""
        super().__init__(
            name="dns",
            rate_limit=100,  # DNS queries can be faster
            timeout=10
        )
        
        # Configure DNS resolver
        self.resolver = dns.resolver.Resolver()
        self.resolver.timeout = 5
        self.resolver.lifetime = 10
    
    def get_name(self) -> str:
        """Return the provider name."""
        return "dns"
    
    def is_available(self) -> bool:
        """DNS is always available - no API key required."""
        return True
    
    def query_domain(self, domain: str) -> List[Tuple[str, str, RelationshipType, float, Dict[str, Any]]]:
        """
        Query DNS records for the domain to discover relationships.
        
        Args:
            domain: Domain to investigate
            
        Returns:
            List of relationships discovered from DNS analysis
        """
        if not self._is_valid_domain(domain):
            return []
        
        relationships = []
        
        # Query A records
        relationships.extend(self._query_a_records(domain))
        
        # Query AAAA records (IPv6)
        relationships.extend(self._query_aaaa_records(domain))
        
        # Query CNAME records
        relationships.extend(self._query_cname_records(domain))
        
        # Query MX records
        relationships.extend(self._query_mx_records(domain))
        
        # Query NS records
        relationships.extend(self._query_ns_records(domain))
        
        return relationships
    
    def query_ip(self, ip: str) -> List[Tuple[str, str, RelationshipType, float, Dict[str, Any]]]:
        """
        Query reverse DNS for the IP address.
        
        Args:
            ip: IP address to investigate
            
        Returns:
            List of relationships discovered from reverse DNS
        """
        if not self._is_valid_ip(ip):
            return []
        
        relationships = []
        
        try:
            # Perform reverse DNS lookup
            reverse_name = dns.reversename.from_address(ip)
            response = self.resolver.resolve(reverse_name, 'PTR')
            
            for ptr_record in response:
                hostname = str(ptr_record).rstrip('.')
                
                if self._is_valid_domain(hostname):
                    raw_data = {
                        'query_type': 'PTR',
                        'ip_address': ip,
                        'hostname': hostname,
                        'ttl': response.ttl
                    }
                    
                    relationships.append((
                        ip,
                        hostname,
                        RelationshipType.A_RECORD,  # Reverse relationship
                        RelationshipType.A_RECORD.default_confidence,
                        raw_data
                    ))
                    
                    self.log_relationship_discovery(
                        source_node=ip,
                        target_node=hostname,
                        relationship_type=RelationshipType.A_RECORD,
                        confidence_score=RelationshipType.A_RECORD.default_confidence,
                        raw_data=raw_data,
                        discovery_method="reverse_dns_lookup"
                    )
        
        except Exception as e:
            self.logger.logger.debug(f"Reverse DNS lookup failed for {ip}: {e}")
        
        return relationships
    
    def _query_a_records(self, domain: str) -> List[Tuple[str, str, RelationshipType, float, Dict[str, Any]]]:
        """Query A records for the domain."""
        relationships = []
        
        #if not DNS_AVAILABLE:
        #    return relationships
        
        try:
            response = self.resolver.resolve(domain, 'A')
            
            for a_record in response:
                ip_address = str(a_record)
                
                raw_data = {
                    'query_type': 'A',
                    'domain': domain,
                    'ip_address': ip_address,
                    'ttl': response.ttl
                }
                
                relationships.append((
                    domain,
                    ip_address,
                    RelationshipType.A_RECORD,
                    RelationshipType.A_RECORD.default_confidence,
                    raw_data
                ))
                
                self.log_relationship_discovery(
                    source_node=domain,
                    target_node=ip_address,
                    relationship_type=RelationshipType.A_RECORD,
                    confidence_score=RelationshipType.A_RECORD.default_confidence,
                    raw_data=raw_data,
                    discovery_method="dns_a_record"
                )
        
        except Exception as e:
            self.logger.logger.debug(f"A record query failed for {domain}: {e}")
        
        return relationships
    
    def _query_aaaa_records(self, domain: str) -> List[Tuple[str, str, RelationshipType, float, Dict[str, Any]]]:
        """Query AAAA records (IPv6) for the domain."""
        relationships = []
        
        #if not DNS_AVAILABLE:
        #    return relationships
        
        try:
            response = self.resolver.resolve(domain, 'AAAA')
            
            for aaaa_record in response:
                ip_address = str(aaaa_record)
                
                raw_data = {
                    'query_type': 'AAAA',
                    'domain': domain,
                    'ip_address': ip_address,
                    'ttl': response.ttl
                }
                
                relationships.append((
                    domain,
                    ip_address,
                    RelationshipType.A_RECORD,  # Using same type for IPv6
                    RelationshipType.A_RECORD.default_confidence,
                    raw_data
                ))
                
                self.log_relationship_discovery(
                    source_node=domain,
                    target_node=ip_address,
                    relationship_type=RelationshipType.A_RECORD,
                    confidence_score=RelationshipType.A_RECORD.default_confidence,
                    raw_data=raw_data,
                    discovery_method="dns_aaaa_record"
                )
        
        except Exception as e:
            self.logger.logger.debug(f"AAAA record query failed for {domain}: {e}")
        
        return relationships
    
    def _query_cname_records(self, domain: str) -> List[Tuple[str, str, RelationshipType, float, Dict[str, Any]]]:
        """Query CNAME records for the domain."""
        relationships = []
        
        #if not DNS_AVAILABLE:
        #    return relationships
        
        try:
            response = self.resolver.resolve(domain, 'CNAME')
            
            for cname_record in response:
                target_domain = str(cname_record).rstrip('.')
                
                if self._is_valid_domain(target_domain):
                    raw_data = {
                        'query_type': 'CNAME',
                        'source_domain': domain,
                        'target_domain': target_domain,
                        'ttl': response.ttl
                    }
                    
                    relationships.append((
                        domain,
                        target_domain,
                        RelationshipType.CNAME_RECORD,
                        RelationshipType.CNAME_RECORD.default_confidence,
                        raw_data
                    ))
                    
                    self.log_relationship_discovery(
                        source_node=domain,
                        target_node=target_domain,
                        relationship_type=RelationshipType.CNAME_RECORD,
                        confidence_score=RelationshipType.CNAME_RECORD.default_confidence,
                        raw_data=raw_data,
                        discovery_method="dns_cname_record"
                    )
        
        except Exception as e:
            self.logger.logger.debug(f"CNAME record query failed for {domain}: {e}")
        
        return relationships
    
    def _query_mx_records(self, domain: str) -> List[Tuple[str, str, RelationshipType, float, Dict[str, Any]]]:
        """Query MX records for the domain."""
        relationships = []
        
        #if not DNS_AVAILABLE:
        #    return relationships
        
        try:
            response = self.resolver.resolve(domain, 'MX')
            
            for mx_record in response:
                mx_host = str(mx_record.exchange).rstrip('.')
                
                if self._is_valid_domain(mx_host):
                    raw_data = {
                        'query_type': 'MX',
                        'domain': domain,
                        'mx_host': mx_host,
                        'priority': mx_record.preference,
                        'ttl': response.ttl
                    }
                    
                    relationships.append((
                        domain,
                        mx_host,
                        RelationshipType.MX_RECORD,
                        RelationshipType.MX_RECORD.default_confidence,
                        raw_data
                    ))
                    
                    self.log_relationship_discovery(
                        source_node=domain,
                        target_node=mx_host,
                        relationship_type=RelationshipType.MX_RECORD,
                        confidence_score=RelationshipType.MX_RECORD.default_confidence,
                        raw_data=raw_data,
                        discovery_method="dns_mx_record"
                    )
        
        except Exception as e:
            self.logger.logger.debug(f"MX record query failed for {domain}: {e}")
        
        return relationships
    
    def _query_ns_records(self, domain: str) -> List[Tuple[str, str, RelationshipType, float, Dict[str, Any]]]:
        """Query NS records for the domain."""
        relationships = []
        
        #if not DNS_AVAILABLE:
        #    return relationships
        
        try:
            response = self.resolver.resolve(domain, 'NS')
            
            for ns_record in response:
                ns_host = str(ns_record).rstrip('.')
                
                if self._is_valid_domain(ns_host):
                    raw_data = {
                        'query_type': 'NS',
                        'domain': domain,
                        'ns_host': ns_host,
                        'ttl': response.ttl
                    }
                    
                    relationships.append((
                        domain,
                        ns_host,
                        RelationshipType.NS_RECORD,
                        RelationshipType.NS_RECORD.default_confidence,
                        raw_data
                    ))
                    
                    self.log_relationship_discovery(
                        source_node=domain,
                        target_node=ns_host,
                        relationship_type=RelationshipType.NS_RECORD,
                        confidence_score=RelationshipType.NS_RECORD.default_confidence,
                        raw_data=raw_data,
                        discovery_method="dns_ns_record"
                    )
        
        except Exception as e:
            self.logger.logger.debug(f"NS record query failed for {domain}: {e}")
        
        return relationships