Spaces:
Running on Zero
Running on Zero
| """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'(<body\b[^>]*>)', 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) | |