Syeds' spam.py, updated

spam.py

#!/usr/local/bin/python

# Script used to check /var/log/iptables.log and penalize mail activity outside an allowed amount.
# you need to edit /etc/rsyslog.conf and add: "kern.warning     /var/log/iptables.log"
# only supported by CentOS 6 and CentOS 7 currently.

from datetime import datetime
from datetime import timedelta
from collections import defaultdict
import subprocess
import os
from shutil import copyfile
import sys

with open('/etc/dhcpd.conf') as gg:
    dhcp = gg.read()

# Let's check if our iptables chain exists and is added to FORWARD
ipt = subprocess.check_output(['/sbin/iptables', '-S', 'FORWARD'])
if "MAILBLOCKED" not in ipt:
        print "Creating iptables chain and adding it to FORWARD - since we need it"
        subprocess.check_output(['/sbin/iptables', '-I', 'FORWARD', '3', '-p', 'tcp', '-m', 'tcp', '--dport', '25', '-j', 'LOG', '--log-prefix', '"Mail Packet: "'])
        subprocess.check_output(['/sbin/iptables', '-N', 'MAILBLOCKED'])
        subprocess.check_output(['/sbin/iptables', '-I', 'FORWARD', '4', '-j', 'MAILBLOCKED'])

# This is essentially a patch to already running scripts since the system will live push these changes
if not os.path.isfile("/etc/rsyslog.d/iptables.conf"):
        # Create the iptables conf file we need.
        with open('/etc/rsyslog.d/iptables.conf','w') as f:
                f.write(':msg, contains, "Mail Packet: " -/var/log/iptables.log\n& ~')

        # Clear the full log files now.
        open('/var/log/messages', 'w').close()
        open('/var/log/iptables.log', 'w').close()

        # Backup the rsyslog file and remove the previous change we made.
        copyfile('/etc/rsyslog.conf', '/etc/rsyslog.conf.bak') # Let's back up the file first.

        # Remove the previous line
        with open("/etc/rsyslog.conf","r") as input:
                input = input.readlines()
                with open("/etc/rsyslog.conf","wb") as output:
                        for line in input:
                                if not line.startswith("kern.warning"):
                                        output.write(line)
        # Restart rsyslogd service
        subprocess.check_output(['service', 'rsyslog', 'restart'])

# Current time so that we don't miss a few packets in the beginning by doing it within the loop
now = datetime.now()

# create the default dict.
spamdict = defaultdict(int)
spamkvms = {}

# Remove the file if it's too big before attempting to process it.
if os.path.getsize('/var/log/iptables.log') > 9990000:
        print "File is too big - clearing the log."
        with open('/var/log/iptables.log', 'w'): pass
        # Restart rsyslogd service - SSH
        subprocess.check_output(['service', 'rsyslog', 'restart'])
        sys.exit() # Exit script
        
# Get the mail packet log and then do the rest.
with open('/var/log/iptables.log') as f:
    iplog = f.readlines()

# Let's do some parsing
for line in iplog:
        if "Mail Packet" not in line:
                continue
        line = line.split()
        # Doing a bunch of date shit so we get the lines we need.
        date = datetime.strptime(str(now.year) + ' ' + ' '.join(line[0:3]), "%Y %b %d %H:%M:%S")
        # Check past hour for any spammers.
        if date > now - timedelta(minutes=60):
                try: # wrap it in a try.
                        # Get source IP and length of packet.
                        src, length = line[11], line[13]

                        # verify the packet, else we continue.
                        if not src.startswith("SRC") and not length.startswith("LEN="):
                                continue # not a line we can handle.

                        src, length = src.split("SRC=")[1], int(length.split("LEN=")[1])
                        spamdict[src] += length

                        # Looks like we're seeing DST IPs be the right one, we might as well add these - this should only be outbound anyways, it's the FORWARD rule chain.
                        dst = line[12].split("DST=")[1]
                        spamdict[dst] += length

                        # note down the iface into a dict - checked against dhcpd conf later.
                        spamkvms[src] = line[9].split('PHYSIN=')[1]
                        spamkvms[dst] = line[10].split('PHYSOUT=')[1]
                except:
                        continue

# We've parsed the file now. Let's do some checking.
for k,v in spamdict.items():
        # 3500000 bytes per hour / 100 35KB emails, based on this: https://securelist.com/spam-and-phishing-in-q1-2017/78221/#spam-email-size
        if v > 1000000:
                ipr = subprocess.check_output(['/sbin/iptables', '-S', 'MAILBLOCKED'])
                # first we check if this is in the dhcpd.conf file - so we avoid adding needless rules. also we check current rules.
                try:
                        if spamkvms[k] in dhcp and not k in ipr:
                                print "Added " + k + " to iptables. We're now blocking port 25 on this."
                                subprocess.check_output(['/sbin/iptables', '-A', 'MAILBLOCKED', '-m', 'tcp', '-p', 'tcp', '-s', k + '/32', '--dport', '25', '-j', 'REJECT'])
                except:
                        print "Failed looking up this server. Continuing"
                        continue
# print the dict
print spamdict