import functools
import asyncio


def connection_timeout(f):
    """Decorator to timeout an asyncio TCP connection handler after 60 seconds."""

    @functools.wraps(f)
    async def wrapper(*args, **kwds):
        try:
            await asyncio.wait_for(f(*args, **kwds), 60)
        except asyncio.TimeoutError:
            writer = args[-1]
            writer.write("\nconnection timed out.\n".encode())
            await writer.drain()
            writer.close()

    return wrapper