#!/usr/bin/env python3
"""TCP forwarding with trusted PROXY v1 metadata, before any TLS bytes."""
import argparse
import asyncio
import ipaddress
import os

from env_file import load, env_int


async def forward(reader, writer, target_host, target_port):
    upstream = None
    try:
        peer = writer.get_extra_info('peername')
        local = writer.get_extra_info('sockname')
        family = 'TCP6' if ipaddress.ip_address(peer[0]).version == 6 else 'TCP4'
        remote, upstream = await asyncio.wait_for(asyncio.open_connection(target_host, target_port), 5)
        upstream.write(f'PROXY {family} {peer[0]} {local[0]} {peer[1]} {local[1]}\r\n'.encode('ascii'))
        await upstream.drain()
        async def copy(source, destination):
            while data := await asyncio.wait_for(source.read(65536), 60):
                destination.write(data)
                await asyncio.wait_for(destination.drain(), 10)
            if destination.can_write_eof():
                destination.write_eof()
        await asyncio.wait_for(asyncio.gather(copy(reader, upstream), copy(remote, writer)), 120)
    except (OSError, asyncio.TimeoutError):
        pass
    finally:
        if upstream:
            upstream.close()
        writer.close()


async def serve(args):
    active = set()
    allowed = [ipaddress.ip_network(item.strip(), strict=False)
               for item in (args.allow_from or '').split(',') if item.strip()]

    def permitted(peer):
        return not allowed or any(ipaddress.ip_address(peer[0]) in network for network in allowed)

    async def accept(reader, writer):
        if len(active) >= 64 or not permitted(writer.get_extra_info('peername')):
            writer.close()
            return
        task = asyncio.current_task()
        active.add(task)
        try:
            await forward(reader, writer, args.target, args.target_port)
        finally:
            active.discard(task)
    server = await asyncio.start_server(accept, args.bind, args.port, limit=65536)
    async with server:
        await server.serve_forever()


if __name__ == '__main__':
    load()
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--bind', default=os.environ.get('INGRESS_BIND', '127.0.0.1'))
    parser.add_argument('--port', type=int, default=env_int('INGRESS_PORT', 18443))
    parser.add_argument('--target', default=os.environ.get('INGRESS_TARGET', '127.0.0.1'))
    parser.add_argument('--target-port', type=int, default=env_int('INGRESS_TARGET_PORT', 19443))
    parser.add_argument('--allow-from', default=os.environ.get('INGRESS_ALLOW_FROM', ''),
                        help='comma-separated CIDRs allowed to connect (empty = any)')
    asyncio.run(serve(parser.parse_args()))