firewall_test.py

import sys, os, stat

zones_hosts = {}
open_zones = {}
servers = {}
clients = {}

zones_hosts_file = sys.argv[1]
rules_file = sys.argv[2]
dir_output = "deploy/"

with open(zones_hosts_file) as f:
  for line in f:
    fields = line.rstrip('\n').split(',')
    zone = fields[0]
    hostname = fields[1]
    ip = fields[2]
    if zone not in zones_hosts:
      zones_hosts[zone] = []
    zones_hosts[zone].append([hostname,ip])
    #sys.stdout.write('line is: {0} {1} {2}\n'.format(fields[0], fields[1], fields[2]))

linenumber=0
with open(rules_file) as f:
  for line in f:
    linenumber += 1
    fields = line.rstrip('\n').split(',')
    from_zone = fields[0]
    to_zone = fields[1]
    port = fields[2]
    if from_zone not in zones_hosts:
       print "error on line " , linenumber , ": from " , from_zone , " not found in " + zones_hosts_file
    if to_zone not in zones_hosts:
       print "error on line " , linenumber , ": to " , to_zone , "not found in " + zones_hosts_file
    for from_host in zones_hosts[from_zone]:
       if from_host[0] not in clients:
         clients[from_host[0]] = {}
       for to_host in zones_hosts[to_zone]:
         if to_host[0] not in servers:
           servers[to_host[0]] = {}
         from_to_port = from_host[1] + '|' + to_host[1] + '|' + port
         to_from_port = to_host[1] + '|' + from_host[1] + '|' + port
         servers[to_host[0]][from_to_port] = 1
         #print "rule - from host: ", from_host[0] , " to host " , to_host[0], " on port ", port
         clients[from_host[0]][to_from_port] = 1

for host in servers:
  print "\n\n*****************"
  print "generating server script for HOST: " , host
  print "******************"

  filename = dir_output + host + "_firewall_test_server_stop.sh"
  if os.path.isfile(filename):
    os.remove(filename)
  file = open(filename, 'w')
  file.write("ps -ef | grep firewall_test | grep -v grep | grep -v sh | awk '{print $2}' | xargs kill -9 \n")
  os.chmod(filename, stat.S_IRWXU)
  file.close

  filename = dir_output + host + "_firewall_test_server_start.sh"
  if os.path.isfile(filename):
    os.remove(filename)
  file = open(filename, 'w')
  file.write("python -u " + host + "_firewall_test_server_start.py > " + host + "_firewall_test_server.log &\n")
  os.chmod(filename, stat.S_IRWXU)
  file.close

  filename = dir_output + host + "_firewall_test_server_start.py"
  if os.path.isfile(filename):
    os.remove(filename)
  file = open(filename, 'w')
  file.write("from firewall_test_include import *\n")
  ports = []
  for server in servers[host]:
     in_fields = server.split('|')
     from_host = in_fields[0]
     to_host = in_fields[1]
     port = in_fields[2]
     if port not in ports:
       ports.append(port)
  file.write("ports=" + str(ports) + "\n")
  file.write("start_server(ports)\n")
  file.close()


for host in clients:
  print "\n\n*****************"
  print "generating client script for HOST: " , host
  print "******************"

  filename = dir_output + host + "_firewall_test_client.sh"
  if os.path.isfile(filename):
    os.remove(filename)
  file = open(filename, 'w')
  file.write("python -u " + host + "_firewall_test_client.py > " + host + "_firewall_test_client.log &\n")
  os.chmod(filename, stat.S_IRWXU)
  file.close

  filename = dir_output + host + "_firewall_test_client.py"
  if os.path.isfile(filename):
    os.remove(filename)
  file = open(filename, 'w')
  file.write("from firewall_test_include import *\n")
  ports = []
  for client in clients[host]:
     in_fields = client.split('|')
     to_host = in_fields[0]
     from_host = in_fields[1]
     port = in_fields[2]
     file.write("status,msg = connect('" + to_host + "', " + port  + ")\n")
     file.write("print ('CLIENT STATUS:host=" + to_host + "|port=" + port + "|status=' + str(status) + '|message:' + msg)\n")
  file.write("print 'end of tests for host'\n")
  file.close()

firewall_test_include.py

import socket, sys, select, threading

stop = False

def listen(port, stop):
  s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  success = 0
  error = ""
  try:
    s.bind(('', port))
  except socket.error as msg:
    error = 'SERVER ERROR:hostname=' + socket.gethostname() + '|port=' + str(port) + '|error=' +  str(msg)
    print error
  s.listen(10)
  success = 1
  while not stop():
    ready = select.select([s], [], [], 1.0)
    if ready[0]:
      conn, addr = s.accept()
      #print 'SERVER: Connection open - from:  ' + addr[0] + ':' + str(addr[1])
      data = conn.recv(1024)
      conn.send(data)
      conn.close()
  s.close()
  return success,error

def connect(host, port):
  s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  success = 0
  msgout  = ""
  try:
    s.settimeout(2)
    s.connect((host, port))
    s.sendall(b'ping')
    data = s.recv(1024)
    success = 1
    msgout = data
    #print 'CLIENT: Received', repr(data)
  except socket.error as msg:
    msgout = str(msg)
    #print 'CLIENT: open failed: hostname = ' + host + ', port = ' + str(port) + ', error = ' +  str(msg)
  s.close()
  return success,msgout


def start_server(ports):
  threads = []
  for port in ports:
    thr = threading.Thread(target=listen, args=([int(port), lambda: stop]), kwargs={})
    thr.start()
    threads.append(thr)

##start server
#server_ports = [8888,8889,8888]
#start_server(server_ports)

#start server - option2
#stop = False
#threads = []
#for port in server_ports:
#  thr = threading.Thread(target=listen, args=([port, lambda: stop]), kwargs={})
#  thr.start()
#  threads.append(thr)

#clients
#connect('localhost', 8888)
#connect('localhost', 22)
#connect('localhost', 777)

#stop = True
#for thr in threads:
#  thr.join

#print "out"

usage.sh

#cat zones_hosts.txt
#ambari,ambari,10.0.1.1
#hadoop,dn1,10.0.0.1
#hadoop,dn2,10.0.0.2
#hadoop,dn3,10.0.0.3

#cat rules.txt
#ambari,hadoop,8670
#hadoop,ambari,8440
#hadoop,ambari,8441
#hadoop,ambari,8440

#cat copy_firewall.sh
#for host in `cat hosts.txt | paste -s -d' '`
#   do
#     ssh ${host} rm -rf /tmp/firewall_test/
#     ssh ${host} mkdir /tmp/firewall_test
#     scp /tmp/firewall_test/* ${host}:/tmp/firewall_test/
#     scp /tmp/firewall_test/deploy/* ${host}:/tmp/firewall_test/
#   done

python firewall_test.py zones_hosts.txt rules.txt
./copy_firewall.sh
clush -a "cd /tmp/firewall_test/;/tmp/firewall_test/\`hostname -f\`_firewall_test_server_start.sh"
clush -a "cd /tmp/firewall_test/;tail -1 /tmp/firewall_test/\`hostname -f\`_firewall_test_client.log"
clush -a "cd /tmp/firewall_test/;grep -H 'status=0' /tmp/firewall_test/\`hostname -f\`_firewall_test_client.log" | sort


clush -a "cd /tmp/firewall_test;/tmp/firewall_test/\`hostname -f\`_firewall_test_server_stop.sh"