#!/usr/bin/env python
# Copyright Trimble 2017-2018
#
# Desc:
#  Scans the local subnet using UPnP. By default prints all the
#  receivers if finds (Model, IP, FW Version, Friendly name). If
#  the user user adds a model number on the command line it will 
#  see if that string is contained in the model number, the string
#  can be a partial string. For example:
#
#  The following will find all BDxxx receivers
#  python upnpScan.py -m BD
# 
#  The following will find BD935 and BD935-INS
#  python upnpScan.py -m BD935
#

import sys
import socket
from six.moves.urllib.request import urlopen, Request
from six.moves.urllib.error import HTTPError, URLError
from six.moves.queue import Queue
from six import BytesIO
import base64
from lxml import etree
import argparse
import threading
import re

parser = argparse.ArgumentParser(description='Find receivers on the current subnet using UPnP')
group = parser.add_mutually_exclusive_group()
group.add_argument('-m','--model', help='Filter using the model string or a subset (e.g. BD)')
group.add_argument('-f','--friendly', help='Filter using the friendly string or a subset (e.g. Stuart)')
parser.add_argument('-u','--uptime', help='Only provide receivers with uptime *less* than the uptime value in seconds, e.g. -u 3600')
parser.add_argument('-v','--verbose', help='Show extra diags?', action='store_true')
args = parser.parse_args()

# UPnP message
msg = \
  'M-SEARCH * HTTP/1.1\r\n' \
  'HOST:239.255.255.250:1900\r\n' \
  'ST:upnp:rootdevice\r\n' \
  'MX:10\r\n' \
  'MAN:"ssdp:discover"\r\n' \
  '\r\n'

# Start a separate thread for reading UDP data to help prevent socket
# buffer overflows.
def recv_func(q):
  # Set up UDP socket - this triggers the discovery
  s = socket.socket(socket.AF_INET,
                    socket.SOCK_DGRAM,
                    socket.IPPROTO_UDP)
  s.setsockopt(socket.SOL_SOCKET,socket.SO_RCVBUF,4000000)
  if s.getsockopt(socket.SOL_SOCKET,socket.SO_RCVBUF) < 4000000:
    print("NOTE: please increase socket OS buffer size for best performance")
    print("For example, on Linux set net.core.rmem_max=26214400")
  s.settimeout(10)
  # Special IP address
  s.sendto(msg.encode(), ('239.255.255.250', 1900) )
  try:
    while True:
      data, addr = s.recvfrom(4096)
      q.put( (data,addr) )
  except socket.timeout:
    q.put( ('', 0) )
    pass
  s.close()

q = Queue()
recv_thread = threading.Thread(target=recv_func,args=(q,))
recv_thread.daemon = True
recv_thread.start()

def get_and_parse_xml(location):
  request = Request(location)
  
  for password in ['password','tr.imble']:
    login = bytes("admin:" + password,'utf-8')
    base64string = base64.b64encode(login).decode('ascii')
    try:
      request.add_header("Authorization", "Basic %s" % base64string)
      response = urlopen(request, timeout = 2)
      if(response.getcode() == 200):
        ret = response.read()
        data = BytesIO(ret)
        #root_xml = etree.parse(data)
        try:
          root_xml = etree.parse(data)
        except etree.XMLSyntaxError:
          print("************ Invalid XML for %s"%location)
          parser = etree.XMLParser(recover=True)
          root_xml = etree.parse(data, parser)
        response.close()
        return root_xml
        break
    except socket.timeout:
      pass

  return None

