#!/usr/bin/env python -O
# Rick van der Zwet, (c) Dec 2006
# <rick@wzoeterwoude.net>

#Simple module to probe the network and draw a picture about it
# Requirements
# -pysnmp 2.x <http://pysnmp.sourceforge.net/>

import sys
import re

# Import PySNMP modules
from pysnmp import asn1, v1, v2c
from pysnmp import role

from common import *

import binascii
import pickle

import getopt

# Initialize defaults
port = 161
retries = 1
timeout = 1
version = '2c'
community = 'public'
loglevel = 1

#Based on code Written by Ilya Etingof <ilya@glas.net>, 2000-2002
def snmp_call(hostname,head_oid):
    result = {}
    
    # Create SNMP manager object
    client = role.manager((hostname, port))
    
    # Pass it a few options
    client.timeout = timeout
    client.retries = retries
    
    # Create a SNMP request&response objects from protocol version
    # specific module.
    try:
        req = eval('v' + version).GETREQUEST()    
        nextReq = eval('v' + version).GETNEXTREQUEST()
        rsp = eval('v' + version).GETRESPONSE()
    
    except (NameError, AttributeError):
        print 'Unsupported SNMP protocol version: %s\n%s' % (version, usage)
        sys.exit(-1)
    
    # Store tables headers
    head_oids = [ head_oid ]
    
    # BER encode initial SNMP Object IDs to query
    encoded_oids = map(asn1.OBJECTID().encode, head_oids)
    
    # Traverse agent MIB
    while 1:
        # Encode SNMP request message and try to send it to SNMP agent
        # and receive a response
        (answer, src) = client.send_and_receive(\
                        req.encode(community=community, encoded_oids=encoded_oids))
    
        # Attempt to decode SNMP response
        rsp.decode(answer)
    
        # Make sure response matches request (request IDs, communities, etc)
        if req != rsp:
            raise 'Unmatched response: %s vs %s' % (str(req), str(rsp))
    
        # Decode BER encoded Object IDs.
        oids = map(lambda x: x[0], map(asn1.OBJECTID().decode, \
                                       rsp['encoded_oids']))
    
        # Decode BER encoded values associated with Object IDs.
        vals = map(lambda x: x[0](), map(asn1.decode, rsp['encoded_vals']))
    
        # Check for remote SNMP agent failure
        if rsp['error_status']:
            # SNMP agent reports 'no such name' when walk is over
            if rsp['error_status'] == 2:
                # Switch over to GETNEXT req on error
                # XXX what if one of multiple vars fails?
                if not (req is nextReq):
                    req = nextReq                
                    continue
                # One of the tables exceeded
                for l in oids, vals, head_oids:
                    del l[rsp['error_index']-1]
            else:
                raise 'SNMP error #' + str(rsp['error_status']) + ' for OID #' \
                      + str(rsp['error_index'])
    
        # Exclude completed OIDs
        while 1:
            for idx in range(len(head_oids)):
                if not asn1.OBJECTID(head_oids[idx]).isaprefix(oids[idx]):
                    # One of the tables exceeded
                    for l in oids, vals, head_oids:
                        del l[idx]
                    break
            else:
                break
    
        if not head_oids:
            return result
            
        # Print out results
        for (oid, val) in map(None, oids, vals):
            if val:
                oid = oid.replace(head_oid,'',1)
                result[oid] = val
    
        # BER encode next SNMP Object IDs to query
        encoded_oids = map(asn1.OBJECTID().encode, oids)
    
        # Update request object
        req['request_id'] = req['request_id'] + 1
    
        # Switch over GETNEXT PDU for if not done
        if not (req is nextReq):
            req = nextReq


#credits to http://hachoir.org/browser/hachoir-core/trunk/hachoir_core/bits.py
#
#print binascii.b2a_hex(value)
#with http://aspn.activestate.com/ASPN/Cookbook/Python/Recipe/496784
def str2hex(value, prefix="", glue=u"", format="%02X"):
    r"""
    Convert binary string in hexadecimal (base 16).

    >>> str2hex("ABC")
    u'414243'
    >>> str2hex("\xF0\xAF", glue=" ")
    u'F0 AF'
    >>> str2hex("ABC", prefix="0x")
    u'0x414243'
    >>> str2hex("ABC", format=r"\x%02X")
    u'\\x41\\x42\\x43'
    """
    if isinstance(glue, str):
        glue = unicode(glue)
    if 0 < len(prefix):
        text = [prefix]
    else:
        text = []
    for character in value:
        text.append(format % ord(character))
    return glue.join(text)

def visit_allowed(key):
    if key[0:7] == '172.16.':
        return True
    else:
        return False




def log(string, level=1):
    if level <= loglevel:
        sys.stdout.write("[%i] %s \n" % (level, string))

def nodedisplay(ip):
    try:
        name = ip2node[ip]
    except:
        name = ''
    return ("%s [%s]" % (ip, name))


