first commit
This commit is contained in:
duffyduck
+197
@@ -0,0 +1,197 @@
|
||||
"""Gemeinsame Bausteine: Verlauf, aktive Verbindungen, Relay, Ausgangsverbindungen."""
|
||||
import asyncio
|
||||
import ipaddress
|
||||
import itertools
|
||||
import logging
|
||||
import socket
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from . import config
|
||||
from .db import db
|
||||
from .policy import Decision, policy
|
||||
|
||||
log = logging.getLogger("proxy")
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- Verlauf
|
||||
class HistoryLogger:
|
||||
COLUMNS = ("ts", "username", "client_ip", "proto", "method", "host", "port", "url",
|
||||
"status", "detail", "bytes_up", "bytes_down", "duration")
|
||||
|
||||
def __init__(self):
|
||||
self.buffer: list[tuple] = []
|
||||
|
||||
def add(self, **e) -> None:
|
||||
if policy.settings.get("log_enabled", "1") != "1":
|
||||
return
|
||||
e.setdefault("ts", time.time())
|
||||
self.buffer.append(tuple(e.get(c, 0 if c in ("bytes_up", "bytes_down", "duration") else None)
|
||||
for c in self.COLUMNS))
|
||||
if len(self.buffer) > 50_000: # Notbremse, falls die DB hängt
|
||||
del self.buffer[:10_000]
|
||||
|
||||
def flush(self) -> None:
|
||||
if not self.buffer:
|
||||
return
|
||||
rows, self.buffer = self.buffer, []
|
||||
db.executemany(f"INSERT INTO history({','.join(self.COLUMNS)}) VALUES({','.join('?' * len(self.COLUMNS))})",
|
||||
rows)
|
||||
|
||||
|
||||
history = HistoryLogger()
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- aktive Verbindungen
|
||||
_ids = itertools.count(1)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Conn:
|
||||
username: str | None
|
||||
client_ip: str
|
||||
proto: str
|
||||
host: str
|
||||
port: int
|
||||
method: str = ""
|
||||
url: str = ""
|
||||
id: int = field(default_factory=lambda: next(_ids))
|
||||
started: float = field(default_factory=time.time)
|
||||
bytes_up: int = 0
|
||||
bytes_down: int = 0
|
||||
tasks: list = field(default_factory=list)
|
||||
closers: list = field(default_factory=list)
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {"id": self.id, "username": self.username, "client_ip": self.client_ip, "proto": self.proto,
|
||||
"host": self.host, "port": self.port, "method": self.method, "url": self.url,
|
||||
"started": self.started, "bytes_up": self.bytes_up, "bytes_down": self.bytes_down}
|
||||
|
||||
def kill(self) -> None:
|
||||
for t in self.tasks:
|
||||
t.cancel()
|
||||
for c in self.closers:
|
||||
try:
|
||||
c()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def finish(self, status: str = "ok", detail: str = "") -> None:
|
||||
registry.pop(self.id, None)
|
||||
history.add(ts=self.started, username=self.username, client_ip=self.client_ip, proto=self.proto,
|
||||
method=self.method, host=self.host, port=self.port, url=self.url, status=status,
|
||||
detail=detail, bytes_up=self.bytes_up, bytes_down=self.bytes_down,
|
||||
duration=round(time.time() - self.started, 3))
|
||||
|
||||
|
||||
registry: dict[int, Conn] = {}
|
||||
|
||||
|
||||
def register(conn: Conn) -> Conn:
|
||||
registry[conn.id] = conn
|
||||
return conn
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- Relay
|
||||
async def _pipe(reader: asyncio.StreamReader, writer: asyncio.StreamWriter, conn: Conn, up: bool) -> None:
|
||||
try:
|
||||
while True:
|
||||
data = await reader.read(65536)
|
||||
if not data:
|
||||
break
|
||||
writer.write(data)
|
||||
if up:
|
||||
conn.bytes_up += len(data)
|
||||
else:
|
||||
conn.bytes_down += len(data)
|
||||
await writer.drain()
|
||||
except (ConnectionError, OSError):
|
||||
pass
|
||||
finally:
|
||||
try:
|
||||
if writer.can_write_eof() and not writer.is_closing():
|
||||
writer.write_eof()
|
||||
except (OSError, RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
async def relay(conn: Conn, c_reader, c_writer, r_reader, r_writer) -> None:
|
||||
"""Leitet Daten in beide Richtungen weiter, bis eine Seite schließt."""
|
||||
up = asyncio.create_task(_pipe(c_reader, r_writer, conn, True))
|
||||
down = asyncio.create_task(_pipe(r_reader, c_writer, conn, False))
|
||||
conn.tasks += [up, down]
|
||||
conn.closers += [c_writer.close, r_writer.close]
|
||||
try:
|
||||
done, pending = await asyncio.wait({up, down}, return_when=asyncio.FIRST_COMPLETED)
|
||||
if pending:
|
||||
# Halb geschlossen: der anderen Richtung noch etwas Zeit geben
|
||||
await asyncio.wait(pending, timeout=60)
|
||||
for t in pending:
|
||||
t.cancel()
|
||||
finally:
|
||||
for w in (c_writer, r_writer):
|
||||
try:
|
||||
w.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- Ausgang
|
||||
class Blocked(Exception):
|
||||
def __init__(self, decision: Decision):
|
||||
super().__init__(decision.reason)
|
||||
self.decision = decision
|
||||
|
||||
|
||||
async def resolve(host: str, port: int, kind=socket.SOCK_STREAM) -> list[tuple]:
|
||||
loop = asyncio.get_running_loop()
|
||||
infos = await loop.getaddrinfo(host.strip("[]"), port, type=kind)
|
||||
seen, out = set(), []
|
||||
for family, _, _, _, sockaddr in infos:
|
||||
if sockaddr[0] not in seen:
|
||||
seen.add(sockaddr[0])
|
||||
out.append((family, sockaddr))
|
||||
return out
|
||||
|
||||
|
||||
async def check_and_resolve(user: dict | None, host: str, port: int, kind=socket.SOCK_STREAM) -> list[tuple]:
|
||||
"""Prüft Regeln für Hostname + jede aufgelöste IP. Gibt erlaubte Adressen zurück."""
|
||||
first = policy.decide(user, host, port)
|
||||
if not first.allowed:
|
||||
raise Blocked(first)
|
||||
try:
|
||||
addrs = await resolve(host, port, kind)
|
||||
except socket.gaierror as e:
|
||||
raise OSError(f"DNS-Auflösung fehlgeschlagen: {e.strerror or e}")
|
||||
allowed, last = [], first
|
||||
for family, sockaddr in addrs:
|
||||
d = policy.decide(user, host, port, ipaddress.ip_address(sockaddr[0].split("%")[0]))
|
||||
if d.allowed:
|
||||
allowed.append((family, sockaddr))
|
||||
else:
|
||||
last = d
|
||||
if not allowed:
|
||||
raise Blocked(last)
|
||||
return allowed
|
||||
|
||||
|
||||
async def open_upstream(user: dict | None, host: str, port: int):
|
||||
addrs = await check_and_resolve(user, host, port)
|
||||
err: Exception | None = None
|
||||
for family, sockaddr in addrs:
|
||||
try:
|
||||
return await asyncio.wait_for(asyncio.open_connection(sockaddr[0], port, family=family),
|
||||
timeout=config.CONNECT_TIMEOUT)
|
||||
except (OSError, asyncio.TimeoutError) as e:
|
||||
err = e
|
||||
if isinstance(err, asyncio.TimeoutError):
|
||||
raise OSError("Zeitüberschreitung beim Verbindungsaufbau")
|
||||
raise err or OSError("Verbindung fehlgeschlagen")
|
||||
|
||||
|
||||
def peer_ip(writer: asyncio.StreamWriter) -> str:
|
||||
peer = writer.get_extra_info("peername")
|
||||
ip = peer[0] if peer else "?"
|
||||
if ip.startswith("::ffff:"):
|
||||
ip = ip[7:]
|
||||
return ip
|
||||
Reference in New Issue
Block a user