#!/usr/bin/python
# A python replacement for arptables-restore

import os
import re
import subprocess
import sys

arptables = "{}/arptables".format( os.path.dirname( os.path.realpath( __file__ ) ) )

def runCmd( cmd ):
   process = subprocess.Popen( cmd.split(), stdout=subprocess.PIPE )
   output = process.communicate()[ 0 ]
   return process.returncode, output

def runArptablesCmd( cmd ):
   returnCode, output = runCmd( "{} {}".format( arptables, cmd ) )

   if returnCode != 0:
      print "ERROR:", output
      sys.exit( 1 )

   return returnCode, output

def clearArptables():
   # set policies to default ACCEPT
   runArptablesCmd( "-P INPUT ACCEPT" )
   runArptablesCmd( "-P FORWARD ACCEPT" )
   runArptablesCmd( "-P OUTPUT ACCEPT" )

   # remove all rules in chains
   runArptablesCmd( "-F" )

   # remove all non-default chains
   _, rules = runArptablesCmd( "-L" )
   for rule in rules.splitlines():
      match = re.search( "Chain\s(.*?)\s\(.*references\)", rule )
      if match:
         runArptablesCmd( "-X {}".format( match.group( 1 ) ) )

if not os.access( arptables, os.X_OK ):
   print "ERROR: {} isn't executable".format( arptables )
   sys.exit( 1 )

clearArptables()

lineNum = 1
table = ""
for rule in sys.stdin:
   rule = rule.strip()

   if len( rule ) == 0:
      continue

   # skip comments
   if re.search( "^#", rule ):
      continue

   # setting the current table
   tableMatch = re.search( "^\*(.*)", rule )
   if tableMatch:
      table = tableMatch.group( 1 )
      continue

   # setting the chain default policy
   directiveMatch = re.search( "^\:(.*?)\s(.*)", rule )
   if directiveMatch:
      if directiveMatch.group( 2 ) == "-":
         runArptablesCmd( "-t {} -N {}".format( table, directiveMatch.group( 1 ) ) )
      else:
         runArptablesCmd( "-t {} -P {} {}".format( table,
                                                   directiveMatch.group( 1 ),
                                                   directiveMatch.group( 2 ) ) )
      continue

   # all other rules
   runArptablesCmd( "-t {} {}".format( table, rule ) )