class Node:
    def addlocal(self,ip):
        self.local_ip.append(ip)
    def addexternal(self,ip):
        self.exteral_ip.append(ip)
    def linked(self,e_node):
        for l_ip in self.local_ip:
            for e_ip in e_node.exteral_ip:
                if l_ip == e_ip:
                    return True
        return False
    def display(self):
        for l_ip in self.local_ip:
            print "local: %s" % nodedisplay(l_ip)
        for e_ip in self.exteral_ip:
            print "exteral: %s" % nodedisplay(e_ip)

    def __init__(self):
        self.local_ip = []
        self.exteral_ip = []


def ip2name(infile):
    log("Using %s for conversion" % infile,2)
    ip2name= {}
    try:
        input = open(infile, 'r')
    except:
        return ip2name
    for line in  input.readlines():
        (k, l ) =  line.rstrip("\n").split("|")
        ip2node[k] = l
    input.close()
    return ip2name


def draw(infile,outfile,ip2name_infile):
    ip2node = ip2name(ip2name_infile)
    G=AGraph(directed=True)
    G.node_attr['style'] = 'filled'
    G.node_attr['shape'] = 'box'
    G.node_attr['color'] = 'green'
    G.graph_attr['splines']='true'
    G.graph_attr['overlap']='false'
    G.graph_attr['fixedsize']='false'

    file = open(infile, 'rb')
    nodes = pickle.load(file)
    file.close()
    list = nodes.keys()
    for k in list:
        k_node = nodedisplay(k)
        print "Node %s" % nodedisplay(k)
        for l in list:
            l_node = nodedisplay(l)
            if nodes[k].linked(nodes[l]):
                G.add_edge(l_node,k_node)
                print "-%s" % nodedisplay(l)
    
    G.layout()
    G.draw(outfile)
    log("File saved in '%s'" % outfile,2)


def probe(hostname,outfile):
    to_visit = [ hostname ]
    log_l = 3
    visited = { }
    ok_visit = []
    failed_visit = []
    ip2local = {}
    ip2external = {}
    nodes = {}
    while to_visit != []:
        hostname = to_visit.pop(0)
        if visited.has_key(hostname):
            log("%s Already visited, skipping" % hostname,log_l)
        else:
            log("---Processing host '%s', %i left, %i/%i ok/failed visited ---" %\
            (hostname, len(to_visit), len(ok_visit), len(failed_visit)),log_l)
    
            #local ips
            try:
                local_ips = snmp_call(hostname,'.1.3.6.1.2.1.4.20.1.1.')
                #mac ips =
                mac_ips = snmp_call(hostname,'.1.3.6.1.2.1.3.1.1.2.')
            except:
                visited[hostname] = True
                failed_visit.append(hostname)
                log("Error: reading data",log_l)
                continue
            
            nodes[hostname] = Node()
            ok_visit.append(hostname)
            for key in local_ips.keys():
                nodes[hostname].addlocal(key)
                visited[key] = True
            for key,value in mac_ips.iteritems():
                key = re.sub('^[0-9]\.1\.','',key)
                mac = str2hex(value,'',':')
                log('%s - %16s : ' % (mac, key),log_l)
                string = ''
                if local_ips.has_key(key):
                    string += "internal"
                else:
                    nodes[hostname].addexternal(key)
                    string += "external "
                    if visited.has_key(key):
                        string += "visited"
                    elif key in to_visit:
                        string += "already planned"
                    else:
                        if visit_allowed(key):
                            string += "added to the list"
                            to_visit.append(key)
                        else:
                            string += "not allowed to visit"
                log(string,log_l)
    
    o_file = open(outfile, 'wb')
    pickle.dump(nodes, o_file)
    o_file.close()
    log("Wrote output to '%s'" % outfile,log_l)

def usage():
    print '''usage
%(prog)s --probe <ip-number> [--outfile=<file>]
%(prog)s --help

''' % { 'prog': sys.argv[0] }


def main():
    global loglevel
    try:
        opts, args = getopt.getopt(sys.argv[1:], 'hp::o::v::',\
        ["help", "probe", "outfile=", "infile=","verbose="])
    except getopt.GetoptError:
        usage()
        sys.exit(128)

    action = False
    hostname = '172.27.129.65'
    infile = False
    outfile = False

    for o, a in opts:
        log("Command '%s' = '%s'" % (o, a),2)
        if o in ("-v", "--verbose"):
            if a.isdigit():
                loglevel = a
            else:
                loglevel = a.count('v')
            log("Loglevel now at %s" % loglevel,4)
        if o in ("-h", "--help"):
            usage()
            sys.exit()
        if o in ("-o", "--outfile"):
            outfile = a
        if o in ("-p", "--probe"):
            action = 'probe'
            print a
            if a:
                hostname = a

    if action == 'probe':
        if not outfile:
            outfile = 'data.pkl'
        probe(hostname,outfile)
    else:
        usage()
        sys.exit(128)

if __name__ == "__main__":
    main()