hosts/src/hosts/core/dns.py

"""DNS resolution service for hosts manager.

Provides manual DNS resolution capabilities with timeout handling,
batch processing, and status tracking for hostname to IP address resolution.
"""

import asyncio
import socket
from datetime import datetime
from enum import Enum
from dataclasses import dataclass
from typing import Optional, List
import logging

logger = logging.getLogger(__name__)


@dataclass
class DNSResolutionStatus(Enum):
    """Status of DNS resolution for an entry."""
    NOT_RESOLVED = "not_resolved"
    RESOLVING = "resolving"
    RESOLVED = "resolved"
    RESOLUTION_FAILED = "failed"
    IP_MISMATCH = "mismatch"
    IP_MATCH = "match"


@dataclass
class DNSResolution:
    """Result of DNS resolution for a hostname."""
    hostname: str
    resolved_ip: Optional[str]
    status: DNSResolutionStatus
    resolved_at: datetime
    error_message: Optional[str] = None

    def is_success(self) -> bool:
        """Check if resolution was successful."""
        return self.status == DNSResolutionStatus.RESOLVED and self.resolved_ip is not None

    def get_age_seconds(self) -> float:
        """Get age of resolution in seconds."""
        return (datetime.now() - self.resolved_at).total_seconds()


async def resolve_hostname(hostname: str, timeout: float = 5.0) -> DNSResolution:
    """Resolve a single hostname to IP address with timeout.
    
    Args:
        hostname: Hostname to resolve
        timeout: Maximum time to wait for resolution in seconds
        
    Returns:
        DNSResolution with result and status
    """
    start_time = datetime.now()
    
    try:
        # Use asyncio DNS resolution with timeout
        loop = asyncio.get_event_loop()
        result = await asyncio.wait_for(
            loop.getaddrinfo(hostname, None, family=socket.AF_UNSPEC),
            timeout=timeout
        )
        
        if result:
            # Get first result (usually IPv4)
            ip_address = result[0][4][0]
            return DNSResolution(
                hostname=hostname,
                resolved_ip=ip_address,
                status=DNSResolutionStatus.RESOLVED,
                resolved_at=start_time
            )
        else:
            return DNSResolution(
                hostname=hostname,
                resolved_ip=None,
                status=DNSResolutionStatus.RESOLUTION_FAILED,
                resolved_at=start_time,
                error_message="No address found"
            )
            
    except asyncio.TimeoutError:
        return DNSResolution(
            hostname=hostname,
            resolved_ip=None,
            status=DNSResolutionStatus.RESOLUTION_FAILED,
            resolved_at=start_time,
            error_message=f"Timeout after {timeout}s"
        )
    except Exception as e:
        return DNSResolution(
            hostname=hostname,
            resolved_ip=None,
            status=DNSResolutionStatus.RESOLUTION_FAILED,
            resolved_at=start_time,
            error_message=str(e)
        )


async def resolve_hostnames_batch(hostnames: List[str], timeout: float = 5.0) -> List[DNSResolution]:
    """Resolve multiple hostnames concurrently.
    
    Args:
        hostnames: List of hostnames to resolve
        timeout: Maximum time to wait for each resolution
        
    Returns:
        List of DNSResolution results
    """
    if not hostnames:
        return []
    
    tasks = [resolve_hostname(hostname, timeout) for hostname in hostnames]
    results = await asyncio.gather(*tasks, return_exceptions=True)
    
    # Convert exceptions to failed resolutions
    resolutions = []
    for i, result in enumerate(results):
        if isinstance(result, Exception):
            resolutions.append(DNSResolution(
                hostname=hostnames[i],
                resolved_ip=None,
                status=DNSResolutionStatus.RESOLUTION_FAILED,
                resolved_at=datetime.now(),
                error_message=str(result)
            ))
        else:
            resolutions.append(result)
    
    return resolutions


class DNSService:
    """DNS resolution service for hosts entries."""
    
    def __init__(
        self,
        enabled: bool = True,
        timeout: float = 5.0
    ):
        """Initialize DNS service.
        
        Args:
            enabled: Whether DNS resolution is enabled
            timeout: Timeout for individual DNS queries
        """
        self.enabled = enabled
        self.timeout = timeout
                        
    async def resolve_entry_async(self, hostname: str) -> DNSResolution:
        """Resolve DNS for a hostname asynchronously.
        
        Args:
            hostname: Hostname to resolve
            
        Returns:
            DNSResolution result
        """
        if not self.enabled:
            return DNSResolution(
                hostname=hostname,
                resolved_ip=None,
                status=DNSResolutionStatus.NOT_RESOLVED,
                resolved_at=datetime.now(),
                error_message="DNS resolution is disabled"
            )
        
        return await resolve_hostname(hostname, self.timeout)
            
    async def refresh_entry(self, hostname: str) -> DNSResolution:
        """Manually refresh DNS resolution for hostname.
        
        Args:
            hostname: Hostname to refresh
            
        Returns:
            Fresh DNSResolution result
        """
        return await self.resolve_entry_async(hostname)
        
    async def refresh_all_entries(self, hostnames: List[str]) -> List[DNSResolution]:
        """Manually refresh DNS resolution for multiple hostnames.
        
        Args:
            hostnames: List of hostnames to refresh
            
        Returns:
            List of fresh DNSResolution results
        """
        if not self.enabled:
            return [
                DNSResolution(
                    hostname=hostname,
                    resolved_ip=None,
                    status=DNSResolutionStatus.NOT_RESOLVED,
                    resolved_at=datetime.now(),
                    error_message="DNS resolution is disabled"
                )
                for hostname in hostnames
            ]
        
        return await resolve_hostnames_batch(hostnames, self.timeout)


def compare_ips(stored_ip: str, resolved_ip: str) -> DNSResolutionStatus:
    """Compare stored IP with resolved IP to determine status.
    
    Args:
        stored_ip: IP address stored in hosts entry
        resolved_ip: IP address resolved from DNS
        
    Returns:
        DNSResolutionStatus indicating match or mismatch
    """
    if stored_ip == resolved_ip:
        return DNSResolutionStatus.IP_MATCH
    else:
        return DNSResolutionStatus.IP_MISMATCH