#!/usr/bin/python

import sys,os,string

def MAC(mac):
  mac = mac.split(':')
  mac = map(lambda x: x.zfill(2), mac)
  mac = string.join(mac, ':')
  return mac

def clean_chain(iface, mac, chain):
  p = os.popen('ebtables -L %s | grep %s' % (chain, mac))
  for x in p.readlines():
    os.system('ebtables -D %s -i %s -s %s -j ACCEPT' % (chain, iface, mac))
  p.close()
 
def clean_mac(iface, mac, chains):
  for chain in chains:
    clean_chain(iface, mac, chain)

def if_list(iface):
  ports=os.popen('brctl showmacs %s | grep yes' % iface).readlines()
  ifaces = {}
  for port in ports:
    port=port.split()[0]
    ifa = os.popen('brctl showstp %s | egrep -o "^eth[0-9]+ \(%s\)"' % (iface, port)).readline()
    ifa = ifa.split()[0]
    ifaces[port]= ifa
  return ifaces

def get_mac_port(iface, mac):
  mac = MAC(mac)
  port=os.popen('brctl showmacs %s | grep %s' % (iface, mac)).readline()
  port=port.split()[0]
  iface = os.popen('brctl showstp %s | egrep -o "^eth[0-9]+ \(%s\)"' % (iface, port)).readline()
  iface = iface.split()[0]
  return iface

try:
  op = sys.argv[1]
  if not op in ['-a','-p','-b','-i']:
    raise ValueError('Invalid option %s' % op)
  if op in ['-p','-b']:
    mac = sys.argv[2]
    iface = sys.argv[3]
  else:
    iface = sys.argv[2]
  #print if_list(iface)
except:
  print 'Use: %s {-p|-b} MAC IFACE | {-i | -a} IFACE' % sys.argv[0]
  sys.exit(1)

open('/tmp/log','a').write('%s\n' % sys.argv)

fwd_chain = 'fwd_%s' % iface
inp_chain = 'inp_%s' % iface

if op == '-i':
  os.system('ebtables -N %s\nebtables -N %s\n' % (inp_chain, fwd_chain))
  l = os.popen('ebtables -L %s | grep 0x888e' % inp_chain).readlines()
  for x in l:
    os.system('ebtables -D %s -p 0x888e -j ACCEPT > /dev/null 2>&1' % inp_chain)
  os.system('ebtables -P %s DROP; ebtables -P %s DROP; ebtables -A %s -p 0x888e -j ACCEPT' % (inp_chain, fwd_chain, inp_chain))
  os.system('ebtables -A INPUT -i %s -j %s;ebtables -A FORWARD -i %s -j %s' % (iface, inp_chain, iface, fwd_chain))
elif op == '-a':
  os.system('ebtables -D INPUT -i %s -j %s;ebtables -D FORWARD -i %s -j %s' % (iface, inp_chain, iface, fwd_chain))
  os.system('ebtables -X inp_%s\nebtables -X fwd_%s\n' % (iface, iface))
else:
  iface = get_mac_port(iface, mac)

  clean_mac(iface, mac, (fwd_chain, inp_chain))
  if op == '-p':
    os.system('ebtables -A %s -i %s -s %s -j ACCEPT; ebtables -A %s -i %s -s %s -j ACCEPT' % (inp_chain, iface, mac, fwd_chain, iface, mac))

sys.exit(0)
