push: code changes x2
This commit is contained in:
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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"
|
||||
}}}
|
||||
],[
|
||||
|
||||
Reference in New Issue
Block a user