#!/usr/bin/env python3
"""A TCP proxy in front of a plain-HTTP OpenBao that holds requests, and can refuse some.

    vaultproxy.py LISTEN_PORT UPSTREAM_PORT CONTROL_DIR

Everything a client sends is held for the number of seconds in CONTROL_DIR/delay
(none while the file is missing) before it goes on to 127.0.0.1:UPSTREAM_PORT, so
each Vault call of the gateway behind it takes that much longer, as against a
distant OpenBao. While CONTROL_DIR/refuse exists, a request that starts with its
contents (say "POST /v1/s3gw-iam/data/") never reaches OpenBao: the proxy answers
it 503 and closes the connection. Both files are read at every request, so a test
can change them while the proxy runs. Python standard library only.
"""

import asyncio
import sys

LISTEN, UPSTREAM, CONTROL = int(sys.argv[1]), int(sys.argv[2]), sys.argv[3]


def control(name, default):
    try:
        with open(f"{CONTROL}/{name}") as f:
            return f.read().strip()
    except FileNotFoundError:
        return default


async def handle(client_r, client_w):
    try:
        server_r, server_w = await asyncio.open_connection("127.0.0.1", UPSTREAM)
    except OSError:
        client_w.close()
        return

    async def upstream():
        try:
            while data := await client_r.read(65536):
                refuse = control("refuse", "")
                if refuse and data.startswith(refuse.encode()):
                    body = b'{"errors":["refused by the proof proxy"]}'
                    client_w.write(b"HTTP/1.1 503 Service Unavailable\r\nContent-Type: application/json\r\n"
                                   b"Connection: close\r\nContent-Length: %d\r\n\r\n%s" % (len(body), body))
                    await client_w.drain()
                    break
                delay = float(control("delay", "0") or 0)
                if delay:
                    await asyncio.sleep(delay)
                server_w.write(data)
                await server_w.drain()
        finally:
            server_w.close()
            client_w.close()

    async def downstream():
        try:
            while data := await server_r.read(65536):
                client_w.write(data)
                await client_w.drain()
        finally:
            client_w.close()

    await asyncio.gather(upstream(), downstream(), return_exceptions=True)


async def main():
    server = await asyncio.start_server(handle, "127.0.0.1", LISTEN)
    async with server:
        await server.serve_forever()


if __name__ == "__main__":
    asyncio.run(main())