Rewrote in bash

tmpufw.py

@@ -0,0 +1,162 @@
+#!/usr/bin/env python3
+"""
+Temporarily apply `ufw` rules
+
+This script allows you to add rules to `ufw` (Uncomplicated Firewall) with a
+time to live. You can then run the script as a cronjob (with the --clean flag)
+to clean up (remove) the expired rules.
+
+Arguments:
+	-h, --help                       show the help message and exit
+	-s, --status                     show rule list with expirations
+	-c, --clean                      clean up expired rules
+	-r RULE, --rule RULE             rule to be added to `ufw`
+	-p POSITION, --position POSITION position to add the rule
+	-t TTL, --ttl TTL                time to live for the rule
+"""
+__author__  = 'Joshua Sherman'
+__file__    = 'tmpufw'
+__license__ = 'MIT'
+__status__  = 'Production'
+__version__ = '1.0.0'
+
+from argparse import ArgumentParser
+from datetime import datetime
+from os import getpid, makedirs, path, remove
+from parsedatetime import Calendar
+from shutil import move
+from subprocess import CalledProcessError, check_output, STDOUT
+from sys import exit
+from time import mktime, time
+
+class tmpufw(object):
+	parser = ArgumentParser(description = 'Temporarily apply `ufw` rules')
+
+	def __init__(self):
+		self.parser.add_argument('-s', '--status', action = 'store_true', help = 'show rule list with expirations')
+		self.parser.add_argument('-c', '--clean', action = 'store_true', help = 'clean up expired rules')
+		self.parser.add_argument('-r', '--rule', help = 'rule to be added to `ufw`')
+		self.parser.add_argument('-p', '--position', default = 1, help = 'position to add the rule')
+		self.parser.add_argument('-t', '--ttl', default = '30 days', help = 'time to live for the rule')
+		args = self.parser.parse_args()
+
+		# Our file names
+		pid_file       = '/var/run/' + __file__ + '.pid'
+		rules_file     = '/usr/local/share/' + __file__ + '/rules'
+		tmp_rules_file = '/tmp/' + __file__ + '-rules'
+
+		if args.status:
+			if path.exists(rules_file):
+				try:
+					print("Expiration\t\tRule")
+					print('=' * 80)
+
+					# Loops through the rules lines
+					for line in open(rules_file, 'r'):
+						# Breaks apart line into expiration timestamp and rule
+						timestamp, rule = line.strip("\n").split(' ', 1)
+
+						print(str(datetime.fromtimestamp(float(timestamp))) + "\t" + rule)
+				except IOError:
+					self.error('unable to read from the rules file: ' + rules_file)
+			else:
+				self.error('there are no rules to display')
+		elif args.clean:
+			# Checks for PID file
+			if path.exists(pid_file):
+				self.error(__file__ + ' is already running')
+			else:
+				# Creates the PID file
+				try:
+					handle = open(pid_file, 'w')
+					handle.write(str(getpid()))
+					handle.close()
+				except IOError:
+					self.error('unable to create PID file: ' + pid_file)
+
+				# Checks for the rules file
+				if path.exists(rules_file):
+					# Opens the temporary rules file
+					try:
+						handle = open(tmp_rules_file, 'a')
+					except IOError:
+						self.error('unable to write to the tmp rules file: ' + tmp_rules_file)
+
+					try:
+						current_time = time()
+
+						# Loops through the rules lines
+						for line in open(rules_file, 'r'):
+							# Breaks apart line into expiration timestamp and rule
+							timestamp, rule = line.strip("\n").split(' ', 1)
+
+							# Checks if rule has expired
+							if current_time < float(timestamp):
+								handle.write(line)
+								print(str(datetime.fromtimestamp(time())) + "\tskipped rule\t" + rule)
+							else:
+								try:
+									self.ufw_execute('delete ' + rule)
+									print(str(datetime.fromtimestamp(time())) + "\tdeleted rule\t" + rule)
+								except CalledProcessError as error:
+									self.ufw_error(error)
+
+						handle.close()
+
+						# Moves the tmp file to the rules file
+						move(tmp_rules_file, rules_file)
+					except IOError:
+						self.error('unable to from the read rules file: ' + rules_file)
+
+				# Removes the PID
+				remove(pid_file)
+		elif args.rule:
+			rules_path = path.dirname(rules_file)
+
+			if not path.exists(rules_path):
+				makedirs(rules_path)
+
+			# Converts the TTL to a timestamp
+			cal       = Calendar()
+			timestamp = mktime(cal.parse(args.ttl)[0])
+
+			# Writes the rule to the rules file
+			try:
+				# TODO Check if rule already exists and update it instead of adding it again
+				handle = open(rules_file, 'a')
+				handle.write(str(timestamp) + ' ' + args.rule)
+				handle.write("\n")
+				handle.close()
+			except IOError:
+				self.error('unable to write to the rules file: ' + rules_file)
+
+			# Attempts to add the rule to `ufw`
+			try:
+				self.ufw_execute('insert ' + str(args.position) + ' ' + args.rule)
+			except CalledProcessError as error:
+				# Catches an error when attempting to add a rule to an empty database
+				if error.output == b"ERROR: Invalid position '1'\n":
+					try:
+						self.ufw_execute(args.rule)
+					except CalledProcessError as error:
+						self.ufw_error(error)
+				else:
+					self.ufw_error(error)
+		else:
+			self.error('no arguments specified')
+
+	def error(self, message):
+		self.parser.print_usage()
+		print(__file__ + ': error: ' + message)
+		exit(2)
+
+	def ufw_execute(self, rule):
+		for arg in [' --dry-run ', ' ']:
+			command = 'ufw' + arg + rule
+			check_output(command, stderr = STDOUT, shell = True)
+
+	def ufw_error(self, error):
+		self.error('ufw: ' + error.output.decode(encoding = 'UTF-8'))
+
+if __name__ == '__main__':
+	tmpufw()