1import asyncio
  2import logging
  3from utils import find_free_port
  4from flag import random_flag_index
  5from aiohttp import web
  6
  7closers = []
  8runners = []
  9
 10
 11async def homepage_handler(request):
 12    return web.HTTPFound("/index.html")
 13
 14
 15async def homepage(port=8001):
 16    app = web.Application()
 17    app.router.add_route("*", "/", homepage_handler)
 18    app.router.add_static(prefix="/", path="/var/www/html/srs/")
 19    runner = web.AppRunner(app)
 20    await runner.setup()
 21    site = web.TCPSite(runner, "127.0.0.1", port)
 22    logging.info(f"{site.name} starting")
 23    await site.start()
 24    runners.append(runner)
 25    return
 26
 27
 28class ClosingHandler:
 29    def __init__(self, exit_event):
 30        self.exit_event = exit_event
 31
 32    def next_port():
 33        return self.next_port
 34
 35    async def default(self, request):
 36        self.exit_event.set()
 37        self.next_port = find_free_port()
 38        data = {
 39            "now": request.url.port,
 40            "next": self.next_port,
 41            "flag-slice": random_flag_index(),
 42        }
 43        return web.json_response(data)
 44
 45
 46# single request server
 47async def srs(loop, port=None):
 48
 49    # create app handler and closing event
 50    closing_time = asyncio.Event()
 51    closing_task = asyncio.create_task(closing_time.wait())
 52    closers.append(closing_time)
 53    handler = ClosingHandler(closing_time)
 54
 55    # add handler to new application runner
 56    app = web.Application()
 57    app.add_routes([web.get("/", handler.default)])
 58    runner = web.AppRunner(app)
 59    await runner.setup()
 60
 61    # start app on a specific tcp port
 62    if port == None:
 63        port = find_free_port()
 64    site = web.TCPSite(runner, "0.0.0.0", port)
 65    logging.info(f"{site.name} starting")
 66    try:
 67        await site.start()
 68    except OSError:
 69        logging.warn(
 70            f"port {port} already in use, trying a different one"
 71        )
 72        await srs(loop)
 73
 74    # wait for closing event
 75    await closing_task
 76    logging.info(f"{site.name} closing")
 77    await runner.cleanup()
 78
 79    loop.create_task(srs(loop, handler.next_port))
 80
 81
 82if __name__ == "__main__":
 83    logging.basicConfig(level=logging.DEBUG)
 84
 85    loop = asyncio.get_event_loop()
 86
 87    loop.create_task(homepage())
 88
 89    for i in range(100):
 90        loop.create_task(srs(loop))
 91
 92    try:
 93        loop.run_forever()
 94    except:
 95        logging.info("\nexiting...")
 96    finally:
 97        for c in closers:
 98            c.set()  # trigger close event which runs runner.cleanup()
 99        for r in runners:
100            loop.create_task(r.cleanup())