198 lines
6.7 KiB
Python
198 lines
6.7 KiB
Python
"""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
|