#!/data/data/com.termux/files/usr/bin/python3
"""
Resuelve hostname.local via mDNS multicast.
Uso: mdns-resolve <hostname.local>
Retorna: IP address o exit 1
"""
import sys
import socket
import struct
import select

def build_mdns_query(hostname: str) -> bytes:
    """Construye query DNS para hostname."""
    # Header: ID=0, Flags=0, QDCOUNT=1, ANCOUNT=0, NSCOUNT=0, ARCOUNT=0
    header = struct.pack('>HHHHHH', 0, 0, 1, 0, 0, 0)

    # Question section
    question = b''
    for part in hostname.rstrip('.').split('.'):
        question += bytes([len(part)]) + part.encode()
    question += b'\x00'  # End of name
    question += struct.pack('>HH', 1, 1)  # Type A, Class IN

    return header + question

def parse_mdns_response(data: bytes, hostname: str) -> str | None:
    """Parsea respuesta mDNS y extrae IP."""
    if len(data) < 12:
        return None

    # Skip header, find answers
    qdcount = struct.unpack('>H', data[4:6])[0]
    ancount = struct.unpack('>H', data[6:8])[0]

    if ancount == 0:
        return None

    # Skip questions
    offset = 12
    for _ in range(qdcount):
        while offset < len(data) and data[offset] != 0:
            if data[offset] & 0xc0 == 0xc0:  # Compression
                offset += 2
                break
            offset += data[offset] + 1
        else:
            offset += 1
        offset += 4  # Type + Class

    # Parse answers
    for _ in range(ancount):
        # Skip name (handle compression)
        while offset < len(data):
            if data[offset] & 0xc0 == 0xc0:
                offset += 2
                break
            elif data[offset] == 0:
                offset += 1
                break
            else:
                offset += data[offset] + 1

        if offset + 10 > len(data):
            break

        rtype, rclass, ttl, rdlength = struct.unpack('>HHIH', data[offset:offset+10])
        offset += 10

        # Type A (1) = IPv4 address
        if rtype == 1 and rdlength == 4 and offset + 4 <= len(data):
            ip = socket.inet_ntoa(data[offset:offset+4])
            return ip

        offset += rdlength

    return None

def resolve_mdns(hostname: str, timeout: float = 2.0) -> str | None:
    """Resuelve hostname via mDNS multicast."""
    if not hostname.endswith('.local'):
        hostname += '.local'

    MDNS_ADDR = '224.0.0.251'
    MDNS_PORT = 5353

    try:
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        sock.setblocking(False)

        # Enviar query
        query = build_mdns_query(hostname)
        sock.sendto(query, (MDNS_ADDR, MDNS_PORT))

        # Esperar respuesta
        end_time = __import__('time').time() + timeout
        while __import__('time').time() < end_time:
            ready, _, _ = select.select([sock], [], [], 0.1)
            if ready:
                try:
                    data, addr = sock.recvfrom(4096)
                    ip = parse_mdns_response(data, hostname)
                    if ip:
                        sock.close()
                        return ip
                except:
                    pass

        sock.close()
    except Exception as e:
        pass

    return None

def main():
    if len(sys.argv) < 2:
        print("Uso: mdns-resolve <hostname.local>", file=sys.stderr)
        sys.exit(1)

    hostname = sys.argv[1]
    timeout = float(sys.argv[2]) if len(sys.argv) > 2 else 2.0

    ip = resolve_mdns(hostname, timeout)
    if ip:
        print(ip)
        sys.exit(0)
    else:
        sys.exit(1)

if __name__ == "__main__":
    main()