push: code changes

This commit is contained in:
Domingo Dirutigliano
2025-02-25 23:53:04 +01:00
parent 8652f40235
commit 6a11dd0d16
37 changed files with 1306 additions and 640 deletions

View File

@@ -1,11 +1,10 @@
from modules.nfproxy.nftables import FiregexTables
from utils import run_func
from modules.nfproxy.models import Service, PyFilter
import os
import asyncio
from utils import DEBUG
import traceback
from fastapi import HTTPException
import time
nft = FiregexTables()
@@ -13,29 +12,37 @@ class FiregexInterceptor:
def __init__(self):
self.srv:Service
self._stats_updater_cb:callable
self.filter_map_lock:asyncio.Lock
self.filter_map: dict[str, PyFilter]
self.pyfilters: set[PyFilter]
self.update_config_lock:asyncio.Lock
self.process:asyncio.subprocess.Process
self.update_task: asyncio.Task
self.server_task: asyncio.Task
self.sock_path: str
self.unix_sock: asyncio.Server
self.ack_arrived = False
self.ack_status = None
self.ack_fail_what = "Unknown"
self.ack_fail_what = "Queue response timed-out"
self.ack_lock = asyncio.Lock()
async def _call_stats_updater_callback(self, filter: PyFilter):
if self._stats_updater_cb:
await run_func(self._stats_updater_cb(filter))
self.sock_reader:asyncio.StreamReader = None
self.sock_writer:asyncio.StreamWriter = None
self.sock_conn_lock:asyncio.Lock
self.last_time_exception = 0
@classmethod
async def start(cls, srv: Service, stats_updater_cb:callable):
async def start(cls, srv: Service):
self = cls()
self._stats_updater_cb = stats_updater_cb
self.srv = srv
self.filter_map_lock = asyncio.Lock()
self.update_config_lock = asyncio.Lock()
self.sock_conn_lock = asyncio.Lock()
if not self.sock_conn_lock.locked():
await self.sock_conn_lock.acquire()
self.sock_path = f"/tmp/firegex_nfproxy_{srv.id}.sock"
if os.path.exists(self.sock_path):
os.remove(self.sock_path)
self.unix_sock = await asyncio.start_unix_server(self._server_listener,path=self.sock_path)
self.server_task = asyncio.create_task(self.unix_sock.serve_forever())
queue_range = await self._start_binary()
self.update_task = asyncio.create_task(self.update_stats())
nft.add(self.srv, queue_range)
@@ -46,19 +53,20 @@ class FiregexInterceptor:
async def _start_binary(self):
proxy_binary_path = os.path.join(os.path.dirname(os.path.abspath(__file__)),"../cpproxy")
self.process = await asyncio.create_subprocess_exec(
proxy_binary_path,
stdout=asyncio.subprocess.PIPE, stdin=asyncio.subprocess.PIPE,
proxy_binary_path, stdin=asyncio.subprocess.DEVNULL,
env={
"NTHREADS": os.getenv("NTHREADS","1"),
"FIREGEX_NFQUEUE_FAIL_OPEN": "1" if self.srv.fail_open else "0",
"FIREGEX_NFPROXY_SOCK": self.sock_path
},
)
line_fut = self.process.stdout.readuntil()
try:
line_fut = await asyncio.wait_for(line_fut, timeout=3)
async with asyncio.timeout(3):
await self.sock_conn_lock.acquire()
line_fut = await self.sock_reader.readuntil()
except asyncio.TimeoutError:
self.process.kill()
raise Exception("Invalid binary output")
raise Exception("Binary don't returned queue number until timeout")
line = line_fut.decode()
if line.startswith("QUEUE "):
params = line.split()
@@ -67,25 +75,45 @@ class FiregexInterceptor:
self.process.kill()
raise Exception("Invalid binary output")
async def _server_listener(self, reader:asyncio.StreamReader, writer:asyncio.StreamWriter):
if self.sock_reader or self.sock_writer:
writer.write_eof() # Technically never reached
writer.close()
reader.feed_eof()
return
self.sock_reader = reader
self.sock_writer = writer
self.sock_conn_lock.release()
async def update_stats(self):
try:
while True:
line = (await self.process.stdout.readuntil()).decode()
if DEBUG:
print(line)
try:
line = (await self.sock_reader.readuntil()).decode()
except Exception as e:
self.ack_arrived = False
self.ack_status = False
self.ack_fail_what = "Can't read from nfq client"
self.ack_lock.release()
await self.stop()
raise HTTPException(status_code=500, detail="Can't read from nfq client") from e
if line.startswith("BLOCKED "):
filter_id = line.split()[1]
filter_name = line.split()[1]
print("BLOCKED", filter_name)
async with self.filter_map_lock:
if filter_id in self.filter_map:
self.filter_map[filter_id].blocked_packets+=1
await self.filter_map[filter_id].update()
print("LOCKED MAP LOCK")
if filter_name in self.filter_map:
print("ADDING BLOCKED PACKET")
self.filter_map[filter_name].blocked_packets+=1
await self.filter_map[filter_name].update()
if line.startswith("MANGLED "):
filter_id = line.split()[1]
filter_name = line.split()[1]
async with self.filter_map_lock:
if filter_id in self.filter_map:
self.filter_map[filter_id].edited_packets+=1
await self.filter_map[filter_id].update()
if filter_name in self.filter_map:
self.filter_map[filter_name].edited_packets+=1
await self.filter_map[filter_name].update()
if line.startswith("EXCEPTION"):
self.last_time_exception = time.time()
print("TODO EXCEPTION HANDLING") # TODO
if line.startswith("ACK "):
self.ack_arrived = True
@@ -101,22 +129,29 @@ class FiregexInterceptor:
traceback.print_exc()
async def stop(self):
self.server_task.cancel()
self.update_task.cancel()
self.unix_sock.close()
if os.path.exists(self.sock_path):
os.remove(self.sock_path)
if self.process and self.process.returncode is None:
self.process.kill()
async def _update_config(self, code):
async with self.update_config_lock:
self.process.stdin.write(len(code).to_bytes(4, byteorder='big')+code.encode())
await self.process.stdin.drain()
try:
async with asyncio.timeout(3):
await self.ack_lock.acquire()
except TimeoutError:
pass
if not self.ack_arrived or not self.ack_status:
await self.stop()
raise HTTPException(status_code=500, detail=f"NFQ error: {self.ack_fail_what}")
if self.sock_writer:
self.sock_writer.write(len(code).to_bytes(4, byteorder='big')+code.encode())
await self.sock_writer.drain()
try:
async with asyncio.timeout(3):
await self.ack_lock.acquire()
except TimeoutError:
self.ack_fail_what = "Queue response timed-out"
if not self.ack_arrived or not self.ack_status:
await self.stop()
raise HTTPException(status_code=500, detail=f"NFQ error: {self.ack_fail_what}")
else:
raise HTTPException(status_code=400, detail="Socket not ready")
async def reload(self, filters:list[PyFilter]):
async with self.filter_map_lock:
@@ -125,12 +160,13 @@ class FiregexInterceptor:
filter_file = f.read()
else:
filter_file = ""
self.filter_map = {ele.name: ele for ele in filters}
await self._update_config(
"global __firegex_pyfilter_enabled\n" +
filter_file + "\n\n" +
"__firegex_pyfilter_enabled = [" + ", ".join([repr(f.name) for f in filters]) + "]\n" +
"__firegex_proto = " + repr(self.srv.proto) + "\n" +
"import firegex.nfproxy.internals\n\n" +
filter_file + "\n\n" +
"firegex.nfproxy.internals.compile()"
"import firegex.nfproxy.internals\n" +
"firegex.nfproxy.internals.compile(globals())\n"
)

