Python aiohttp.web.middleware() Examples

def error_pages(overrides):
    @web.middleware
    async def middleware(request, handler):
        try:
            response = await handler(request)
            override = overrides.get(response.status)
            if override is None:
                return response
            else:
                return await override(request, response)
        except web.HTTPException as ex:
            override = overrides.get(ex.status)
            if override is None:
                raise
            else:
                return await override(request, ex)

    return middleware 

def middleware(self, request, handler):
        try:
            overrides = {}
            headers = request.headers

            forwarded_for = self.get_forwarded_for(headers)
            if forwarded_for:
                overrides['remote'] = str(forwarded_for[-self._num])

            proto = self.get_forwarded_proto(headers)
            if proto:
                overrides['scheme'] = proto[-self._num]

            host = self.get_forwarded_host(headers)
            if host is not None:
                overrides['host'] = host

            request = request.clone(**overrides)

            return await handler(request)
        except RemoteError as exc:
            exc.log(request)
            await self.raise_error(request) 

def setup(self, app):
        if self._client is not None:  # pragma: no branch
            client = self._client
        else:
            client = aiohttp.ClientSession()  # pragma: no cover
        try:
            async with client.get(
                    'https://www.cloudflare.com/ips-v4') as response:
                self._ip_networks |= self._parse_mask(await response.text())
            async with client.get(
                    'https://www.cloudflare.com/ips-v6') as response:
                self._ip_networks |= self._parse_mask(await response.text())
        finally:
            if self._client is None:  # pragma: no cover
                await client.close()

        if not self._ip_networks:
            raise RuntimeError("No networks are available")
        app.middlewares.append(self.middleware) 

def middleware(self, request, handler):
        overrides = {}

        for elem in reversed(request.forwarded[-self._num:]):
            for_ = elem.get('for')
            if for_:
                overrides['remote'] = for_
            proto = elem.get('proto')
            if proto is not None:
                overrides['scheme'] = proto
            host = elem.get('host')
            if host is not None:
                overrides['host'] = host

        request = request.clone(**overrides)
        return await handler(request) 

def create_error_middleware(self, overrides):

        @web.middleware
        async def error_middleware(request, handler):
            try:
                response = await handler(request)
                status = response.status
                override = overrides.get(status)
                if override:
                    response = await override(request)
                    response.headers.update(self.headers)
                    response.set_status(status)
                    return response
                return response
            except web.HTTPException as ex:
                override = overrides.get(ex.status)
                if override:
                    return await override(request)
                raise

        return error_middleware 

def error_pages_handler(self, overrides):
        @web.middleware
        async def middleware(request, handler):
            try:
                response = await handler(request)
                override = overrides.get(response.status)
                if override is None:
                    return response
                else:
                    return await override(request, response)
            except aiohttp.web.HTTPException as ex:
                override = overrides.get(ex.status)
                if override is None:
                    raise
                else:
                    return await override(request, ex)

        return middleware 

def get_session(request):
    session = request.get(SESSION_KEY)
    if session is None:
        storage = request.get(STORAGE_KEY)
        if storage is None:
            raise RuntimeError(
                "Install aiohttp_session middleware "
                "in your aiohttp.web.Application")
        else:
            session = await storage.load_session(request)
            if not isinstance(session, Session):
                raise RuntimeError(
                    "Installed {!r} storage should return session instance "
                    "on .load_session() call, got {!r}.".format(storage,
                                                                session))
            request[SESSION_KEY] = session
    return session 

def validation_middleware(request: web.Request, handler) -> web.Response:
    """
    Validation middleware for aiohttp web app

    Usage:

    .. code-block:: python

        app.middlewares.append(validation_middleware)


    """
    orig_handler = request.match_info.handler
    if not hasattr(orig_handler, "__schemas__"):
        if not issubclass_py37fix(orig_handler, web.View):
            return await handler(request)
        sub_handler = getattr(orig_handler, request.method.lower(), None)
        if sub_handler is None:
            return await handler(request)
        if not hasattr(sub_handler, "__schemas__"):
            return await handler(request)
        schemas = sub_handler.__schemas__
    else:
        schemas = orig_handler.__schemas__
    result = {}
    for schema in schemas:
        data = await request.app["_apispec_parser"].parse(
            schema["schema"], request, locations=schema["locations"]
        )
        if schema["put_into"]:
            request[schema["put_into"]] = data
        elif data:
            try:
                result.update(data)
            except (ValueError, TypeError):
                result = data
                break
    request[request.app["_apispec_request_data_name"]] = result
    return await handler(request) 

