#!/usr/bin/env python3
"""Cliente que habla PROXY protocol v1 y v2 contra un PostgreSQL parchado.

    proxy_client.py <puerto> <caso>

No usa libpq: arma el header PROXY a mano y despues el StartupMessage, que es
lo unico que hace falta para comprobar que el servidor toma la direccion del
header y no la del socket.  Asi se puede probar tambien lo que un cliente
normal no puede mandar: headers invalidos, truncados o de familias raras.

Casos:
  v1            header de texto valido, IPv4
  v1_ipv6       header de texto valido, IPv6
  v2            header binario valido, IPv4
  v2_local      header binario con comando LOCAL (el proxy habla de si mismo)
  sin_header    conexion directa, sin header
  v1_basura     "PROXY " seguido de basura
  v1_corto      header v1 cortado a la mitad
  v2_mal_sig    header binario con la firma equivocada
  v2_largo      v2 que declara mas bytes de los que manda
"""
import socket
import struct
import sys

V2_SIG = b"\r\n\r\n\x00\r\nQUIT\n"


def startup(user=b"postgres", db=b"postgres"):
    body = struct.pack("!i", 196608) + b"user\x00" + user + b"\x00" \
        + b"database\x00" + db + b"\x00\x00"
    return struct.pack("!i", len(body) + 4) + body


def header(caso):
    if caso == "v1":
        return b"PROXY TCP4 198.51.100.7 203.0.113.9 51234 5432\r\n"
    if caso == "v1_ipv6":
        return b"PROXY TCP6 2001:db8::7 2001:db8::9 51234 5432\r\n"
    if caso == "v2":
        # ver 2 / PROXY, TCP over IPv4, 12 bytes de direcciones
        addr = socket.inet_aton("198.51.100.7") + socket.inet_aton("203.0.113.9") \
            + struct.pack("!HH", 51234, 5432)
        return V2_SIG + bytes([0x21, 0x11]) + struct.pack("!H", len(addr)) + addr
    if caso == "v2_local":
        return V2_SIG + bytes([0x20, 0x00]) + struct.pack("!H", 0)
    if caso == "sin_header":
        return b""
    if caso == "v1_basura":
        return b"PROXY estoesbasura\r\n"
    if caso == "v1_corto":
        return b"PROXY TCP4 198.51.100.7 203."
    if caso == "v2_mal_sig":
        addr = socket.inet_aton("198.51.100.7") + socket.inet_aton("203.0.113.9") \
            + struct.pack("!HH", 51234, 5432)
        return b"\x00" * 12 + bytes([0x21, 0x11]) + struct.pack("!H", len(addr)) + addr
    if caso == "v2_largo":
        return V2_SIG + bytes([0x21, 0x11]) + struct.pack("!H", 200) + b"\x00" * 4
    raise SystemExit(f"caso desconocido: {caso}")


def main(port, caso):
    s = socket.create_connection(("127.0.0.1", int(port)), timeout=5)
    h = header(caso)
    if h:
        s.sendall(h)
    s.sendall(startup())
    try:
        data = s.recv(4096)
    except socket.timeout:
        print(f"{caso}: TIMEOUT (el servidor no contesto)")
        return
    if not data:
        print(f"{caso}: el servidor cerro sin responder")
        return
    tipo = chr(data[0])
    if tipo == "E":
        # mensaje de error: campos separados por \0
        campos = data[5:].split(b"\x00")
        msg = next((c[1:].decode(errors="replace") for c in campos
                    if c[:1] in (b"M",)), "?")
        print(f"{caso}: ERROR -> {msg}")
    elif tipo == "R":
        print(f"{caso}: conexion aceptada (mensaje de autenticacion)")
    else:
        print(f"{caso}: respuesta tipo '{tipo}'")
    s.close()


if __name__ == "__main__":
    main(sys.argv[1], sys.argv[2])