View File

@@ -15,18 +15,18 @@ class ServiceManager:
self.srv = srv
self.db = db
self.status = STATUS.STOP
self.filters: dict[int, FiregexFilter] = {}
self.filters: dict[str, FiregexFilter] = {}
self.lock = asyncio.Lock()
self.interceptor = None
async def _update_filters_from_db(self):
pyfilters = [
PyFilter.from_dict(ele) for ele in
PyFilter.from_dict(ele, self.db) for ele in
self.db.query("SELECT * FROM pyfilter WHERE service_id = ? AND active=1;", self.srv.id)
]
#Filter check
old_filters = set(self.filters.keys())
new_filters = set([f.id for f in pyfilters])
new_filters = set([f.name for f in pyfilters])
#remove old filters
for f in old_filters:
if f not in new_filters:
@@ -34,7 +34,7 @@ class ServiceManager:
#add new filters
for f in new_filters:
if f not in old_filters:
self.filters[f] = [ele for ele in pyfilters if ele.id == f][0]
self.filters[f] = [ele for ele in pyfilters if ele.name == f][0]
if self.interceptor:
await self.interceptor.reload(self.filters.values())
@@ -43,16 +43,11 @@ class ServiceManager:
async def next(self,to):
async with self.lock:
if (self.status, to) == (STATUS.ACTIVE, STATUS.STOP):
if to == STATUS.STOP:
await self.stop()
self._set_status(to)
# PAUSE -> ACTIVE
elif (self.status, to) == (STATUS.STOP, STATUS.ACTIVE):
if to == STATUS.ACTIVE:
await self.restart()
def _stats_updater(self,filter:PyFilter):
self.db.query("UPDATE pyfilter SET blocked_packets = ?, edited_packets = ? WHERE filter_id = ?;", filter.blocked_packets, filter.edited_packets, filter.id)
def _set_status(self,status):
self.status = status
self.__update_status_db(status)
@@ -60,7 +55,7 @@ class ServiceManager:
async def start(self):
if not self.interceptor:
nft.delete(self.srv)
self.interceptor = await FiregexInterceptor.start(self.srv, self._stats_updater)
self.interceptor = await FiregexInterceptor.start(self.srv)
await self._update_filters_from_db()
self._set_status(STATUS.ACTIVE)
@@ -69,6 +64,7 @@ class ServiceManager:
if self.interceptor:
await self.interceptor.stop()
self.interceptor = None
self._set_status(STATUS.STOP)
async def restart(self):
await self.stop()