def file_address(self):
        """Address to the file being served."""
        return f"http://{self._address}:{self._port}/{self.path.name}"

    # This middleware makes sure only the specified file is accessible. This is needed
    # since aiohttp only supports serving an entire directory. 

def middleware(self, request, handler):
        if request.path not in self._white_paths:
            auth_header = request.headers.get(hdrs.AUTHORIZATION)

            if auth_header is None or not auth_header.startswith('Basic '):
                return await self.raise_error(request)

            try:
                secret = auth_header[6:].encode('utf-8')

                auth_decoded = base64.decodebytes(secret).decode('utf-8')
            except (UnicodeDecodeError, UnicodeEncodeError,
                    binascii.Error):
                await self.raise_error(request)

            credentials = auth_decoded.split(':')

            if len(credentials) != 2:
                await self.raise_error(request)

            username, password = credentials

            if username != self._username or password != self._password:
                await self.raise_error(request)

        return await handler(request) 

def setup(self, app):
        app.middlewares.append(self.middleware) 

def middleware(self, request, handler):
        if (
            request.path not in self._white_paths and
            request.host not in self._allowed_hosts
        ):
            await self.raise_error(request)

        return await handler(request) 

def middleware(self, request, handler):
        token = request.headers.get('Authorization')
        if token and token.startswith('Bearer'):
            token = token[7:]
        else:
            token = request.rel_url.query.get('token')
            if not token:
                token = request.headers.get('token')
        request.verified = False

        if token:
            try:
                payload = self.decode(token)
                request.verified = True
            except jwt.JWTError:
                raise web.HTTPUnauthorized()
        else:
            payload = {}
        request.session = payload
        return await handler(request) 

def middleware(self, request, handler):
        remote_ip = ip_address(request.remote)

        for network in self._ip_networks:
            if remote_ip in network:
                request = request.clone(
                    remote=request.headers['CF-CONNECTING-IP'])
                return await handler(request)

        msg = "Not cloudflare: %(remote_ip)s"
        context = {'remote_ip': remote_ip}
        logger.error(msg, context)

        await self.raise_error(request) 

def session_middleware(request, handler):
    """Wrapper to Session Middleware factory.
    """
    # Do the trick, by passing app & handler back to original session
    # middleware factory. Do not forget to await on results here as original
    # session middleware factory is also awaitable.
    app = request.app
    storage = RedisStorage(app['redis'], httponly=False)
    middleware = session_middleware_(storage)
    return await middleware(request, handler) 

def middleware(self, request, handler):
        if request.path in self._white_paths:
            return await handler(request)
        try:
            overrides = {}
            headers = request.headers

            forwarded_for = self.get_forwarded_for(headers)
            peer_ip, *_ = request.transport.get_extra_info('peername')
            ips = [ip_address(peer_ip)] + list(reversed(forwarded_for))
            ip = remote_ip(self._trusted, ips)
            overrides['remote'] = str(ip)

            proto = self.get_forwarded_proto(headers)
            if proto:
                if len(proto) > len(self._trusted):
                    raise IncorrectProtoCount(len(self._trusted), proto)
                overrides['scheme'] = proto[0]

            host = self.get_forwarded_host(headers)
            if host is not None:
                overrides['host'] = host

            request = request.clone(**overrides)

            return await handler(request)

        except RemoteError as exc:
            exc.log(request)
            await self.raise_error(request) 

def setup(self, app):
        app.middlewares.append(self.middleware) 

def ratelimit(get_storage):
    @web.middleware
    async def ratelimit_middleware(request, handler):
        storage = get_storage(handler)
        if not hasattr(storage, "_ratelimit"):
            storage.setdefault("ratelimit", collections.defaultdict(lambda: 0))
            storage.setdefault("ratelimit_last", collections.defaultdict(lambda: 1))
            storage.setdefault("last_request", collections.defaultdict(lambda: 0))
        if storage["last_request"][request.remote] > time.time() - 30:
            # Maybe ratelimit, was requested within 30 seconds
            last = storage["ratelimit_last"][request.remote]
            storage["ratelimit_last"][request.remote] = storage["ratelimit"][request.remote]
            storage["ratelimit"][request.remote] += last
            if storage["ratelimit"][request.remote] > 50:
                # If they have to wait more than 5 seconds (10 requests), kill em.
                return web.Response(status=429)
            await asyncio.sleep(storage["ratelimit"][request.remote] / 10)
        else:
            try:
                del storage["ratelimit"][request.remote]
                del storage["ratelimit_last"][request.remote]
            except KeyError:
                pass
        storage["last_request"][request.remote] = time.time()
        return await handler(request)
    return ratelimit_middleware 

