#!/usr/bin/python
# Copyright (C) 2006-2007 XenSource Ltd.
# Copyright (C) 2008-2009 Citrix Ltd.
#
# This program is free software; you can redistribute it and/or modify 
# it under the terms of the GNU Lesser General Public License as published 
# by the Free Software Foundation; version 2.1 only.
#
import sys
sys.path.insert(0, "/opt/xensource/sm")
sys.path.insert(0, "/opt/xensource/sm/snapwatchd")
import util
import xslib
import random

def help():
    print "Usage: blockunblockiscsipaths <block/unblock> <noOfPaths> <IP1>,<IP2>,..."
    sys.exit(-1)
    
def blockIP(ip):
    try:
	cmd = ['iptables', '-A', 'OUTPUT', '-d', ip, '-j', 'DROP']
        util.pread(cmd)
	cmd = ['iptables', '-A', 'INPUT', '-s', ip, '-j', 'DROP']
        util.pread(cmd)
    except Exception, e:
        util.SMlog("There was an exception in blocking ip: %s" % ip)

def unblockIP(ip):
    try:
	cmd = ['iptables', '-D', 'OUTPUT', '-d', ip, '-j', 'DROP']
        util.pread(cmd)
	cmd = ['iptables', '-D', 'INPUT', '-s', ip, '-j', 'DROP']
        util.pread(cmd)
    except Exception, e:
        util.SMlog("There was an exception in blocking ip: %s" % ip)
	
# Test Cmdline args
if len(sys.argv) != 4:
    help()

op = sys.argv[1]
no = sys.argv[2]
ipList = sys.argv[3].split(',')
    
if op == 'block':
    newList = random.sample(ipList, int(no))
else:
    newList = ipList
    
paths = ''
for ip in newList:    
    if op == 'block':
	paths += ip
	blockIP(ip)
    elif op == 'unblock':
	unblockIP(ip)
    else:
	continue
    paths += ','

paths = paths.strip(',')

xs_handle = xslib.xs_daemon_open()
xslib.setval(xs_handle, '/xencert/block-unblock-over', '1')
xslib.xs_daemon_close(xs_handle)
if op == 'block':
    sys.stdout.write(paths)
sys.exit(0)