bl-imp/main.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import argparse
import os
import subprocess
import sys

import requests
from ruamel.yaml import YAML, scalarstring


class BlacklistImporter:
	def __init__(self, args):
		self.outfile = args.outfile
		self.dryrun = args.dryrun
		self.path = os.path.dirname(__file__)
		self.url = "https://raw.githubusercontent.com/JabberSPAM/blacklist/master/blacklist.txt"
		self.blacklist = ""
		self.change = False

	def request(self):
		"""
		determine if the download is required
		"""
		etag_path = "/".join([self.path, ".etag"])
		blacklist_path = "/".join([self.path, "blacklist.txt"])

		# check if etag header is present if not set local_etag to ""
		if os.path.isfile(etag_path):
			# catch special case were etag file is present and blacklist.txt is not
			if not os.path.isfile(blacklist_path):
				local_etag = ""
			else:
				# if both files are present continue normally
				with open(etag_path, "r") as local_file:
					local_etag = local_file.read()
		else:
			local_etag = ""

		with requests.Session() as s:
			# head request to check etag
			head = s.head(self.url)
			etag = head.headers['etag']

			# if etags match up or if the connection is not possible fall back to local cache
			if local_etag == etag or head.status_code != 200:
				# if local cache is present overwrite blacklist var
				if os.path.isfile(blacklist_path):
					with open(blacklist_path, "r", encoding="utf-8") as local_file:
						self.blacklist = local_file.read()

			# in any other case request a new file
			else:
				r = s.get(self.url)
				r.encoding = 'utf-8'
				local_etag = head.headers['etag']
				self.blacklist = r.content.decode()

				with open(blacklist_path, "w") as local_file:
					local_file.write(self.blacklist)

				with open(etag_path, 'w') as local_file:
					local_file.write(local_etag)

	def main(self):
		# first check if blacklist is updated
		self.request()

		# only output the selected outfile
		if self.dryrun:
			print("outfile selected: %s" % self.outfile)

		# blacklist processing
		self.process()

		# reload config if changes have been applied
		if self.change:
			# catch ejabberdctl missing
			if os.path.isfile('/usr/sbin/ejabberdctl'):
				subprocess.call(['/usr/sbin/ejabberdctl', 'reload_config'], shell=False)

			# report missing ejabberdctl reload_config
			else:
				print('/usr/sbin/ejabberdctl was not found', file=sys.stderr)
				print('blacklist changes have been applied\nejabberd config was not reloaded', file=sys.stderr)
				sys.exit(1)

	def process(self):
		"""
		function to build and compare the local yaml file to the remote file
		if the remote file is different, the local file gets overwritten
		"""
		# init new YAML variable
		local_file = YAML(typ="safe")

		# None catch
		if self.outfile is not None:
			# prevent FileNotFoundError on first run or file missing
			if os.path.isfile(self.outfile):
				local_file = local_file.load(open(self.outfile, "r", encoding="utf-8"))

		# blacklist frame
		remote_file = {
			"acl": {
				"spamblacklist": {
					"server": []
				}
			}
		}

		# build the blacklist with the given frame to compare to local blacklist
		for entry in self.blacklist.split():
			entry = scalarstring.DoubleQuotedScalarString(entry)
			remote_file["acl"]["spamblacklist"]["server"].append(entry)

		yml = YAML()
		yml.indent(offset=2)
		yml.default_flow_style = False

		# if dry-run true print expected content
		if self.dryrun:
			yml.dump(remote_file, sys.stdout)

		# only if the local_file and remote_file are unequal write new file
		elif local_file != remote_file:

			# prevent FileNotFoundError if self.outfile is not assigned
			if self.outfile is None:
				print("no outfile assigned", file=sys.stderr)
				print(parser.format_help(), file=sys.stderr)
				sys.exit(2)

			# proceed to update the defined outfile
			elif self.outfile is not None:
				self.change = True
				yml.dump(remote_file, open(self.outfile, "w"))

			# if that's impossible break and display help message
			else:
				print(parser.format_help(), file=sys.stderr)
				sys.exit(1)


if __name__ == "__main__":
	parser = argparse.ArgumentParser()
	parser.add_argument('-o', '--outfile', help='set path to output file', dest='outfile', default=None)
	parser.add_argument('-dr', '--dry-run', help='perform a dry run', action='store_true', dest='dryrun', default=False)
	args = parser.parse_args()

	# run
	BlacklistImporter(args).main()