def setup(self, app):
        app.on_response_prepare.append(self.on_response_prepare)
        app.middlewares.append(self.middleware) 

def middleware(self, request, handler):
        if request.path in self._white_paths:
            return await handler(request)
        try:
            overrides = {}

            forwarded = request.forwarded
            if len(self._trusted) != len(forwarded):
                raise IncorrectForwardedCount(len(self._trusted),
                                              len(forwarded))

            peer_ip, *_ = request.transport.get_extra_info('peername')
            ips = [ip_address(peer_ip)]

            for elem in reversed(request.forwarded):
                for_ = elem.get('for')
                if for_:
                    ips.append(ip_address(for_))
                proto = elem.get('proto')
                if proto is not None:
                    overrides['scheme'] = proto
                host = elem.get('host')
                if host is not None:
                    overrides['host'] = host

                overrides['remote'] = str(remote_ip(self._trusted, ips))

            request = request.clone(**overrides)
            return await handler(request)
        except RemoteError as exc:
            exc.log(request)
            await self.raise_error(request) 

def setup(self, app):
        app.middlewares.append(self.middleware) 

def setup(self, app):
        app.middlewares.append(self.middleware) 

def sentry_middleware_factory(sentry_kwargs=None):
    """
    Aiohttp Middleware for logging error messages automatically on Sentry
    """
    if sentry_kwargs is None:
        sentry_kwargs = {}

    sentry_kwargs = {
        'dsn': sentry_option("dsn"),
        'transport': AioHttpTransport,
        'enable_breadcrumbs': False,
        **sentry_kwargs,
    }

    # Create a new raven client with the specified configuration
    _client = raven.Client(**sentry_kwargs)

    # Set the tags context with an application ID for all reported events
    _client.tags_context({"app_id": APP_NAME})

    @web.middleware
    async def impl(request: web.Request, handler: Callable) -> web.Response:
        try:
            return await handler(request)
        except:
            _client.captureException(data={
                'request': {
                    'query_string': request.query_string,
                    'cookies': request.headers.get('Cookie', ''),
                    'headers':  dict(request.headers),
                    'url': request.path,
                    'method': request.method,
                    'env': {
                        'REMOTE_ADDR': request.transport.get_extra_info('peername')[0],
                    }
                }
            })
            raise

    return impl 

def sentry_middleware_factory(sentry_kwargs=None):
    """
    Aiohttp Middleware for logging error messages automatically on Sentry
    """
    if sentry_kwargs is None:
        sentry_kwargs = {}

    sentry_kwargs = {
        'dsn': sentry_option("dsn"),
        'transport': AioHttpTransport,
        'enable_breadcrumbs': False,
        **sentry_kwargs,
    }

    # Create a new raven client with the specified configuration
    _client = raven.Client(**sentry_kwargs)

    # Set the tags context with an application ID for all reported events
    _client.tags_context({"app_id": APP_NAME})

    @web.middleware
    async def impl(request: web.Request, handler: Callable) -> web.Response:
        try:
            return await handler(request)
        except:
            _client.captureException(data={
                'request': {
                    'query_string': request.query_string,
                    'cookies': request.headers.get('Cookie', ''),
                    'headers':  dict(request.headers),
                    'url': request.path,
                    'method': request.method,
                    'env': {
                        'REMOTE_ADDR': request.transport.get_extra_info('peername')[0],
                    }
                }
            })
            raise

    return impl 

def test_validation_exception_middleware(swagger_docs, aiohttp_client):
    @web.middleware
    async def middleware(request, handler):
        try:
            return await handler(request)
        except RequestValidationFailed as exc:
            assert exc.errors == {"query": "value should be type of int"}
            raise exc

    async def handler(request):
        """
        ---
        parameters:

          - name: query
            in: query
            required: true
            schema:
              type: integer

        responses:
          '200':
            description: OK.
        """
        return web.json_response()

    swagger = swagger_docs()
    swagger._app.middlewares.append(middleware)
    swagger.add_get("/r", handler)

    client = await aiohttp_client(swagger._app)

    params = {"query": "abc"}
    resp = await client.get("/r", params=params)
    assert resp.status == 400 

