#!/usr/bin/python
#
# My solution to Email worms
#
#

import StringIO,os,re,select,time

class IPTables:
	def __init__(self):
		self._tick=0
		self.table={}
	def iptables_add(self,ip):
		c1='iptables -A INPUT -m state --state NEW -s "%s" -p tcp --dport 25 -j LOG  --log-prefix "dynablocked "'%ip
		c2='iptables -A INPUT -m state --state NEW -s "%s" -p tcp --dport 25 -j DROP'%ip
		os.system(c1)
		os.system(c2)
	def iptables_del(self,ip):
		c1='iptables -D INPUT -m state --state NEW -s "%s" -p tcp --dport 25 -j LOG  --log-prefix "dynablocked "'%ip
		c2='iptables -D INPUT -m state --state NEW -s "%s" -p tcp --dport 25 -j DROP'%ip
		os.system(c1)
		os.system(c2)
		print "Dynablock: Unblocking %s"%(ip)
	def block(self,ip,t=300):
		exp=time.time()+t
		if not self.table.has_key(ip):
			self.table[ip]={'ip':ip,"expire":exp}
			self.iptables_add(ip)
	def clean(self):
		n=len(self.table)
		t=time.time()
		for k in self.table.keys():
			if t>=self.table[k]["expire"]:
				del self.table[k]
		if n!=len(self.table):
			print "Dynablock:Email cleaning number IPTABLES rules %d to %d"%(n,len(self.table))
		ips=os.popen("iptables-save")
		for l in ips.readlines():
			mo=re.match(r"-A INPUT -s ([0-9.]+) -p tcp -m state --state NEW -m tcp --dport 25 -j DROP",l)
			if mo:
				ip=mo.group(1)
				if not self.table.has_key(ip):
					self.iptables_del(ip)
	def tick(self):
		t=time.time()
		if t>=self._tick:
			self.clean()
			self._tick=t+30

class DynaBlock:
	def __init__(self,ipt):
		self._tick=0
		self.ipt=ipt
		self.email={};
	def log_smtpd_client(self,id,ip,dns):
		exp=time.time()+300
		if ip!='127.0.0.1' and ip!='213.186.38.138':
			self.email[id]={"id":id,"ip":ip,"dns":dns,"expire":exp}
	def log_smtp_status_sent(self,id,virus):
		if self.email.has_key(id):
			d=self.email[id]
			del self.email[id]
			if virus==1:
				print "Dynablock: Virus on mail %s, blocking %s (%s)"%(id,d["ip"],d["dns"])
				self.ipt.block(d["ip"])
			elif virus==2:
				print "Dynablock: BIG spam on mail %s, could block %s (%s)"%(id,d["ip"],d["dns"])
#				self.ipt.block(d["ip"],120)
	def log_line(self,l):
		mo=re.match(r"(.*) postfix/smtpd\[([0-9]+)\]: ([0-9A-F]+): client=(.*)\[([0-9.]+)\]",l)
		if mo:
			id=mo.group(3)
			ip=mo.group(5)
			dns=mo.group(4)
			self.log_smtpd_client(id,ip,dns)
			return
		mo=re.match(r"(.*) postfix/smtp\[([0-9]+)\]: ([0-9A-F]+): .* status=sent \(250 .* from MTA:",l)
		if mo:
			id=mo.group(3)
			self.log_smtp_status_sent(id,0)
			return
		mo=re.match(r"(.*) postfix/smtp\[([0-9]+)\]: ([0-9A-F]+): .* status=sent \(250 .* VIRUS:",l)
		if mo:
			id=mo.group(3)
			self.log_smtp_status_sent(id,1)
			return
		mo=re.match(r"(.*) postfix/smtp\[([0-9]+)\]: ([0-9A-F]+): .* status=sent \(250 .* discarded, UBE,",l)
		if mo:
			id=mo.group(3)
			self.log_smtp_status_sent(id,2)
			return
	def clean(self):
		n=len(self.email)
		t=time.time()
		for k in self.email.keys():
			if t>=self.email[k]["expire"]:
				del self.email[k]
		if n!=len(self.email):
			print "Dynablock:Email cleaning number watched mail from %d to %d"%(n,len(self.email))
	def tick(self):
		t=time.time()
		if t>=self._tick:
			self.clean()
			self._tick=t+30
		self.ipt.tick()

dyna=DynaBlock(IPTables())
log=os.popen("tail -n10 --follow=name /var/log/mail.log")
while 1:
	(i,o,e)=select.select([log],[],[],1)
	if len(i):
		l=log.readline()
		if not l:
			break
		dyna.log_line(l)
	dyna.tick()