removed fastapi_socketio + general improves
This commit is contained in:
Domingo Dirutigliano
@@ -1,97 +1,61 @@
|
||||
from modules.nfregex.nftables import FiregexTables
|
||||
from modules.nfproxy.nftables import FiregexTables
|
||||
from utils import run_func
|
||||
from modules.nfregex.models import Service, Regex
|
||||
import re
|
||||
from modules.nfproxy.models import Service, PyFilter
|
||||
import os
|
||||
import asyncio
|
||||
import traceback
|
||||
from utils import DEBUG
|
||||
from fastapi import HTTPException
|
||||
|
||||
#TODO copied file, review
|
||||
import socket
|
||||
import shutil
|
||||
|
||||
nft = FiregexTables()
|
||||
|
||||
class RegexFilter:
|
||||
def __init__(
|
||||
self, regex,
|
||||
is_case_sensitive=True,
|
||||
input_mode=False,
|
||||
output_mode=False,
|
||||
blocked_packets=0,
|
||||
id=None,
|
||||
update_func = None
|
||||
):
|
||||
self.regex = regex
|
||||
self.is_case_sensitive = is_case_sensitive
|
||||
if input_mode == output_mode:
|
||||
input_mode = output_mode = True # (False, False) == (True, True)
|
||||
self.input_mode = input_mode
|
||||
self.output_mode = output_mode
|
||||
self.blocked = blocked_packets
|
||||
self.id = id
|
||||
self.update_func = update_func
|
||||
self.compiled_regex = self.compile()
|
||||
|
||||
@classmethod
|
||||
def from_regex(cls, regex:Regex, update_func = None):
|
||||
return cls(
|
||||
id=regex.id, regex=regex.regex, is_case_sensitive=regex.is_case_sensitive,
|
||||
blocked_packets=regex.blocked_packets,
|
||||
input_mode = regex.mode in ["C","B"], output_mode=regex.mode in ["S","B"],
|
||||
update_func = update_func
|
||||
)
|
||||
def compile(self):
|
||||
if isinstance(self.regex, str):
|
||||
self.regex = self.regex.encode()
|
||||
if not isinstance(self.regex, bytes):
|
||||
raise Exception("Invalid Regex Paramether")
|
||||
re.compile(self.regex) # raise re.error if it's invalid!
|
||||
case_sensitive = "1" if self.is_case_sensitive else "0"
|
||||
if self.input_mode:
|
||||
yield case_sensitive + "C" + self.regex.hex()
|
||||
if self.output_mode:
|
||||
yield case_sensitive + "S" + self.regex.hex()
|
||||
|
||||
async def update(self):
|
||||
if self.update_func:
|
||||
await run_func(self.update_func, self)
|
||||
|
||||
class FiregexInterceptor:
|
||||
|
||||
def __init__(self):
|
||||
self.srv:Service
|
||||
self.filter_map_lock:asyncio.Lock
|
||||
self.filter_map: dict[str, RegexFilter]
|
||||
self.regex_filters: set[RegexFilter]
|
||||
self.update_config_lock:asyncio.Lock
|
||||
self._stats_updater_cb:callable
|
||||
self.process:asyncio.subprocess.Process
|
||||
self.update_task: asyncio.Task
|
||||
self.ack_arrived = False
|
||||
self.ack_status = None
|
||||
self.ack_fail_what = ""
|
||||
self.ack_lock = asyncio.Lock()
|
||||
self.base_dir = os.path.join(
|
||||
os.path.dirname(os.path.abspath(__file__)),
|
||||
"socks", self.srv.id
|
||||
)
|
||||
self.n_threads = int(os.getenv("NTHREADS","1"))
|
||||
|
||||
self.connection_socket = os.path.join(self.base_dir, "connection.sock")
|
||||
self.vedict_sockets = [os.path.join(self.base_dir, f"vedict{i}.sock") for i in range(self.n_threads)]
|
||||
self.socks = []
|
||||
|
||||
def add_sock(self, path):
|
||||
sock = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
|
||||
sock.bind(path)
|
||||
self.socks.append(sock)
|
||||
return sock
|
||||
|
||||
async def _call_stats_updater_callback(self, filter: PyFilter):
|
||||
if self._stats_updater_cb:
|
||||
await run_func(self._stats_updater_cb(filter))
|
||||
|
||||
@classmethod
|
||||
async def start(cls, srv: Service):
|
||||
async def start(cls, srv: Service, stats_updater_cb:callable):
|
||||
self = cls()
|
||||
self.srv = srv
|
||||
self.filter_map_lock = asyncio.Lock()
|
||||
self.update_config_lock = asyncio.Lock()
|
||||
self._stats_updater_cb = stats_updater_cb
|
||||
os.makedirs(self.base_dir, exist_ok=True)
|
||||
self.add_sock(self.connection_socket)
|
||||
for path in self.vedict_sockets:
|
||||
self.add_sock(path)
|
||||
queue_range = await self._start_binary()
|
||||
self.update_task = asyncio.create_task(self.update_blocked())
|
||||
# TODO starts python workers
|
||||
nft.add(self.srv, queue_range)
|
||||
if not self.ack_lock.locked():
|
||||
await self.ack_lock.acquire()
|
||||
return self
|
||||
|
||||
async def _start_binary(self):
|
||||
proxy_binary_path = os.path.join(os.path.dirname(os.path.abspath(__file__)),"../cppqueue")
|
||||
proxy_binary_path = os.path.join(os.path.dirname(os.path.abspath(__file__)),"../cppproxy")
|
||||
self.process = await asyncio.create_subprocess_exec(
|
||||
proxy_binary_path,
|
||||
stdout=asyncio.subprocess.PIPE, stdin=asyncio.subprocess.PIPE,
|
||||
env={"MATCH_MODE": "stream" if self.srv.proto == "tcp" else "block", "NTHREADS": os.getenv("NTHREADS","1")},
|
||||
)
|
||||
self.process.stdin.write(self.base_dir.encode().hex().encode()+b" 3\n")
|
||||
await self.process.stdin.drain()
|
||||
line_fut = self.process.stdout.readuntil()
|
||||
try:
|
||||
line_fut = await asyncio.wait_for(line_fut, timeout=3)
|
||||
@@ -106,68 +70,14 @@ class FiregexInterceptor:
|
||||
self.process.kill()
|
||||
raise Exception("Invalid binary output")
|
||||
|
||||
async def update_blocked(self):
|
||||
try:
|
||||
while True:
|
||||
line = (await self.process.stdout.readuntil()).decode()
|
||||
if DEBUG:
|
||||
print(line)
|
||||
if line.startswith("BLOCKED "):
|
||||
regex_id = line.split()[1]
|
||||
async with self.filter_map_lock:
|
||||
if regex_id in self.filter_map:
|
||||
self.filter_map[regex_id].blocked+=1
|
||||
await self.filter_map[regex_id].update()
|
||||
if line.startswith("ACK "):
|
||||
self.ack_arrived = True
|
||||
self.ack_status = line.split()[1].upper() == "OK"
|
||||
if not self.ack_status:
|
||||
self.ack_fail_what = " ".join(line.split()[2:])
|
||||
self.ack_lock.release()
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
except asyncio.IncompleteReadError:
|
||||
pass
|
||||
except Exception:
|
||||
traceback.print_exc()
|
||||
|
||||
async def stop(self):
|
||||
self.update_task.cancel()
|
||||
if self.process and self.process.returncode is None:
|
||||
self.process.kill()
|
||||
|
||||
async def _update_config(self, filters_codes):
|
||||
async with self.update_config_lock:
|
||||
self.process.stdin.write((" ".join(filters_codes)+"\n").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:
|
||||
raise HTTPException(status_code=500, detail=f"NFQ error: {self.ack_fail_what}")
|
||||
|
||||
|
||||
async def reload(self, filters:list[RegexFilter]):
|
||||
async with self.filter_map_lock:
|
||||
self.filter_map = self.compile_filters(filters)
|
||||
filters_codes = self.get_filter_codes()
|
||||
await self._update_config(filters_codes)
|
||||
|
||||
def get_filter_codes(self):
|
||||
filters_codes = list(self.filter_map.keys())
|
||||
filters_codes.sort(key=lambda a: self.filter_map[a].blocked, reverse=True)
|
||||
return filters_codes
|
||||
|
||||
def compile_filters(self, filters:list[RegexFilter]):
|
||||
res = {}
|
||||
for filter_obj in filters:
|
||||
try:
|
||||
raw_filters = filter_obj.compile()
|
||||
for filter in raw_filters:
|
||||
res[filter] = filter_obj
|
||||
except Exception:
|
||||
pass
|
||||
return res
|
||||
for sock in self.socks:
|
||||
sock.close()
|
||||
shutil.rmtree(self.base_dir)
|
||||
|
||||
async def reload(self, filters:list[PyFilter]):
|
||||
# filters are the functions to use in the workers (other functions are disabled or not flagged as filters)
|
||||
# TODO update filters in python workers (prob for new filters added) (reading from file????)
|
||||
pass
|
||||
Reference in New Issue
Block a user