Added client's args: --ban {...} --uban {...}

This commit is contained in:
Artyom Titov 2026-08-27 14:55:16 +03:00
commit 1b857495dc
2 changed files with 119 additions and 18 deletions

View file

@ -1,19 +1,61 @@
#!/usr/bin/env python
from enum import unique
import ipaddress
import os
import signal
import socket
import json
from collections.abc import Iterable
import argparse
def conctruct_json(ban_list: Iterable[ipaddress.IPv4Address] | Iterable[ipaddress.IPv6Address],
uban_list: Iterable[ipaddress.IPv4Address] | Iterable[ipaddress.IPv6Address],)-> str:
json_map = {"ban": {}, "uban": {} }
for ip in ban_list:
json_map["ban"][str(ip)] = {
"remark": "",
"reason": ""
}
print(json_map)
for ip in uban_list:
json_map["uban"][str(ip)] = {
"reason": ""
}
print(json_map)
return json.dumps(json_map)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("-b","--ban",nargs="+",type=str,
help="ban ip's")
parser.add_argument("-u", "--uban", nargs="+", type=str,
help="unban ip's")
args = parser.parse_args()
conflicts_ban_uban = False
if args.ban != None and args.uban != None:
conflicts_ban_uban = not (set(args.ban).isdisjoint(args.uban))
if conflicts_ban_uban:
print("--ban and --uban must not contain common element(s)")
exit(1)
if args.ban != None:
ip_to_ban = [ipaddress.IPv4Address(ip) for ip in args.ban]
else:
ip_to_ban = []
if args.uban != None:
ip_to_uban = [ipaddress.IPv4Address(ip) for ip in args.uban]
else:
ip_to_uban = []
request = "action\n" + conctruct_json(ip_to_ban, ip_to_uban)
socket_path = "/run/ipban/ipban.sock"
client = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
client.connect(socket_path)
client.sendall(b"""action
req = """action
{
"ban": {
"1.2.3.4": {
@ -26,7 +68,7 @@ if __name__ == "__main__":
}
},
"uban": {
"1.3.3.7": {
"1.4.4.6": {
"reason": "Misstake"
},
"1.0.0.1": {
@ -34,6 +76,9 @@ if __name__ == "__main__":
}
}
}
""")
"""
req = request
client.sendall(req.encode())
client.shutdown(socket.SHUT_WR)
print(client.recv(4096).decode())

View file

@ -43,7 +43,7 @@ def flush_nft(ips : Iterable[ipaddress.IPv4Address]):
elements = {{ {ips_str} }}; }}""")
#add list of ips to set in nft
def add_to_end_nft(ips: Iterable[ipaddress.IPv4Address]):
def add_to_set_nft(ips: Iterable[ipaddress.IPv4Address]):
nft = Nftables()
ips_str = ""
for ip in ips:
@ -53,13 +53,25 @@ def add_to_end_nft(ips: Iterable[ipaddress.IPv4Address]):
return
rc, output, error = nft.cmd(f"add element inet ipban banned_ipv4 {{ {ips_str} }}")
def delete_from_set_nft(ips: Iterable[ipaddress.IPv4Address]):
nft = Nftables()
ips_str = ""
for ip in ips:
ips_str += f"{str(ip)}, "
ips_str = ips_str[:-2]
if ips_str == "":
return
rc, output, error = nft.cmd(f"delete element inet ipban banned_ipv4 {{ {ips_str} }}")
def init_nft(sql_conn: sqlite3.Connection):
rows = sql_conn.execute("""SELECT ip FROM banned
""").fetchall()
flush_nft([row[0] for row in rows])
#add selected ip to db table banned
def make_ban(sql_conn: sqlite3.Connection, ip: ipaddress.IPv4Address, remark = "", reason = "") -> bool:
def sql_ban(sql_conn: sqlite3.Connection, ip: ipaddress.IPv4Address, remark = "", reason = "") -> bool:
#checking if ip is already banned
row = sql_conn.execute(f"""SELECT id FROM {ban_t_name}
WHERE ip = ? LIMIT 1""", (str(ip),)).fetchone()
@ -74,6 +86,34 @@ def make_ban(sql_conn: sqlite3.Connection, ip: ipaddress.IPv4Address, remark = "
sql_conn.commit()
return True
def sql_uban(sql_conn: sqlite3.Connection, ip: ipaddress.IPv4Address, uban_reason = "") -> bool:
ban_t_name="banned"
uban_t_name="banned_archive"
row = sql_conn.execute(f"""SELECT id FROM {ban_t_name}
WHERE ip = ? LIMIT 1""", (str(ip),)).fetchone()
print(row)
if row is None:
return False #not banned
uban_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
query = f"""INSERT INTO {uban_t_name}
(ban_time, uban_time, ip, ban_remark, ban_reason, uban_reason)
SELECT ban_time, ?, ip, remark, reason, ?
FROM {ban_t_name}
WHERE ip = ?
"""
sql_conn.execute(query, (uban_time, uban_reason, str(ip)))
sql_conn.execute(f"""DELETE FROM {ban_t_name} WHERE ip= ?""", (str(ip),))
sql_conn.commit()
return True
#what to do in case of req_type=="action"
#TODO: protect from situation: ban "1.1.1.1", uban "1.1.1.1"
def action(sql_conn: sqlite3.Connection, conn: socket.socket, data: str):
@ -81,17 +121,25 @@ def action(sql_conn: sqlite3.Connection, conn: socket.socket, data: str):
data_json = json.loads(data)
ips_to_ban = list()
for ip,creds in data_json["ban"].items():
ip_addr = ipaddress.IPv4Address(ip)
remark = creds["remark"]
reason = creds["reason"]
make_ban(sql_conn, ip_addr, remark, reason)
ips_to_ban.append(ip)
if "ban" in data_json:
for ip,creds in data_json["ban"].items():
ip_addr = ipaddress.IPv4Address(ip)
remark = creds["remark"]
reason = creds["reason"]
sql_ban(sql_conn, ip_addr, remark, reason)
ips_to_ban.append(ip)
add_to_set_nft(ips_to_ban)
ips_to_uban = list()
if "uban" in data_json:
for ip, creds in data_json["uban"].items():
ip_addr = ipaddress.IPv4Address(ip)
reason = creds["reason"]
sql_uban(sql_conn,ip_addr,reason)
ips_to_uban.append(ip)
print(ips_to_uban)
delete_from_set_nft(ips_to_uban)
for ip, creds in data_json["uban"].items():
print("Not implemented yet")
break
add_to_end_nft(ips_to_ban)
conn.sendall(b"ok")
def list_ips(sql_conn: sqlite3.Connection, conn: socket.socket, data: str):
@ -155,13 +203,21 @@ if __name__ == "__main__":
signal.signal(signal.SIGTERM, stop_handler)
sql_conn = sqlite3.connect("/var/lib/ipban/db.file")
ban_t_name="banned"
sql_conn.execute(f"""CREATE TABLE IF NOT EXISTS {ban_t_name}
uban_t_name="banned_archive"
sql_conn.execute(f"""CREATE TABLE IF NOT EXISTS {ban_t_name}
(id INTEGER PRIMARY KEY AUTOINCREMENT,
ban_time TEXT NOT NULL,
ip TEXT NOT NULL,
remark TEXT,
reason TEXT)""")
sql_conn.execute(f"""CREATE TABLE IF NOT EXISTS {uban_t_name}
(id INTEGER PRIMARY KEY AUTOINCREMENT,
ban_time TEXT NOT NULL,
uban_time TEXT NOT NULL,
ip TEXT NOT NULL,
ban_remark TEXT,
ban_reason TEXT,
uban_reason TEXT)""")
create_nft_table()
init_nft(sql_conn)
#flush_nft([ ipaddress.IPv4Address("123.134.124.1"), ipaddress.IPv4Address("8.8.8.8")])