# Copyright 2015 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Module that enumerates mDNS services."""

from __future__ import print_function

import collections
import dpkt
import select
import socket
import time


MDNS_MULTICAST_ADDRESS = '224.0.0.251'
MDNS_PORT = 5353
PACKET_BUFFER_SIZE = 2048


Service = collections.namedtuple(
    'Service',
    ('hostname', 'ip', 'port', 'ptrname', 'text'))


def _ParseResponse(data):
  """Parse mDNS network response into Service tuple.

  Args:
    data: Byte-blob of mDNS response.

  Returns:
    None if the mDNS response is invalid or incomplete, else a fully populated
    |Service| tuple.
  """
  hostname = None
  ip = None
  port = None
  ptrname = None
  text = None

  try:
    dns_response = dpkt.dns.DNS(data)
  except (dpkt.Error, dpkt.NeedData, dpkt.UnpackError):
    # Ignore bad mDNS response.
    return None

  for rr in dns_response.an:
    if rr.type == dpkt.dns.DNS_A:
      ip = socket.inet_ntoa(rr.ip)
    elif rr.type == dpkt.dns.DNS_PTR:
      ptrname = rr.ptrname
    elif rr.type == dpkt.dns.DNS_SRV:
      hostname = rr.srvname
      port = rr.port
    elif rr.type == dpkt.dns.DNS_TXT:
      text = dict(entry.split('=', 1) for entry in rr.text)
  service = Service(hostname, ip, port, ptrname, text)

  # Ignore incomplete responses.
  if any(x is None for x in service):
    return None
  return service


def FindServices(source_ip, service_name, should_add_func=None,
                 should_continue_func=None, timeout_seconds=1):
  """Find all instances of |service_name| on the network.

  For each service found, |should_add_func| is called with the service
  information to determine wheter the service should be added to the results
  list.  This method exits after |timeout_seconds| or earlier if instructed by
  the return value from |should_continue_func|.

  Args:
    source_ip: IP address of the network interface to use for service discovery.
    service_name: Name of mDNS service to discover (eg. '_ssh._tcp.local').
    should_add_func: Function called for each service found to determine if the
      service should be added to the results list.  If None is specified, all
      services found are added to the results list.
    should_continue_func: Function called for each service found to determine
      whether to continue with service discovery.  If None is specified,
      service discovery continues until the timeout expires.
    timeout_seconds: Number of seconds to wait for services to respond to
      discovery request.

  Returns:
    List of |Service| found.

  """
  # Default callback functions that add all services to the results list and
  # continue service discovery until the timeout expires.
  if not should_add_func:
    should_add_func = lambda service: True
  if not should_continue_func:
    should_continue_func = lambda service: True

  net_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
  net_socket.bind((source_ip, 0))
  query = [dpkt.dns.DNS.Q(name=service_name, type=dpkt.dns.DNS_PTR)]
  dns_packet = dpkt.dns.DNS(op=dpkt.dns.DNS_QUERY, qd=query)
  net_socket.sendto(str(dns_packet),
                    (MDNS_MULTICAST_ADDRESS, MDNS_PORT))

  results = []
  remaining_time = timeout_seconds
  end_time = time.time() + remaining_time
  while remaining_time > 0:
    read_ready, _, _ = select.select([net_socket], [], [], remaining_time)
    if not read_ready:
      break

    data, _ = net_socket.recvfrom(PACKET_BUFFER_SIZE)
    service = _ParseResponse(data)
    if service:
      if should_add_func(service):
        results.append(service)
      if not should_continue_func(service):
        break

    remaining_time = end_time - time.time()

  return results