push: code changes x2

This commit is contained in:
Domingo Dirutigliano
2025-02-28 21:14:09 +01:00
parent 6a11dd0d16
commit 8ae533e8f7
31 changed files with 544 additions and 397 deletions

View File

@@ -5,9 +5,12 @@ import asyncio
import traceback
from fastapi import HTTPException
import time
from utils import run_func
nft = FiregexTables()
OUTSTREAM_BUFFER_SIZE = 1024*10
class FiregexInterceptor:
def __init__(self):
@@ -28,14 +31,20 @@ class FiregexInterceptor:
self.sock_writer:asyncio.StreamWriter = None
self.sock_conn_lock:asyncio.Lock
self.last_time_exception = 0
self.outstrem_function = None
self.expection_function = None
self.outstrem_task: asyncio.Task
self.outstrem_buffer = ""
@classmethod
async def start(cls, srv: Service):
async def start(cls, srv: Service, outstream_func=None, exception_func=None):
self = cls()
self.srv = srv
self.filter_map_lock = asyncio.Lock()
self.update_config_lock = asyncio.Lock()
self.sock_conn_lock = asyncio.Lock()
self.outstrem_function = outstream_func
self.expection_function = exception_func
if not self.sock_conn_lock.locked():
await self.sock_conn_lock.acquire()
self.sock_path = f"/tmp/firegex_nfproxy_{srv.id}.sock"
@@ -50,16 +59,37 @@ class FiregexInterceptor:
await self.ack_lock.acquire()
return self
async def _stream_handler(self):
while True:
try:
line = (await self.process.stdout.readuntil()).decode(errors="ignore")
print(line, end="")
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
self.outstrem_buffer+=line
if len(self.outstrem_buffer) > OUTSTREAM_BUFFER_SIZE:
self.outstrem_buffer = self.outstrem_buffer[-OUTSTREAM_BUFFER_SIZE:]+"\n"
if self.outstrem_function:
await run_func(self.outstrem_function, self.srv.id, line)
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, stdin=asyncio.subprocess.DEVNULL,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.STDOUT,
env={
"NTHREADS": os.getenv("NTHREADS","1"),
"FIREGEX_NFQUEUE_FAIL_OPEN": "1" if self.srv.fail_open else "0",
"FIREGEX_NFPROXY_SOCK": self.sock_path
},
)
self.outstrem_task = asyncio.create_task(self._stream_handler())
try:
async with asyncio.timeout(3):
await self.sock_conn_lock.acquire()
@@ -101,9 +131,7 @@ class FiregexInterceptor:
filter_name = line.split()[1]
print("BLOCKED", filter_name)
async with self.filter_map_lock:
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 "):
@@ -113,8 +141,9 @@ class FiregexInterceptor:
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
self.last_time_exception = int(time.time()*1000) #ms timestamp
if self.expection_function:
await run_func(self.expection_function, self.srv.id, self.last_time_exception)
if line.startswith("ACK "):
self.ack_arrived = True
self.ack_status = line.split()[1].upper() == "OK"
@@ -132,6 +161,7 @@ class FiregexInterceptor:
self.server_task.cancel()
self.update_task.cancel()
self.unix_sock.close()
self.outstrem_task.cancel()
if os.path.exists(self.sock_path):
os.remove(self.sock_path)
if self.process and self.process.returncode is None:

View File

@@ -3,6 +3,7 @@ from modules.nfproxy.firegex import FiregexInterceptor
from modules.nfproxy.nftables import FiregexTables, FiregexFilter
from modules.nfproxy.models import Service, PyFilter
from utils.sqlite import SQLite
from utils import run_func
class STATUS:
STOP = "stop"
@@ -11,13 +12,20 @@ class STATUS:
nft = FiregexTables()
class ServiceManager:
def __init__(self, srv: Service, db):
def __init__(self, srv: Service, db, outstream_func=None, exception_func=None):
self.srv = srv
self.db = db
self.status = STATUS.STOP
self.filters: dict[str, FiregexFilter] = {}
self.lock = asyncio.Lock()
self.interceptor = None
self.outstream_function = outstream_func
self.last_exception_time = 0
async def excep_internal_handler(srv, exc_time):
self.last_exception_time = exc_time
if exception_func:
await run_func(exception_func, srv, exc_time)
self.exception_function = excep_internal_handler
async def _update_filters_from_db(self):
pyfilters = [
@@ -52,10 +60,16 @@ class ServiceManager:
self.status = status
self.__update_status_db(status)
def read_outstrem_buffer(self):
if self.interceptor:
return self.interceptor.outstrem_buffer
else:
return ""
async def start(self):
if not self.interceptor:
nft.delete(self.srv)
self.interceptor = await FiregexInterceptor.start(self.srv)
self.interceptor = await FiregexInterceptor.start(self.srv, outstream_func=self.outstream_function, exception_func=self.exception_function)
await self._update_filters_from_db()
self._set_status(STATUS.ACTIVE)
@@ -75,10 +89,12 @@ class ServiceManager:
await self._update_filters_from_db()
class FirewallManager:
def __init__(self, db:SQLite):
def __init__(self, db:SQLite, outstream_func=None, exception_func=None):
self.db = db
self.service_table: dict[str, ServiceManager] = {}
self.lock = asyncio.Lock()
self.outstream_function = outstream_func
self.exception_function = exception_func
async def close(self):
for key in list(self.service_table.keys()):
@@ -100,7 +116,7 @@ class FirewallManager:
srv = Service.from_dict(srv)
if srv.id in self.service_table:
continue
self.service_table[srv.id] = ServiceManager(srv, self.db)
self.service_table[srv.id] = ServiceManager(srv, self.db, outstream_func=self.outstream_function, exception_func=self.exception_function)
await self.service_table[srv.id].next(srv.status)
def get(self,srv_id) -> ServiceManager:

View File

@@ -28,22 +28,22 @@ class FiregexTables(NFTableManager):
def __init__(self):
super().__init__([
{"add":{"chain":{
{"add":{"chain":{ #Input chain attached before conntrack see it
"family":"inet",
"table":self.table_name,
"name":self.input_chain,
"type":"filter",
"hook":"prerouting",
"prio":-150,
"prio":-301,
"policy":"accept"
}}},
{"add":{"chain":{
{"add":{"chain":{ #Output chain attached after conntrack saw it
"family":"inet",
"table":self.table_name,
"name":self.output_chain,
"type":"filter",
"hook":"postrouting",
"prio":-150,
"prio":-290,
"policy":"accept"
}}}
],[