# Wait for the UPnP response. This will only get info for the current
# subnet.
while True:
  data, addr = q.get()
  if data == '':
    break
  data = data.decode('utf-8')

  lines = data.split("\n")
  for line in lines:
    # We know Trimble devices point to /upnp.xml
    m=re.match('LOCATION.*?http://(.*?):([0-9]+)/upnp.xml',line)
    if m is not None:
      # This may be a Trimble device, get the IP address
      IP_addr = addr[0]  # from socket recvfrom()
      IP_addr_v2 = m.group(1)  # from UPNP "LOCATION" string
      IP_port = m.group(2)
      if IP_addr != IP_addr_v2:
        # There seems to be a bug in the receiver UPnP address sometimes...
        if args.verbose: print('Warning: IP mismatch: %s vs %s' % (IP_addr, IP_addr_v2))
      location = 'http://%s:%s/upnp.xml' % (IP_addr, IP_port)
      
      try:
        # get upnp.xml
        root_xml = get_and_parse_xml( location )

        if(root_xml == None):
          pass

        # get the modelName, could filter now and only process devices
        # we want
        elem = root_xml.find('.//{*}modelName')

        if elem is not None:
          model = elem.text.rstrip('/')
          friendly = root_xml.find('.//{*}friendlyName')
          friendlyTxt = friendly.text.rstrip('/')

          # If filtering is requested filter before any further
          # processing
          if args.model:
            is_match = args.model in model
          elif args.friendly:
            is_match = args.friendly in friendlyTxt
          else:
            is_match = True
          if is_match:
            # Tweak the URL to add "TrimbleServices.xml" We cannot
            # change the UPnP protocol so have to use standard tags in
            # "upnp.xml". However, we have a secondary file
            # "TrimbleServices.xml" that all coreBuild receivers
            # provide. It provides a lot of service descovery
            # information we can use.
            http_addr = location.split("/upnp.xml")[0]
            location = http_addr.rstrip() + "/TrimbleServices.xml"
            try:
              # Try to download TrimbleServices.xml, if we get this we
              # are almost certainly a receiver

              root_xml = get_and_parse_xml( location )

              # Grab a few of the tags we may want
              if(root_xml.find('.//{*}FWReferenceVersion') is not None):
                FWRef = root_xml.find('.//{*}FWReferenceVersion')
              else:
                FWRef = None
              FWVer = root_xml.find('.//{*}FWVersion')
              SN = root_xml.find('.//{*}serialNumber')

              # Warranty:
              # <warranty>
              #    <year>2030</year>
              #    <month>1</month>
              #    <day>1</day>
              #  </warranty>

              warranty_year  = 'XXXX'
              warranty_month = 'XX'
              warranty_day   = 'XX'
              warranty = root_xml.findall('.//{*}warranty')
              for node in warranty:
                if(warranty):
                  warranty_year  = node.findtext('.//{*}year')
                  warranty_month = node.findtext('.//{*}month')
                  warranty_day   = node.findtext('.//{*}day')

              #print("%s %s %s" % (warranty_year, warranty_month,warranty_day))

              if(friendly is not None and SN is not None and FWVer is not None):
                name  = friendly.text
                serial = SN.text
                Version = FWVer.text

                if(FWRef is not None):
                  refFW = FWRef.text
                else:
                  refFW = []

                try:
                  location = http_addr.rstrip() + "/xml/dynamic/sec_merge.xml?powerData=&power=&sysData="

                  root_xml = get_and_parse_xml( location )

                  day  = root_xml.find('.//{*}day')
                  hour = root_xml.find('.//{*}hour')
                  mins = root_xml.find('.//{*}min')
                  sec  = root_xml.find('.//{*}sec')

                  uptime =(    int(day.text) * 86400
                             + int(hour.text) * 3600
                             + int(mins.text) * 60
                             + int(sec.text) )

                  FWDate  = root_xml.find('.//{*}FWDate').text

                  if((not args.uptime) or (uptime < int(args.uptime))):
                    # Now display some basic information about the device
                    if(refFW):
                      print("%s Uptime = %lds %s FWRef = %s-%s(%s) SN=%s Name = %s" % (model,uptime,http_addr,refFW,Version,FWDate,serial,name))
                    else:
                      print("%s Uptime = %lds %s FW = %s(%s) SN=%s Name = %s" % (model,uptime,http_addr,Version,FWDate,serial,name))
                except (AttributeError,HTTPError,URLError,socket.timeout):
                  if(refFW):
                    print("%s %s FWRef = %s-%s SN=%s Name = %s" % (model,http_addr,refFW,Version,serial,name))
                  else:
                    print("%s %s FW = %s SN=%s Name = %s" % (model,http_addr,Version,serial,name))
                  pass
            except (AttributeError,HTTPError,URLError,socket.timeout):
            #except Exception as e:
              #print('Error! Code: {c}, Message, {m}'.format(c = type(e).__name__, m = str(e)))
              if args.verbose: print("Can't open TrimbleServices: " + addr[0])
              pass
      except (AttributeError,HTTPError,URLError,socket.timeout):
        if args.verbose: print("Can't open upnp: " + addr[0] + ", " + location)
        pass