poc2_worked
This commit is contained in:
Mateusz Gruszczyński
@@ -0,0 +1,191 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import pwd
|
||||
import shutil
|
||||
import socket
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
||||
class RedisSupervisor:
|
||||
"""Run the persistent Redis history service inside the IDS container."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
enabled: bool,
|
||||
data_dir: str,
|
||||
port: int = 6379,
|
||||
maxmemory_mb: int = 0,
|
||||
snapshot_seconds: int = 1800,
|
||||
aof: bool = True,
|
||||
) -> None:
|
||||
self.enabled = bool(enabled)
|
||||
self.data_dir = data_dir
|
||||
self.port = int(port)
|
||||
# 0 means unlimited. Traffic retention is time-based; Redis must not evict
|
||||
# arbitrary history just because an old deployment exported a memory cap.
|
||||
self.maxmemory_mb = max(0, int(maxmemory_mb))
|
||||
self.snapshot_seconds = max(300, int(snapshot_seconds))
|
||||
self.aof = bool(aof)
|
||||
self.executable = shutil.which("redis-server")
|
||||
self._lock = threading.RLock()
|
||||
self._stop = threading.Event()
|
||||
self._proc: subprocess.Popen | None = None
|
||||
self._thread = threading.Thread(target=self._run, name="redis-supervisor", daemon=True)
|
||||
self._restarts = 0
|
||||
self._last_error = ""
|
||||
|
||||
def start(self, *, wait_ready_seconds: float = 12.0) -> bool:
|
||||
if not self.enabled:
|
||||
self._last_error = "managed Redis disabled"
|
||||
return False
|
||||
if not self.executable:
|
||||
self._last_error = "redis-server executable not found"
|
||||
return False
|
||||
try:
|
||||
self._prepare_data_dir()
|
||||
except OSError as exc:
|
||||
self._last_error = f"cannot prepare Redis data directory: {exc}"
|
||||
return False
|
||||
self._spawn()
|
||||
if not self.wait_ready(wait_ready_seconds):
|
||||
return False
|
||||
if not self._thread.is_alive():
|
||||
self._thread.start()
|
||||
return True
|
||||
|
||||
def wait_ready(self, timeout: float = 12.0) -> bool:
|
||||
deadline = time.monotonic() + max(0.2, float(timeout))
|
||||
while time.monotonic() < deadline and not self._stop.is_set():
|
||||
with self._lock:
|
||||
proc = self._proc
|
||||
if proc is None:
|
||||
self._last_error = self._last_error or "redis-server did not start"
|
||||
return False
|
||||
code = proc.poll()
|
||||
if code is not None:
|
||||
self._last_error = f"redis-server exited with code {code}"
|
||||
return False
|
||||
if self._ping():
|
||||
self._last_error = ""
|
||||
return True
|
||||
time.sleep(0.1)
|
||||
self._last_error = self._last_error or f"Redis did not become ready on 127.0.0.1:{self.port}"
|
||||
return False
|
||||
|
||||
def stop(self) -> None:
|
||||
self._stop.set()
|
||||
with self._lock:
|
||||
proc = self._proc
|
||||
if proc is not None and proc.poll() is None:
|
||||
proc.terminate()
|
||||
try:
|
||||
proc.wait(timeout=4)
|
||||
except subprocess.TimeoutExpired:
|
||||
proc.kill()
|
||||
proc.wait(timeout=2)
|
||||
if self._thread.is_alive():
|
||||
self._thread.join(timeout=2)
|
||||
|
||||
def status(self) -> dict[str, Any]:
|
||||
with self._lock:
|
||||
proc = self._proc
|
||||
running = bool(proc is not None and proc.poll() is None)
|
||||
pid = proc.pid if running else None
|
||||
return {
|
||||
"managed": self.enabled,
|
||||
"available": bool(self.executable),
|
||||
"running": running,
|
||||
"ready": running and self._ping(),
|
||||
"pid": pid,
|
||||
"restarts": self._restarts,
|
||||
"data_dir": self.data_dir,
|
||||
"maxmemory_mb": self.maxmemory_mb,
|
||||
"snapshot_seconds": self.snapshot_seconds,
|
||||
"aof": self.aof,
|
||||
"persistence": "AOF everysec + RDB" if self.aof else "RDB",
|
||||
"last_error": self._last_error,
|
||||
}
|
||||
|
||||
def _run(self) -> None:
|
||||
while not self._stop.wait(2):
|
||||
with self._lock:
|
||||
proc = self._proc
|
||||
if proc is not None and proc.poll() is None:
|
||||
continue
|
||||
if self._stop.is_set():
|
||||
return
|
||||
self._restarts += 1
|
||||
self._spawn()
|
||||
# A restart is only considered successful once Redis accepts PING.
|
||||
self.wait_ready(8.0)
|
||||
|
||||
def _prepare_data_dir(self) -> None:
|
||||
Path(self.data_dir).mkdir(parents=True, exist_ok=True)
|
||||
try:
|
||||
user = pwd.getpwnam("redis")
|
||||
except KeyError:
|
||||
return
|
||||
for root, dirs, files in os.walk(self.data_dir):
|
||||
os.chown(root, user.pw_uid, user.pw_gid)
|
||||
for name in dirs:
|
||||
os.chown(os.path.join(root, name), user.pw_uid, user.pw_gid)
|
||||
for name in files:
|
||||
os.chown(os.path.join(root, name), user.pw_uid, user.pw_gid)
|
||||
|
||||
def _spawn(self) -> None:
|
||||
if not self.executable:
|
||||
return
|
||||
cmd = [
|
||||
self.executable,
|
||||
"--bind", "127.0.0.1",
|
||||
"--protected-mode", "yes",
|
||||
"--port", str(self.port),
|
||||
"--save", str(self.snapshot_seconds), "100",
|
||||
"--appendonly", "yes" if self.aof else "no",
|
||||
"--appendfsync", "everysec",
|
||||
"--aof-use-rdb-preamble", "yes",
|
||||
"--dir", self.data_dir,
|
||||
"--dbfilename", "traffic.rdb",
|
||||
"--maxmemory-policy", "noeviction",
|
||||
"--loglevel", "warning",
|
||||
]
|
||||
if self.maxmemory_mb > 0:
|
||||
cmd.extend(["--maxmemory", f"{self.maxmemory_mb}mb"])
|
||||
else:
|
||||
cmd.extend(["--maxmemory", "0"])
|
||||
kwargs: dict[str, Any] = {
|
||||
"stdin": subprocess.DEVNULL,
|
||||
"stdout": subprocess.DEVNULL,
|
||||
"stderr": subprocess.DEVNULL,
|
||||
"start_new_session": True,
|
||||
}
|
||||
try:
|
||||
user = pwd.getpwnam("redis")
|
||||
if os.geteuid() == 0:
|
||||
kwargs["user"] = user.pw_uid
|
||||
kwargs["group"] = user.pw_gid
|
||||
except KeyError:
|
||||
pass
|
||||
try:
|
||||
proc = subprocess.Popen(cmd, **kwargs)
|
||||
with self._lock:
|
||||
self._proc = proc
|
||||
self._last_error = ""
|
||||
except OSError as exc:
|
||||
self._last_error = str(exc)
|
||||
with self._lock:
|
||||
self._proc = None
|
||||
|
||||
def _ping(self) -> bool:
|
||||
try:
|
||||
with socket.create_connection(("127.0.0.1", self.port), timeout=0.3) as sock:
|
||||
sock.settimeout(0.3)
|
||||
sock.sendall(b"*1\r\n$4\r\nPING\r\n")
|
||||
return sock.recv(64).startswith(b"+PONG")
|
||||
except OSError:
|
||||
return False
|
||||
Reference in New Issue
Block a user