"""A small HTML entry shell using the public ASGI middleware interface.""" from pathlib import Path import re class EntryShell: """Only decorate the document; never buffer API, SSE, media or asset responses.""" def __init__(self, app, shell_path): self.app = app self.shell = Path(shell_path).read_bytes() async def __call__(self, scope, receive, send): if scope['type'] != 'http' or scope.get('method') != 'GET' or scope.get('path') != '/': await self.app(scope, receive, send) return start = None chunks = [] decorate = False async def intercept(message): nonlocal start, decorate if message['type'] == 'http.response.start': headers = dict(message.get('headers', [])) decorate = (message['status'] == 200 and headers.get(b'content-type', b'').startswith(b'text/html') and b'content-encoding' not in headers) if decorate: start = message else: await send(message) elif message['type'] == 'http.response.body' and decorate: chunks.append(message.get('body', b'')) if not message.get('more_body', False): body = re.sub(rb'(]*>)', lambda match: match[0] + self.shell, b''.join(chunks), count=1, flags=re.I) headers = [(k, v) for k, v in start['headers'] if k.lower() not in {b'content-length', b'etag', b'cache-control'}] headers += [(b'content-length', str(len(body)).encode()), (b'cache-control', b'no-store')] await send({**start, 'headers': headers}) await send({'type': 'http.response.body', 'body': body}) else: await send(message) await self.app(scope, receive, intercept)