#!/usr/bin/env python
# vim:ts=2:et:sw=2:ai
#
# Scan Wireless Leiden Network and report status of links and nodes
#
# Rick van der Zwet <info@rickvanderzwet.nl>

from pprint import pprint
from xml.dom.minidom import parse, parseString
import gformat
import os.path
import re
import subprocess
import sys
import time
import yaml
from datetime import datetime

# When force is used as argument, use this range
DEFAULT_SCAN_RANGE= ['172.16.0.0/21']

# Default node status output
nodemap_status_file = '/tmp/nodemap_status.yaml'

#
# BEGIN nmap XML parser
# XXX: Should properly go to seperate class/module
def get_attribute(node,attr):
  return node.attributes[attr].value

def attribute_from_node(parent,node,attr):
  return parent.getElementsByTagName(node)[0].attributes[attr].value

def parse_port(node):
  item = dict()
  item['protocol'] = get_attribute(node,'protocol')
  item['portid'] = get_attribute(node,'portid')
  item['state'] = attribute_from_node(node,'state','state')
  item['reason'] = attribute_from_node(node,'state','reason')
  return item

def parse_ports(node):
  item = dict()
  for port in node.getElementsByTagName('port'):
    port_item = parse_port(port)
    item[port_item['portid']] = port_item
  return item

def parse_host(node):
  # Host status
  item = dict()
  item['state'] = attribute_from_node(node,'status','state')
  item['reason'] = attribute_from_node(node,'status','reason')
  item['addr'] = attribute_from_node(node,'address','addr') 
  item['addrtype'] = attribute_from_node(node,'address','addrtype')

  # Service status
  ports = node.getElementsByTagName('ports')
  if ports:
    item['port'] = parse_ports(ports[0])
  return item

def parse_nmap(root):
  status = dict()
  for node in root.childNodes[2].getElementsByTagName('host'):
    scan = parse_host(node)
    if not status.has_key(scan['addr']):
      status[scan['addr']] = scan
  return status
#
# END nmap parser
#



def _do_nmap_scan(command, iphosts):
  """ Run/Read nmap XML with various choices"""
  command = "nmap -n -iL - -oX - %s" %(command)
  print "# New run '%s', can take a while to complete" % (command)
  p = subprocess.Popen(command.split(),
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        stdin=subprocess.PIPE, bufsize=-1)

  (stdoutdata, stderrdata) = p.communicate("\n".join(iphosts))
  if p.returncode != 0:
    print "# [ERROR] nmap failed to complete '%s'" % stderrdata
    sys.exit(1) 

  dom = parseString(stdoutdata)
  return (parse_nmap(dom),stdoutdata)



def do_nmap_scan(command, iphosts, result_file=None, forced_scan=False):
  """ Wrapper around _run_nmap to get listing of all hosts, the default nmap
      does not return results for failed hosts"""
  # Get all hosts to be processed
  (init_status, stdoutdata) = _do_nmap_scan(" -sL",iphosts)

  # Return stored file if exists
  if not forced_scan and result_file and os.path.exists(result_file)  \
    and os.path.getsize(result_file) > 0:
    print "# Reading stored NMAP results from '%s'" % (result_file)
    status = parse_nmap(parse(result_file))
  else:
    # New scan
    (status, stdoutdata) = _do_nmap_scan(command, iphosts)

    # Store result if requested
    if result_file:
      print "# Saving results in %s" % (result_file)
      f = file(result_file,'w')
      f.write(stdoutdata)
      f.close()

  init_status.update(status)
  return init_status



def do_snmpwalk(host, oid):
   """ Do snmpwalk, returns (p, stdout, stderr)"""
   # GLobal SNMP walk options
   snmpwalk = ('snmpwalk -r 0 -t 1 -OX -c public -v 2c %s' % host).split()
   p = subprocess.Popen(snmpwalk + [oid],
       stdout=subprocess.PIPE, stderr=subprocess.PIPE)
   (stdoutdata, stderrdata) = p.communicate()
   stdoutdata = stdoutdata.split('\n')[:-1]
   stderrdata = stderrdata.split('\n')[:-1]
   return (p, stdoutdata, stderrdata)




