Automatically ban IPs logged in btmp via firewall-cmd

btmp_ban

#!/usr/local/bin/python3

"""
This is a quick and dirty solution I wrote in 10 minutes, which then I piled more stuffs
without properly reorganizing the code.
It surely works, but is ugly, slow, and a mess.

`ipblack` list includes IP address prefixes that should be excluded from being banned.
this includes internal IP, and maybe IP from your ISP's cellular in case you want to
connect to your server from the outside.
you probably have to change it manually.

This script pings ip-api.com for every IP addresses it checks. this adds a significant
time delay, or cause problems in case of a network outage. it can be disabled by adding
"noquery" in the argument.

this script also manually asks you if the IP needs to be banned or not. if you want to ban
then automatically, add "silent" in the argument.

banned IPs will be written in "bannedip.txt" in case you want to check which one has been banned.

This code is written to be used with CentOS, which utilizes firewall-cmd. if your system is
not CentOS(Which is good for you), you might have to change the zone name (Fedora has a completely
different name for example), or just rewrite the whole code to work with iptables instead of firewalld.
"""

import os
import re
import sys
import time
import requests
import subprocess
from subprocess import PIPE

ipblack = ["192.168", "223.38", "223.62"]

if "silent" in sys.argv:
	silent = True
else:
	silent = False

if "noquery" in sys.argv:
	noquery = True
else:
	noquery = False

def clear():
	return os.system("clear")

def print_geoip(query):
	r = requests.get(f"http://ip-api.com/json/{query}?fields=status,message,country,regionName,city,isp")
	r.raise_for_status()
	data = r.json()
	if data['status'] == "fail":
		print(f"GeoIP Lookup Failed.\nReason: {data['message']}")
		return
	for key in data:
		print(f"{key}: {data[key]}") if (not key == "status") and data[key] else None
	return

def grab_lastb():
	return subprocess.run("lastb", stdout = PIPE).stdout.decode().split("\n")

def import_firewalld_rules():
	return re.findall(
		"""rule family="ipv4" source address="(.*?)" drop""",
		subprocess.run(["firewall-cmd", "--zone=public", "--list-all"], stdout=PIPE).stdout.decode()
	)

def banhammer(query):
	return os.system(f"""firewall-cmd --permanent --zone=public --add-rich-rule='rule family="ipv4" source address="{query}" drop'""")

def reload():
	return os.system("firewall-cmd --reload")

bannedip = import_firewalld_rules()

lastb = grab_lastb()
clear()
print(lastb[-2])
print(silent)
time.sleep(3)

for line in lastb[:-3]:
	linesplit = line.split()
	ip = linesplit[2]
	if ip in bannedip or ".".join(ip.split()[:2]) in ipblack:
		continue
	if not silent:
		print(line)
		if not noquery:
			print_geoip(ip)
		if input("Ban this IP?(y/n): ") == "y":
			if banhammer(ip) != 0:raise
			bannedip.append(ip)
		clear()
	else:
		print(line)
		if banhammer(ip) != 0:
			if silent:continue
			else:raise
		bannedip.append(ip)
clear()

if silent or input("Reload firewall-cmd(y/n)?: ") == "y":  reload()

with open("bannedip.txt", "a") as f:
	f.write("=====begins=====\n")
	f.write("\n".join(bannedip))
	f.write("\n")

print("finished.")