#!/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']

#
# 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):
  """ 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 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()):
  """ SNMP scanning, based on results fould in NMAP scan"""
  mac_to_host = dict()
  host_processed = dict()

  #
  # Gather SNMP data from hosts
  for host, scan in status.iteritems():
    if scan['state'] != "up":
      continue
    
    # Filter set? use it
    if iphosts and not host in iphosts:
      print "## IP '%s' not in specified filter" % host
      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 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(':')])
      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
    host_processed[hostname] = status[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']

  # List of hosts which has some kind of problem
  for host in list(set(configs.keys()) - set(host_processed.keys())):
     print "# Problem in host '%s'" % host

  # Correlation mapping
  for host, details in host_processed.iteritems():
    print "# Working on %s" % host
    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"  % (host, mac_to_host[arpmac])


def usage():
    print "Usage: %s <all|force|stored|host HOST1 [HOST2 ...]>" % sys.argv[0]
    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
  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)
  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)
  else:
    nmap_status = stored_status['nmap_status']

  # Do SNMP discovery
  if opt_snmp_scan:
    snmp_status = do_snmp_scan(iplist, nmap_status, stored_status)
  else:
    snmp_status = stored_status
 
  # 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()