View File

@@ -15,13 +15,19 @@ class Service:
class PyFilter:
def __init__(self, filter_id:int, name: str, blocked_packets: int, edited_packets: int, active: bool, **other):
self.id = filter_id
def __init__(self, name: str, blocked_packets: int, edited_packets: int, active: bool, db, **other):
self.name = name
self.blocked_packets = blocked_packets
self.edited_packets = edited_packets
self.active = active
self.__db = db
async def update(self):
self.__db.query("UPDATE pyfilter SET blocked_packets = ?, edited_packets = ? WHERE name = ?;", self.blocked_packets, self.edited_packets, self.name)
def __repr__(self):
return f"<PyFilter {self.name}>"
@classmethod
def from_dict(cls, var: dict):
return cls(**var)
def from_dict(cls, var: dict, db):
return cls(**var, db=db)

View File

@@ -1,6 +1,14 @@
from modules.nfproxy.models import Service
from utils import ip_parse, ip_family, NFTableManager, nftables_int_to_json
def convert_protocol_to_l4(proto:str):
if proto == "tcp":
return "tcp"
elif proto == "http":
return "tcp"
else:
raise Exception("Invalid protocol")
class FiregexFilter:
def __init__(self, proto:str, port:int, ip_int:str, target:str, id:int):
self.id = id
@@ -11,7 +19,7 @@ class FiregexFilter:
def __eq__(self, o: object) -> bool:
if isinstance(o, FiregexFilter) or isinstance(o, Service):
return self.port == o.port and self.proto == o.proto and ip_parse(self.ip_int) == ip_parse(o.ip_int)
return self.port == o.port and self.proto == convert_protocol_to_l4(o.proto) and ip_parse(self.ip_int) == ip_parse(o.ip_int)
return False
class FiregexTables(NFTableManager):
@@ -61,7 +69,7 @@ class FiregexTables(NFTableManager):
"chain": self.output_chain,
"expr": [
{'match': {'left': {'payload': {'protocol': ip_family(srv.ip_int), 'field': 'saddr'}}, 'op': '==', 'right': nftables_int_to_json(srv.ip_int)}},
{'match': {"left": { "payload": {"protocol": str(srv.proto), "field": "sport"}}, "op": "==", "right": int(srv.port)}},
{'match': {"left": { "payload": {"protocol": convert_protocol_to_l4(str(srv.proto)), "field": "sport"}}, "op": "==", "right": int(srv.port)}},
{"mangle": {"key": {"meta": {"key": "mark"}},"value": 0x1338}},
{"queue": {"num": str(init) if init == end else {"range":[init, end] }, "flags": ["bypass"]}}
]
@@ -72,7 +80,7 @@ class FiregexTables(NFTableManager):
"chain": self.input_chain,
"expr": [
{'match': {'left': {'payload': {'protocol': ip_family(srv.ip_int), 'field': 'daddr'}}, 'op': '==', 'right': nftables_int_to_json(srv.ip_int)}},
{'match': {"left": { "payload": {"protocol": str(srv.proto), "field": "dport"}}, "op": "==", "right": int(srv.port)}},
{'match': {"left": { "payload": {"protocol": convert_protocol_to_l4(str(srv.proto)), "field": "dport"}}, "op": "==", "right": int(srv.port)}},
{"mangle": {"key": {"meta": {"key": "mark"}},"value": 0x1337}},
{"queue": {"num": str(init) if init == end else {"range":[init, end] }, "flags": ["bypass"]}}
]