def do_snmp_scan(iphosts, status, stored_status=dict(), forced_scan=False):
  """ SNMP scanning, based on results fould in NMAP scan"""
  mac_to_host = dict()
  host_processed = dict()

  #
  # Gather SNMP data from hosts
  for host in iphosts:
    # Status might be containing old hosts as well and visa-versa
    if not status.has_key(host):
       print "## [ERROR] No nmap result found"
       continue

    scan = status[host]
    if scan['state'] != "up":
      continue
  
    print '# Processing host %s' % host
    # IP -> Mac addresses found in host ARP table, with key IP
    status[host]['arpmac'] = dict()
    # MAC -> iface addresses, with key MAC
    status[host]['mac'] = dict()
    # Mirrored: iface -> MAC addresses, with key interface name
    status[host]['iface'] = dict()
    try:
      if not forced_scan and stored_status[host]['snmp_retval'] != 0:
        print "## SNMP Connect failed last time, ignoring"
        continue
    except:
      pass
  
    stored_status[host] = dict()
    if not "open" in scan['port']['161']['state']:
      print "## [ERROR] SNMP port not opened"
      continue
  
    (p, output, stderrdata) = do_snmpwalk(host, 'SNMPv2-MIB::sysDescr')
    stored_status[host]['snmp_retval'] = p.returncode
    # Assume host remain reachable troughout all the SNMP queries
    if p.returncode != 0:
       print "## [ERROR] SNMP failed '%s'" % ",".join(stderrdata)
       continue
  
    # Get some host details
    # SNMPv2-MIB::sysDescr.0 = STRING: FreeBSD CNodeSOM1.wLeiden.NET
    # 8.0-RELEASE-p2 FreeBSD 8.0-RELEASE-p2 #2: Fri Feb 19 18:24:23 CET 2010
    # root@80fab2:/usr/obj/nanobsd.wleiden/usr/src/sys/kernel.wleiden i386
    status[host]['sys_desc'] = output[0]
    hostname = output[0].split(' ')[4]
    release = output[0].split(' ')[5]
    stored_status[host]['hostname'] = status[host]['hostname'] = hostname
    stored_status[host]['release'] = status[host]['release'] = release
    print "## %(hostname)s - %(release)s" % stored_status[host]
  
    # Check if the host is already done processing
    # Note: the host is marked done processing at the end
    if host_processed.has_key(hostname):
      print "## Host already processed this run"
      continue  
    
    # Interface list with key the index number
    iface_descr = dict()
    # IF-MIB::ifDescr.1 = STRING: ath0
    r = re.compile('^IF-MIB::ifDescr\[([0-9]+)\] = STRING: ([a-z0-9]+)$')
    (p, output, stderrdata) = do_snmpwalk(host, 'IF-MIB::ifDescr')
    for line in output:
      m = r.match(line)
      iface_descr[m.group(1)] = m.group(2)
  
    # IF-MIB::ifPhysAddress[1] = STRING: 0:80:48:54:bb:52
    r = re.compile('^IF-MIB::ifPhysAddress\[([0-9]+)\] = STRING: ([0-9a-f:]*)$')
    (p, output, stderrdata) = do_snmpwalk(host, 'IF-MIB::ifPhysAddress')
    for line in output:
      m = r.match(line)
      # Ignore lines which has no MAC address
      if not m.group(2): continue
      index = m.group(1)
      # Convert to proper MAC
      mac = ":".join(["%02X" % int(x,16) for x in m.group(2).split(':')])
      if not iface_descr.has_key(index):
        print "## Index cannot be mapped to a key, available:"
        for index, value  in iface_descr.iteritems():
          print "## - %s [%s]" % (value, index)
      else:
        print "## Local MAC %s [index:%s] -> %s" % (iface_descr[index], index, mac)
        status[host]['mac'][mac] = iface_descr[index]
        status[host]['iface'][iface_descr[index]] = mac
      mac_to_host[mac] = hostname
  
    # Process host SNMP status
    (p, output, stderrdata) = do_snmpwalk(host, 'RFC1213-MIB::atPhysAddress')
    # RFC1213-MIB::atPhysAddress[8][1.172.21.160.34] = Hex-STRING: 00 17 C4 CC 5B F2 
    r = re.compile('^RFC1213-MIB::atPhysAddress\[[0-9]+\]\[1\.([0-9\.]+)\] = Hex-STRING: ([0-9A-F\ ]+)$')
    for line in output:
      m = r.match(line)
      ip = m.group(1)
      # Replace spaces in MAC with :
      mac = ":".join(m.group(2).split(' ')[:-1])
      status[host]['arpmac'][ip] = mac
  
      local = '[remote]'
      if mac in status[host]['mac'].keys():
        local = '[local]'
      print "## Arp table MAC %s -> %s %s" % (ip, mac, local)
  
    # Make sure we keep a record of the processed host which ip entry to check
    host_processed[hostname] = host

  stored_status['host_processed'] = host_processed
  stored_status['mac_to_host'] = mac_to_host
  stored_status['nmap_status'] = status
  return stored_status
  



