More Pythoning

#!/usr/bin/python -tt
# Copyright 2014 John Lewis
# This is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 2 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.

import sys
import re
import mysql.connector
# Iptables module.
import iptc

def build_blocked_list():
 table = iptc.Table(iptc.Table.FILTER)
 rule = iptc.Rule()
 chain = iptc.Chain(iptc.Table(iptc.Table.FILTER), 'INPUT')
 blocked_ips = []
 for chain in table.chains:
 for rule in chain.rules:
 blocked_ips.append(re.sub('/255.255.255.255', '', rule.src))
 return blocked_ips

def iptables_block(ip):
 table = iptc.Table(iptc.Table.FILTER)
 rule = iptc.Rule()
 chain = iptc.Chain(iptc.Table(iptc.Table.FILTER), 'INPUT')
 rule.in_interface = 'eth0'
 rule.src = ip
 t = rule.create_target('DROP')
 chain.insert_rule(rule)

def check(db, dbusername, dbpassword, logfile):
 dbconnection = mysql.connector.connect(user=dbusername, password=dbpassword,
 database=db)
 dbcursor = dbconnection.cursor()
 ipblocks_query = ('select ipb_address from ipblocks where ipb_address=%s')
 user_query = ('select user_name from user where user_name=%s')

blocked_ips = build_blocked_list()

with open(logfile, 'r') as f:
 f = f.readlines()

for line in f:
 select = re.search('(^\d\d\d\.\d\d\d\.\d\d\d\.\d\d\d).+User:(\w+)', line)
 if select:
 dbcursor.execute(ipblocks_query, (select.group(2),))
 row = dbcursor.fetchone()
 # If user in ipblocks table, block IP if not already.
 if row is not None and select.group(1) not in blocked_ips:
 blocked_ips.append(select.group(1))
 iptables_block(select.group(1))
 else:
 dbcursor.execute(user_query, (select.group(2),))
 row = dbcursor.fetchone()
 # If user *not* in user table, block IP if not already.
 if row is None and select.group(1) not in blocked_ips:
 blocked_ips.append(select.group(1))
 iptables_block(select.group(1))

def main():
 args = sys.argv[1:]
 if not args:
 print 'usage: db dbusername dbpassword logfile'
 sys.exit(1)

if len(args) < 4:
 print 'error: you must specify a db, dbusername, dbpassword, AND logfile'
 sys.exit(1)

db = args[0]
 dbusername = args[1]
 dbpassword = args[2]
 logfile = args[3]

check(db, dbusername, dbpassword, logfile)

if __name__ == '__main__':
 main()

 

Leave a Reply