#!/usr/bin/env python3

# fail2ban-digest: an email digest aggregator for fail2ban.
# Copyright (C) 2017 Enrico Tagliavini <enrico.tagliavini@gmail.com>
# 
# This program 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, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

from datetime import datetime, timedelta
from email.message import EmailMessage
from smtplib import SMTP
from string import Template
from time import sleep
from traceback import print_exc

import argparse
import atexit
import dbm
import errno
import os
import socket
import re
import sys

db_location = '/var/lib/fail2ban/digest'
db_creation_date_key = 'db_creation_date'
db_date_format = '%Y-%m-%d %H:%M:%S'
default_mail_template = Template('''Hi,\n
   this is a digest email of banned IPs since ${creation_date} UTC and ${date_now}

${digest}

Regards,

Fail2ban digest
''')

class store_yesno(argparse.Action):
	def __init__(self, option_strings, dest, nargs = None, **kwargs):
		if nargs is not None:
			raise ValueError("nargs not allowed")
		super(store_yesno, self).__init__(option_strings, dest, nargs = 0, **kwargs)

	def __call__(self, parser, namespace, values, option_string = None):
		if option_string.startswith('--no-'):
			setattr(namespace, self.dest, False)
		else:
			setattr(namespace, self.dest, True)
		return

ipv4_re = re.compile(r'^(\d){1,3}\.(\d){1,3}\.(\d){1,3}\.(\d){1,3}$')

def ip_address(string):
	valid = 2
	try:
		socket.inet_pton(socket.AF_INET, string)
		# this can match also stuff like '4' or '127.1' or '127.1.256'
		# most likely a typo rather than intentional, so refuse it
		if ipv4_re.search(string) is None:
			valid -= 1
	except socket.error:
		valid -= 1
	try:
		socket.inet_pton(socket.AF_INET6, string)
	except socket.error:
		valid -= 1
	if valid > 0:
		return string
	else:
		raise argparse.ArgumentTypeError('%s is not a valid IPv4 or IPv6 address' % string)

def close_db(db_fd):
	db_fd.close()
	atexit.unregister(close_db)
	return

def db_busy_open(filename, flags, timeout):
	start = datetime.utcnow()
	threshold = timedelta(seconds = timeout)
	while True:
		try:
			db = dbm.open(filename, flags)
			break
		except OSError as e:
			if e.errno == errno.EAGAIN: # errno 11, try again in a moment
				sleep(1)
			else:
				raise e
		if datetime.utcnow() - start > threshold:
			raise TimeoutError(errno.ETIMEDOUT, 'timeout while waiting for database file to be unlocked', filename)
	atexit.register(close_db, db)
	return db

def add(db, ip):
	db = db_busy_open(db_location + '/' + db + '.dbm', 'c', 30)
	if db_creation_date_key not in db.keys():
		db[db_creation_date_key] = datetime.utcnow().strftime(db_date_format).encode('UTF-8')
	event_date = ('%s, ' % datetime.now().strftime(db_date_format)).encode('UTF-8')
	try:
		db[ip] += event_date
	except KeyError:
		db[ip] = event_date
	close_db(db)
	return

def digest(db, delete):
	db_file = db_location + '/' + db + '.dbm'
	db = db_busy_open(db_file, 'r', 30)  # this is just a trick to lock the DB, to have a consistent copy
	if delete:
		#os.rename(db_file, db_file + '.digest') # just in case we want to delay the removal in the future
		os.unlink(db_file)
	try:
		db_creation_date = db[db_creation_date_key].decode('UTF-8')
	except KeyError as e:
		db_creation_date = 'not found'
	event_list = []
	for ip in db.keys():
		if ip.decode('UTF-8') == db_creation_date_key:
			continue
		event_list.append((ip.decode('UTF-8'), db[ip].decode('UTF-8').split(', ')[:-1]))
	close_db(db)
	event_list.sort(key = lambda x: len(x[1]), reverse = True)
	msg = ''
	for ip, events in event_list:
		msg += '%3d event(s) for IP %-42s: %s\n' %(len(events), ip, ', '.join(events))
	return (db_creation_date, msg)

def mail_digest(db, mail_to, mail_from, delete):
	msg = EmailMessage()
	date_now = datetime.utcnow().strftime(db_date_format)
	creation_date, dgst = digest(db, delete)
	msg.set_content(default_mail_template.substitute(
			creation_date = creation_date,
			date_now = date_now,
			digest = dgst,
	))
	msg['To'] = mail_to
	msg['From'] = mail_from
	msg['Subject'] = '[Fail2Ban] %s: digest for %s %s' % (db, socket.gethostname(), date_now)
	mta = SMTP(host = 'localhost')
	mta.send_message(msg)
	mta.quit()
	return

def main(args):
	if args.cmd == 'add':
		add(args.database, args.ip)
	elif args.cmd == 'digest':
		print(digest(args.database, args.delete)[1])
	elif args.cmd == 'maildigest':
		mail_digest(args.database, args.to, args.mail_from, args.delete)
	return

if __name__ == '__main__':
	progname = os.path.basename(sys.argv[0])
	parser = argparse.ArgumentParser(
			prog = progname,
			description = 'Gather fail2ban events to process periodically and generate a digest',
			epilog = 'use `%s command --help\' to get help about a specific command' % progname,
	)
	subparsers = parser.add_subparsers(title = 'available commands')

	subcommands = {}
	sc = 'add'
	subcommands[sc] = subparsers.add_parser(
			sc,
			description = 'Add a new event to the DB'
	)
	subcommands[sc].add_argument(
			'database',
			help = 'database where to store the event'
	)

	subcommands[sc].add_argument(
			'ip',
			type = ip_address,
			help = 'offending IP address, both IPv4 and IPv6 are accepted'
	)
	subcommands[sc].set_defaults(cmd = sc)

	sc = 'digest'
	subcommands[sc] = subparsers.add_parser(
			sc,
			description = 'print database digest to standard output'
	)
	subcommands[sc].add_argument(
			'database',
			help = 'database to generate the digest from'
	)
	subcommands[sc].add_argument(
			'--delete', '--no-delete',
			action = store_yesno,
			default = False,
			help = 'do / don\'t delete current database, next call to add will create a new empty one'
	)
	subcommands[sc].set_defaults(cmd = sc)

	sc = 'maildigest'
	subcommands[sc] = subparsers.add_parser(
			sc,
			description = 'send database digest via email'
	)
	subcommands[sc].add_argument(
			'database',
			help = 'database to generate the digest from'
	)
	subcommands[sc].add_argument(
			'--delete', '--no-delete',
			action = store_yesno,
			default = True,
			help = 'do / don\'t delete current database, next call to add will create a new empty one'
	)
	subcommands[sc].add_argument(
			'--mail-from',
			action = 'store',
			default = 'Fail2ban at {0} <fail2ban@{0}>'.format(socket.gethostname()),
			help = 'Use FROM address for the email From header. Default: Fail2ban at {0} <fail2ban@{0}> (automatically detected system hostname)'.format(socket.gethostname())
	)
	subcommands[sc].add_argument(
			'--to',
			action = 'store',
			default = 'root',
			help = 'send email to specified user / address. Default is root'
	)
	subcommands[sc].set_defaults(cmd = sc)


	args = parser.parse_args(sys.argv[1:])
	#print(args)
	try:
		main(args)
	except KeyboardInterrupt:
		sys.exit(1)
	except Exception:
		print_exc()
		sys.exit(1)