def generate_status(configs, stored_status):
  """ Generate result file from stored_status """
  host_processed = stored_status['host_processed']
  mac_to_host = stored_status['mac_to_host']
  status = stored_status['nmap_status']

  # Data store format used for nodemap generation
  nodemap = { 'node' : {}, 'link' : {}}

  # XXX: Pushed back till we actually store the MAC in the config files automatically
  #configmac_to_host = dict()
  #for host,config in configs.iteritems():
  #  for iface_key in gformat.get_interface_keys(config):
  #    configmac_to_host[config[iface_key]['mac']] = host

  # List of hosts which has some kind of problem
  for host in configs.keys():
     fqdn = host + ".wLeiden.NET"
     if fqdn in host_processed.keys():
       continue
     config = configs[host]
     print "# Problems in host '%s'" % host
     host_down = True
     for ip in gformat.get_used_ips([config]):
       if not gformat.valid_addr(ip):
         continue
       if status[ip]['state'] == "up":
         host_down = False
       print "## -  ", ip, status[ip]['state']
     if host_down:
       print "## HOST is DOWN!"
       nodemap['node'][fqdn] = gformat.DOWN
     else:
       print "## SNMP problems (not reachable, deamon not running, etc)"
       nodemap['node'][fqdn] = gformat.UNKNOWN

       

  # Correlation mapping
  for fqdn, ip in host_processed.iteritems():
    details = status[ip]
    nodemap['node'][fqdn] = gformat.OK
    print "# Working on %s" % fqdn
    for ip, arpmac in details['arpmac'].iteritems():
      if arpmac in details['mac'].keys():
        # Local MAC address
        continue
      if not mac_to_host.has_key(arpmac):
        print "## [WARN] No parent host for MAC %s (%s) found" % (arpmac, ip)
      else:
        print "## Interlink %s - %s"  % (fqdn, mac_to_host[arpmac])
        nodemap['link'][(fqdn,mac_to_host[arpmac])] = gformat.OK

  stream = file(nodemap_status_file,'w')
  yaml.dump(nodemap, stream, default_flow_style=False)
  print "# Wrote nodemap status to '%s'" % nodemap_status_file


def do_merge(files):
  """ Merge all external statuses in our own nodestatus, using optimistic approch """
  try:
    stream = file(nodemap_status_file,'r')
    status = yaml.load(stream)
  except IOError, e:
    # Data store format used for nodemap generation
    status = { 'node' : {}, 'link' : {}}

  for cfile in files:
    try:
      print "# Merging '%s'" % cfile
      stream = file(cfile,'r')
      new_status = yaml.load(stream)
      for item in ['node', 'link']:
        for key, value in new_status[item].iteritems():
          if not status[item].has_key(key):
            # New items always welcome
            status[item][key] = value
            print "## [%s][%s] is new (%s)" % (item, key, value) 
          elif value < status[item][key]:
            # Better values always welcome
            status[item][key] = value
            print "## [%s][%s] is better (%s)" % (item, key, value) 
    except IOError, e:
      print "## ERROR '%s'" % e

  # Save results back to file
  stream = file(nodemap_status_file,'w')
  yaml.dump(status, stream, default_flow_style=False)

