Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Python-iptables how to optimize code

I am just learning python and have written some code to set iptables using the python-iptables library. The problem I'm running in to is that I had to rewrite a lot of the same lines of code over and over. I understand functions somewhat but not OOP. I'm thinking there is a better OOP way of writing this code but I can't get my head around it. Any pointers will be greatly appreciated. The code is below.

import iptc

def dropAllInbound():
    chain = iptc.Chain(iptc.Table(iptc.Table.FILTER), 'INPUT')
    rule = iptc.Rule()
    rule.in_interface = 'eth+'
    rule.target = iptc.Target(rule, 'DROP')
    chain.insert_rule(rule)

def allowLoopback():
    chain = iptc.Chain(iptc.Table(iptc.Table.FILTER), 'INPUT')
    rule = iptc.Rule()
    rule.in_interface = 'lo'
    rule.target = iptc.Target(rule, 'ACCEPT')
    chain.insert_rule(rule)

def allowEstablishedInbound():
    chain = iptc.Chain(iptc.Table(iptc.Table.FILTER), 'INPUT')
    rule = iptc.Rule()
    match = rule.create_match('state')
    match.state = 'RELATED,ESTABLISHED'
    rule.target = iptc.Target(rule, 'ACCEPT')
    chain.insert_rule(rule)

def allowHTTP():
    chain = iptc.Chain(iptc.Table(iptc.Table.FILTER), 'INPUT')
    rule = iptc.Rule()
    rule.in_interface = 'eth+'
    rule.protocol = 'tcp'
    match = rule.create_match('tcp')
    match.dport = '80'
    rule.target = iptc.Target(rule, 'ACCEPT')
    chain.insert_rule(rule)

def allowHTTPS():
    chain = iptc.Chain(iptc.Table(iptc.Table.FILTER), 'INPUT')
    rule = iptc.Rule()
    rule.in_interface = 'eth+'
    rule.protocol = 'tcp'
    match = rule.create_match('tcp')
    match.dport = '443'
    rule.target = iptc.Target(rule, 'ACCEPT')
    chain.insert_rule(rule)

def allowSSH():
    chain = iptc.Chain(iptc.Table(iptc.Table.FILTER), 'INPUT')
    rule = iptc.Rule()
    rule.in_interface = 'eth+'
    rule.protocol = 'tcp'
    match = rule.create_match('tcp')
    match.dport = '22'
    rule.target = iptc.Target(rule, 'ACCEPT')
    chain.insert_rule(rule)

def allowEstablishedOutbound():
    chain = iptc.Chain(iptc.Table(iptc.Table.FILTER), 'OUTPUT')
    rule = iptc.Rule()
    match = rule.create_match('state')
    match.state = 'RELATED,ESTABLISHED'
    rule.target = iptc.Target(rule, 'ACCEPT')
    chain.insert_rule(rule)

def dropAllOutbound():
    chain = iptc.Chain(iptc.Table(iptc.Table.FILTER), 'OUTPUT')
    rule = iptc.Rule()
    rule.in_interface = 'eth+'
    rule.target = iptc.Target(rule, 'DROP')
    chain.insert_rule(rule)

def defaultAction():
    dropAllOutbound()
    dropAllInbound()
    allowLoopback()
    allowEstablishedInbound()
    allowEstablishedOutbound()

def getInput():
        print 'Default action (1) is most secure '
        print 'Default  -  1'
        print 'HTTP     -  2'
        print 'HTTPS    -  3'
        print 'SSH      -  4'
        print 'Exit     -  5'
        choices = raw_input('Enter choices (comma Separated) ').split(',')
        for action in choices:
            if action == "1":
                defaultAction()
                break
            if action == "2":
                allowHTTP()
                break
            if action == "3":
                allowHTTPS()
                break
            if action == "4":
                allowSSH()
                break
            else:
                break
getInput()

Notice how all the rules have similar lines of code. Is there a way to create a rule generator object or something like that to minimize rewriting that code?

I added the following function so and call it every time the script is run so that the rules get flushed.

def startClean():
    chainIn = iptc.Chain(iptc.Table(iptc.Table.FILTER), 'INPUT')
    chainIn.flush()
    chainOut = iptc.Chain(iptc.Table(iptc.Table.FILTER), 'OUTPUT')
    chainOut.flush()
like image 892
h33th3n Avatar asked Dec 26 '13 01:12

h33th3n


1 Answers

OOP is for maintaining the state of something. OOP is for some object that has both properties and methods to manipulate those properties.

class Chair(object):

    MAX_WEIGHT = 300

    def __init__(self):
        super().__init__()

        self.weight = 5
        self.currentWeight = self.weight
        self.holding = None
        self.broken = False

    def hold(self, item):
        self.holding = item
        self.currentWeight = self.weight + item.weight
        self.checkWeight()

    def checkWeight(self):
        if self.holding.weight > self.MAX_WEIGHT:
            self.broken = True
            ...

Your code seems fine; rewriting the code just for OOP may be more work than it's worth. If you really want to use OOP you may wish to do something like below.

class Table(object):
    def __init__(self):
        self.chain = None
        self.rule = None
        self.match = None

    def setInput(self):
        self.chain = iptc.Chain(iptc.Table(iptc.Table.FILTER), 'INPUT')

    def setOutput(self):
        self.chain = iptc.Chain(iptc.Table(iptc.Table.FILTER), 'OUTPUT')

    ...

table = Table()
table.setInput()
...
like image 112
justengel Avatar answered Sep 28 '22 20:09

justengel