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

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

# translate arptables rules to arptables-restore syntax
def processTable( table ):
   currentChain = ""
   chains = []
   customChains = []
   rules = []

   for line in table.split( "\n" ):
      line = line.strip()

      if len( line ) == 0:
         continue

      # add default chains
      chainMatch = re.search( "Chain\s(.*?)\s\(policy\s(.*?)\)", line )
      if chainMatch:
         chain, policy = chainMatch.group( 1 ), chainMatch.group( 2 )
         chains.append( ":{} {}".format( chain, policy ) )
         currentChain = chain
         continue

      # add custom chains
      customChainMatch = re.search( "Chain\s(.*?)\s\(", line )
      if customChainMatch:
         chain = customChainMatch.group( 1 )
         customChains.append( ":{} -".format( chain ) )
         currentChain = chain
         continue

      # add all other rules
      rules.append( "-A {} {}".format( currentChain, line ) )

   return "\n".join( [ "*filter" ] + chains + customChains + rules )

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

_, output = runArptablesCmd( "-t filter -L -n" )
print processTable( output )