def usage():
    print "Usage: %s <arguments>"
    print "Arguments:"
    print "\tall         = scan all known ips, using cached nmap"
    print "\tnmap-only   = scan all known ips, using nmap only"
    print "\tsnmp-only   = scan all known ips, using snmp only"
    print "\tforce       = scan all known ips, no cache used"
    print "\tforced-snmp = scan all known ips, no snmp cache"
    print "\tstored      = generate status file using stored entries"
    print "\thost <HOST1> [HOST2 ...]  = generate status file using stored entries"
    print "\tmerge <FILE1> [FILE2 ...] = merge status file with other status files"
    sys.exit(0)


def main():
  start_time = datetime.now()
  stored_status_file = '/tmp/stored_status.yaml'
  nmap_result_file = '/tmp/test.xml'

  stored_status = dict()
  nmap_status = dict()
  snmp_status = dict()

  opt_nmap_scan = True
  opt_store_scan = True
  opt_snmp_scan = True
  opt_force_snmp = False
  opt_force_scan = False
  opt_force_range = False
  if len(sys.argv) == 1:
    usage()

  if sys.argv[1] == "all":
    pass
  elif sys.argv[1] == "nmap-only":
    opt_snmp_scan = False
  elif sys.argv[1] == "snmp-only":
    opt_nmap_scan = False
  elif sys.argv[1] == "force":
    opt_force_scan = True
  elif sys.argv[1] == "forced-snmp":
    opt_nmap_scan = False
    opt_force_snmp = True
  elif sys.argv[1] == "host":
    opt_force_range = True
    opt_force_scan = True
  elif sys.argv[1] == "stored":
    opt_snmp_scan = False
    opt_nmap_scan = False
    opt_store_scan = False
  elif sys.argv[1] == "merge":
    do_merge(sys.argv[2:])
    sys.exit(0)
  else:
    usage()

  # By default get all IPs defined in config, else own range
  if not opt_force_range:
    configs = gformat.get_all_configs()
    iplist = gformat.get_used_ips(configs.values())
  else:
    iplist = sys.argv[1:]

  # Load data hints from previous run if exists
  if not opt_force_scan and os.path.exists(stored_status_file) and os.path.getsize(stored_status_file) > 0:
    print "## Loading stored data hints from '%s'" % stored_status_file
    stream = file(stored_status_file,'r')
    stored_status = yaml.load(stream)
  else:
    print "[ERROR] '%s' does not exists" % stored_status_file

  # Do a NMAP discovery
  if opt_nmap_scan:
    if not opt_store_scan:
      nmap_result_file = None
    nmap_status = do_nmap_scan(
                   "-p T:ssh,U:domain,T:80,T:ntp,U:snmp,T:8080 -sU -sT ",
                   iplist,nmap_result_file, opt_force_scan)

  else:
    nmap_status = stored_status['nmap_status']

  # Save the MAC -> HOST mappings, by default as it helps indentifing the
  # 'unknown links'
  mac_to_host = {}
  if stored_status:
     mac_to_host = stored_status['mac_to_host']

  # Do SNMP discovery
  if opt_snmp_scan:
    snmp_status = do_snmp_scan(iplist, nmap_status, stored_status, opt_force_snmp)
  else:
    snmp_status = stored_status

  # Include our saved MAC -> HOST mappings
  mac_to_host.update(snmp_status['mac_to_host'])
  snmp_status['mac_to_host'] = mac_to_host
 
  # Store changed data to disk 
  if opt_store_scan:
    stream = file(stored_status_file,'w')
    yaml.dump(snmp_status, stream, default_flow_style=False)
    print "## Stored data hints to '%s'" % stored_status_file

  # Finally generated status
  generate_status(configs, snmp_status)
  print "# Took %s seconds to complete" % (datetime.now() - start_time).seconds



if __name__ == "__main__":
  main()