#!/usr/bin/python
# fence_snmp.py: fabric fencing for RHCS based on setting a network interface
# to admin down.  Will be used for iSCSI connections.
#
# Written by Ross Vandegrift <ross@kallisti.us>
# Released into the public domain

import sys, getopt, random, socket
from pysnmp import role, v2c, asn1

ifAdminStatus = ".1.3.6.1.2.1.2.2.1.7."
up = 1
down = 2
testing = 3

def usage():
    line = "\t%s\t%s"
    print ""
    print "This script fences a node by sending a command via SNMP to set"
    print "ifAdminStatus to down.  It's designed to kill an iSCSI node's"
    print "access to the shared storage.  It supports SNMP v1 and v2c."
    print ""
    print "Usage: fence_snmp [options]"
    print line % ("-h", "\tPrint usage")
    print line % ("-V", "\tRun verbosely")
    print line % ("-c [private]", "Write community string to use")
    print line % ("-v [1|2c]", "Use SNMP version 1 or 2c")
    print line % ("-a [hostname]", "IP/hostname of SNMP agent")
    print line % ("-i [index]", "ifIndex entry of the port ")
    print line % ("-o [action]", "One of down, up, or status")


def vprint(v, s):
    if v:
        print s


def parseargs():
    try:
        opt, arg = getopt.getopt (sys.argv[1:], 'hVc:v:a:i:o:')
    except getopt.GetoptError, e:
        print str (e)
        usage ()
        sys.exit (-1)

    comm = ver = host = index = action = verbose = None

    for o, a in opt:
        if o == '-h':
            usage ()
            sys.exit (-1)
        if o == '-V':
            verbose = True
        if o == '-c':
            comm = a
        if o == '-v':
            ver = a
            if ver not in ('1', '2c'):
                print "version must be one of 1 or 2c"
                usage ()
                sys.exit (-1)
        if o == '-a':
            host = a
        if o == '-i':
            try:
                index = int(a)
            except:
                print "ifIndex must be an integer"
                usage ()
                sys.exit (-1)
        if o == '-o':
            action = a
            if action not in ('up', 'down', 'status'):
                print "action msut be one of up, down, or status"
                usage ()
                sys.exit (-1)

    if comm == None or ver == None or host == None or index == None \
            or action == None:
        print "All args are madatory!"
        usage ()
        sys.exit (-1)

    return (comm, ver, host, index, action, verbose)


def snmpget (host, comm, oid):
    req = v2c.GETREQUEST ()
    encoded_oids = map (asn1.OBJECTID().encode, [oid,])
    req['community'] = comm
    tr = role.manager ((host, 161))
    rsp = v2c.RESPONSE ()
    (rawrsp, src) = tr.send_and_receive (req.encode (encoded_oids=encoded_oids))
    rsp.decode (rawrsp)
    oids = map (lambda x: x[0], map(asn1.OBJECTID ().decode, rsp['encoded_oids']))
    vals = map (lambda x: x[0] (), map(asn1.decode, rsp['encoded_vals']))
    return vals[0]


def snmpset (host, comm, oid, type, value):
    req = v2c.SETREQUEST (request_id=random.randint (1,2**16-1))
    req['community'] = comm
    tr = role.manager ((host, 161))
    rsp = v2c.RESPONSE ()
    encoded_oids = map (asn1.OBJECTID ().encode, [oid,])
    encoded_vals = []
    encoded_vals.append (eval ('asn1.' + type + '()').encode (value))
    (rawrsp, src) = tr.send_and_receive (req.encode (encoded_oids=encoded_oids, encoded_vals=encoded_vals))
    rsp.decode(rawrsp)
    if rsp['error_status']:
        raise IOError('SNMP error while setting ifAdminStatus: check community string')
    oids = map (lambda x: x[0], map (asn1.OBJECTID().decode, rsp['encoded_oids']))
    vals = map (lambda x: x[0] (), map (asn1.decode, rsp['encoded_vals']))
    if vals[0] == value:
        return vals[0]
    else:
        raise IOError('SNMP error while setting ifAdminStatus: different value returned')


def main():
    (comm, ver, host, index, action, verbose) = parseargs ()

    try:
        switch = socket.gethostbyname (host)
    except socket.gaierror, err:
        vprint (verbose, "fence_snmp: %s" % str (err[1]))

    if action == 'up':
        r = snmpset (switch, comm, ifAdminStatus + str (index), 'INTEGER', up)
        vprint (verbose, "Unfenced")
        sys.exit (0)
    elif action == 'down':
        r = snmpset (switch, comm, ifAdminStatus + str (index), 'INTEGER', down)
        vprint (verbose, "Fenced")
        sys.exit (0)
    else: # assume status
        r = int (snmpget (switch, comm, ifAdminStatus + str (index)))
        if r == up:
            vprint (verbose, "up(1)")
            sys.exit (1)
        elif r == down:
            vprint (verbose, "down(2)")
            sys.exit (2)
        elif r == testing:
            vprint (verbose, "testing(3)")
            sys.exit (3)
        else:
            vprint (verbose, "unknown: %s", r)
            sys.exit (-1)

if __name__ == "__main__":
    main()