def session_middleware(storage):

    if not isinstance(storage, AbstractStorage):
        raise RuntimeError("Expected AbstractStorage got {}".format(storage))

    @web.middleware
    async def factory(request, handler):
        request[STORAGE_KEY] = storage
        raise_response = False
        try:
            response = await handler(request)
        except web.HTTPException as exc:
            response = exc
            raise_response = True
        if not isinstance(response, web.StreamResponse):
            raise RuntimeError(
                "Expect response, not {!r}".format(type(response)))
        if not isinstance(response, web.Response):
            # likely got websocket or streaming
            return response
        if response.prepared:
            raise RuntimeError(
                "Cannot save session data into prepared response")
        session = request.get(SESSION_KEY)
        if session is not None:
            if session._changed:
                await storage.save_session(request, response, session)
        if raise_response:
            raise response
        return response

    return factory 

def new_session(request):
    storage = request.get(STORAGE_KEY)
    if storage is None:
        raise RuntimeError(
            "Install aiohttp_session middleware "
            "in your aiohttp.web.Application")
    else:
        session = await storage.new_session()
        if not isinstance(session, Session):
            raise RuntimeError(
                "Installed {!r} storage should return session instance "
                "on .load_session() call, got {!r}.".format(storage, session))
        request[SESSION_KEY] = session
    return session 

def create_app(default_cors_options: CORSOptions) -> Tuple[web.Application, Iterable[WebMiddleware]]:
    app = web.Application()
    app['api_versions'] = (1, 2, 3, 4)
    app.on_startup.append(init)
    app.on_shutdown.append(shutdown)
    # middleware must be wrapped by web.middleware at the outermost level.
    return app, [web.middleware(apartial(rlim_middleware, app))] 

def rlim_middleware(app: web.Application,
                          request: web.Request,
                          handler: WebRequestHandler) -> web.StreamResponse:
    # This is a global middleware: request.app is the root app.
    now = Decimal(time.time()).quantize(_time_prec)
    rr = app['redis_rlim']
    if request['is_authorized']:
        rate_limit = request['keypair']['rate_limit']
        access_key = request['keypair']['access_key']
        ret = await redis.execute_script(
            rr, 'ratelimit', _rlim_script,
            [access_key],
            [str(now), str(_rlim_window)],
        )
        rolling_count = int(ret)
        if rolling_count > rate_limit:
            raise RateLimitExceeded
        remaining = rate_limit - rolling_count
        response = await handler(request)
        response.headers['X-RateLimit-Limit'] = str(rate_limit)
        response.headers['X-RateLimit-Remaining'] = str(remaining)
        response.headers['X-RateLimit-Window'] = str(_rlim_window)
        return response
    else:
        # No checks for rate limiting for non-authorized queries.
        response = await handler(request)
        response.headers['X-RateLimit-Limit'] = '1000'
        response.headers['X-RateLimit-Remaining'] = '1000'
        response.headers['X-RateLimit-Window'] = str(_rlim_window)
        return response 

def setup_management_ui(self) -> None:
        ui_base = self.config["server.ui_base_path"]
        if ui_base == "/":
            ui_base = ""
        directory = (self.config["server.override_resource_path"]
                     or pkg_resources.resource_filename("maubot", "management/frontend/build"))
        self.app.router.add_static(f"{ui_base}/static", f"{directory}/static")
        self.setup_static_root_files(directory, ui_base)

        with open(f"{directory}/index.html", "r") as file:
            index_html = file.read()

        @web.middleware
        async def frontend_404_middleware(request: web.Request, handler) -> web.Response:
            if hasattr(handler, "__self__") and isinstance(handler.__self__, web.StaticResource):
                try:
                    return await handler(request)
                except web.HTTPNotFound:
                    return web.Response(body=index_html, content_type="text/html")
            return await handler(request)

        async def ui_base_redirect(_: web.Request) -> web.Response:
            raise web.HTTPFound(f"{ui_base}/")

        self.app.middlewares.append(frontend_404_middleware)
        self.app.router.add_get(f"{ui_base}/", lambda _: web.Response(body=index_html,
                                                                      content_type="text/html"))
        self.app.router.add_get(ui_base, ui_base_redirect)