class NetworkManager(EventEmitter):
    """NetworkManager class."""

    Events = SimpleNamespace(
        Request='request',
        Response='response',
        RequestFailed='requestfailed',
        RequestFinished='requestfinished',
    )

    def __init__(self, client: CDPSession, frameManager: FrameManager) -> None:
        """Make new NetworkManager."""
        super().__init__()
        self._client = client
        self._frameManager = frameManager
        self._requestIdToRequest: Dict[Optional[str], Request] = dict()
        self._requestIdToResponseWillBeSent: Dict[Optional[str], Dict] = dict()
        self._extraHTTPHeaders: OrderedDict[str, str] = OrderedDict()
        self._offline: bool = False
        self._credentials: Optional[Dict[str, str]] = None
        self._attemptedAuthentications: Set[Optional[str]] = set()
        self._userRequestInterceptionEnabled = False
        self._protocolRequestInterceptionEnabled = False
        self._requestHashToRequestIds = Multimap()
        self._requestHashToInterceptionIds = Multimap()

        self._client.on(
            'Network.requestWillBeSent',
            lambda event: self._client._loop.create_task(
                self._onRequestWillBeSent(event)),
        )
        self._client.on('Network.requestIntercepted',
                        self._onRequestIntercepted)  # noqa: E501
        self._client.on('Network.requestServedFromCache',
                        self._onRequestServedFromCache)  # noqa: #501
        self._client.on('Network.responseReceived', self._onResponseReceived)
        self._client.on('Network.loadingFinished', self._onLoadingFinished)
        self._client.on('Network.loadingFailed', self._onLoadingFailed)

    async def authenticate(self, credentials: Dict[str, str]) -> None:
        """Provide credentials for http auth."""
        self._credentials = credentials
        await self._updateProtocolRequestInterception()

    async def setExtraHTTPHeaders(self, extraHTTPHeaders: Dict[str,
                                                               str]) -> None:
        """Set extra http headers."""
        self._extraHTTPHeaders = OrderedDict()
        for k, v in extraHTTPHeaders.items():
            if not isinstance(v, str):
                raise TypeError(
                    f'Expected value of header "{k}" to be string, '
                    f'but {type(v)} is found.')
            self._extraHTTPHeaders[k.lower()] = v
        await self._client.send('Network.setExtraHTTPHeaders',
                                {'headers': self._extraHTTPHeaders})

    def extraHTTPHeaders(self) -> Dict[str, str]:
        """Get extra http headers."""
        return dict(**self._extraHTTPHeaders)

    async def setOfflineMode(self, value: bool) -> None:
        """Change offline mode enable/disable."""
        if self._offline == value:
            return
        self._offline = value
        await self._client.send(
            'Network.emulateNetworkConditions', {
                'offline': self._offline,
                'latency': 0,
                'downloadThroughput': -1,
                'uploadThroughput': -1,
            })

    async def setUserAgent(self, userAgent: str) -> None:
        """Set user agent."""
        await self._client.send('Network.setUserAgentOverride',
                                {'userAgent': userAgent})

    async def setRequestInterception(self, value: bool) -> None:
        """Enable request interception."""
        self._userRequestInterceptionEnabled = value
        await self._updateProtocolRequestInterception()

    async def _updateProtocolRequestInterception(self) -> None:
        enabled = (self._userRequestInterceptionEnabled
                   or bool(self._credentials))
        if enabled == self._protocolRequestInterceptionEnabled:
            return
        self._protocolRequestInterceptionEnabled = enabled
        patterns = [{'urlPattern': '*'}] if enabled else []
        await asyncio.gather(
            self._client.send(
                'Network.setCacheDisabled',
                {'cacheDisabled': enabled},
            ),
            self._client.send(
                'Network.setRequestInterception',
                {'patterns': patterns},
            ))

    async def _onRequestWillBeSent(self, event: Dict) -> None:
        if self._protocolRequestInterceptionEnabled:
            requestHash = generateRequestHash(event.get('request', {}))
            interceptionId = self._requestHashToInterceptionIds.firstValue(
                requestHash)  # noqa: E501
            if interceptionId:
                self._onRequest(event, interceptionId)
                self._requestHashToInterceptionIds.delete(
                    requestHash, interceptionId)  # noqa: E501
            else:
                self._requestHashToRequestIds.set(
                    requestHash, event.get('requestId'))  # noqa: E501
                self._requestIdToResponseWillBeSent[event.get(
                    'requestId')] = event  # noqa: E501
            return
        self._onRequest(event, None)

    async def _send(self, method: str, msg: dict) -> None:
        try:
            await self._client.send(method, msg)
        except Exception as e:
            debugError(logger, e)

    def _onRequestIntercepted(self, event: dict) -> None:  # noqa: C901
        if event.get('authChallenge'):
            response = 'Default'
            if event['interceptionId'] in self._attemptedAuthentications:
                response = 'CancelAuth'
            elif self._credentials:
                response = 'ProvideCredentials'
                self._attemptedAuthentications.add(event['interceptionId'])
            username = getattr(self, '_credentials', {}).get('username')
            password = getattr(self, '_credentials', {}).get('password')

            self._client._loop.create_task(
                self._send(
                    'Network.continueInterceptedRequest', {
                        'interceptionId': event['interceptionId'],
                        'authChallengeResponse': {
                            'response': response,
                            'username': username,
                            'password': password,
                        }
                    }))
            return

        if (not self._userRequestInterceptionEnabled
                and self._protocolRequestInterceptionEnabled):
            self._client._loop.create_task(
                self._send('Network.continueInterceptedRequest', {
                    'interceptionId': event['interceptionId'],
                }))

        requestHash = generateRequestHash(event['request'])
        requestId = self._requestHashToRequestIds.firstValue(requestHash)
        if requestId:
            requestWillBeSentEvent = self._requestIdToResponseWillBeSent[
                requestId]  # noqa: E501
            self._onRequest(requestWillBeSentEvent,
                            event.get('interceptionId'))  # noqa: E501
            self._requestHashToRequestIds.delete(requestHash, requestId)
            self._requestIdToResponseWillBeSent.pop(requestId, None)
        else:
            self._requestHashToInterceptionIds.set(
                requestHash, event['interceptionId'])  # noqa: E501

    def _onRequest(self, event: Dict, interceptionId: Optional[str]) -> None:
        redirectChain: List[Request] = list()
        if event.get('redirectResponse'):
            request = self._requestIdToRequest.get(event['requestId'])
            if request:
                redirectResponse = event['redirectResponse']
                self._handleRequestRedirect(
                    request,
                    redirectResponse.get('status'),
                    redirectResponse.get('headers'),
                    redirectResponse.get('fromDiskCache'),
                    redirectResponse.get('fromServiceWorker'),
                    redirectResponse.get('SecurityDetails'),
                )
                redirectChain = request._redirectChain

        isNavigationRequest = bool(
            event.get('requestId') == event.get('loaderId')
            and event.get('type') == 'Document')
        self._handleRequestStart(
            event['requestId'],
            interceptionId,
            event.get('request', {}).get('url'),
            isNavigationRequest,
            event.get('type', ''),
            event.get('request', {}),
            event.get('frameId'),
            redirectChain,
        )

    def _onRequestServedFromCache(self, event: Dict) -> None:
        request = self._requestIdToRequest.get(event.get('requestId'))
        if request:
            request._fromMemoryCache = True

    def _handleRequestRedirect(self,
                               request: 'Request',
                               redirectStatus: int,
                               redirectHeaders: Dict,
                               fromDiskCache: bool,
                               fromServiceWorker: bool,
                               securityDetails: Dict = None) -> None:
        response = Response(self._client, request, redirectStatus,
                            redirectHeaders, fromDiskCache, fromServiceWorker,
                            securityDetails)
        request._response = response
        request._redirectChain.append(request)
        response._bodyLoadedPromiseFulfill(
            NetworkError('Response body is unavailable for redirect response'))
        self._requestIdToRequest.pop(request._requestId, None)
        self._attemptedAuthentications.discard(request._interceptionId)
        self.emit(NetworkManager.Events.Response, response)
        self.emit(NetworkManager.Events.RequestFinished, request)

    def _handleRequestStart(self, requestId: str,
                            interceptionId: Optional[str], url: str,
                            isNavigationRequest: bool, resourceType: str,
                            requestPayload: Dict, frameId: Optional[str],
                            redirectChain: List['Request']) -> None:
        frame = None
        if frameId and self._frameManager is not None:
            frame = self._frameManager.frame(frameId)

        request = Request(self._client, requestId, interceptionId,
                          isNavigationRequest,
                          self._userRequestInterceptionEnabled, url,
                          resourceType, requestPayload, frame, redirectChain)
        self._requestIdToRequest[requestId] = request
        self.emit(NetworkManager.Events.Request, request)

    def _onResponseReceived(self, event: dict) -> None:
        request = self._requestIdToRequest.get(event['requestId'])
        # FileUpload sends a response without a matching request.
        if not request:
            return
        _resp = event.get('response', {})
        response = Response(self._client, request, _resp.get('status', 0),
                            _resp.get('headers', {}),
                            _resp.get('fromDiskCache'),
                            _resp.get('fromServiceWorker'),
                            _resp.get('securityDetails'))
        request._response = response
        self.emit(NetworkManager.Events.Response, response)

    def _onLoadingFinished(self, event: dict) -> None:
        request = self._requestIdToRequest.get(event['requestId'])
        # For certain requestIds we never receive requestWillBeSent event.
        # @see https://crbug.com/750469
        if not request:
            return
        response = request.response
        if response:
            response._bodyLoadedPromiseFulfill(None)
        self._requestIdToRequest.pop(request._requestId, None)
        self._attemptedAuthentications.discard(request._interceptionId)
        self.emit(NetworkManager.Events.RequestFinished, request)

    def _onLoadingFailed(self, event: dict) -> None:
        request = self._requestIdToRequest.get(event['requestId'])
        # For certain requestIds we never receive requestWillBeSent event.
        # @see https://crbug.com/750469
        if not request:
            return
        request._failureText = event.get('errorText')
        response = request.response
        if response:
            response._bodyLoadedPromiseFulfill(None)
        self._requestIdToRequest.pop(request._requestId, None)
        self._attemptedAuthentications.discard(request._interceptionId)
        self.emit(NetworkManager.Events.RequestFailed, request)
예제 #2
0
class NetworkManager(EventEmitter):
    """NetworkManager class."""

    Events = SimpleNamespace(
        Request='request',
        Response='response',
        RequestFailed='requestfailed',
        RequestFinished='requestfinished',
    )

    def __init__(self, client: Session) -> None:
        """Make new NetworkManager."""
        super().__init__()
        self._client = client
        self._requestIdToRequest: Dict[Optional[str], Request] = dict()
        self._interceptionIdToRequest: Dict[Optional[str], Request] = dict()
        self._extraHTTPHeaders: OrderedDict[str, str] = OrderedDict()
        self._offline: bool = False
        self._credentials: Optional[Dict[str, str]] = None
        self._attemptedAuthentications: Set[str] = set()
        self._userRequestInterceptionEnabled = False
        self._protocolRequestInterceptionEnabled = False
        self._requestHashToRequestIds = Multimap()
        self._requestHashToInterceptionIds = Multimap()

        self._client.on('Network.requestWillBeSent', self._onRequestWillBeSent)
        self._client.on('Network.requestIntercepted',
                        self._onRequestIntercepted)  # noqa: E501
        self._client.on('Network.responseReceived', self._onResponseReceived)
        self._client.on('Network.loadingFinished', self._onLoadingFinished)
        self._client.on('Network.loadingFailed', self._onLoadingFailed)

    async def authenticate(self, credentials: Dict[str, str]) -> None:
        """Provide credentials for http auth."""
        self._credentials = credentials
        await self._updateProtocolRequestInterception()

    async def setExtraHTTPHeaders(self, extraHTTPHeaders: Dict[str,
                                                               str]) -> None:
        """Set extra http headers."""
        self._extraHTTPHeaders = OrderedDict()
        for k, v in extraHTTPHeaders.items():
            if not isinstance(v, str):
                raise TypeError(f'Expected value of header {k} to be string, '
                                'but {} is found.'.format(type(v)))
            self._extraHTTPHeaders[k.lower()] = v
        await self._client.send('Network.setExtraHTTPHeaders',
                                {'headers': self._extraHTTPHeaders})

    def extraHTTPHeaders(self) -> Dict[str, str]:
        """Get extra http headers."""
        return dict(**self._extraHTTPHeaders)

    async def setOfflineMode(self, value: bool) -> None:
        """Change offline mode enable/disable."""
        if self._offline == value:
            return
        self._offline = value
        await self._client.send(
            'Network.emulateNetworkConditions', {
                'offline': self._offline,
                'latency': 0,
                'downloadThroughput': -1,
                'uploadThroughput': -1,
            })

    async def setUserAgent(self, userAgent: str) -> None:
        """Set user agent."""
        await self._client.send('Network.setUserAgentOverride',
                                {'userAgent': userAgent})

    async def setRequestInterception(self, value: bool) -> None:
        """Enable request intercetion."""
        self._userRequestInterceptionEnabled = value
        await self._updateProtocolRequestInterception()

    async def _updateProtocolRequestInterception(self) -> None:
        enabled = (self._userRequestInterceptionEnabled
                   or bool(self._credentials))
        if enabled == self._protocolRequestInterceptionEnabled:
            return
        self._protocolRequestInterceptionEnabled = enabled
        patterns = [{'urlPattern': '*'}] if enabled else []
        await asyncio.gather(
            self._client.send(
                'Network.setCacheDisabled',
                {'cacheDisabled': enabled},
            ),
            self._client.send(
                'Network.setRequestInterception',
                {'patterns': patterns},
            ))

    def _onRequestIntercepted(self, event: dict) -> None:  # noqa: C901
        if event.get('authChallenge'):
            response = 'Default'
            if event['interceptionId'] in self._attemptedAuthentications:
                response = 'CancelAuth'
            elif self._credentials:
                response = 'ProvideCredentials'
                self._attemptedAuthentications.add(event['interceptionId'])
            username = getattr(self, '_credentials', {}).get('username')
            password = getattr(self, '_credentials', {}).get('password')
            asyncio.ensure_future(
                self._client.send(
                    'Network.continueInterceptedRequest', {
                        'interceptionId': event['interceptionId'],
                        'authChallengeResponse': {
                            'response': response,
                            'username': username,
                            'password': password,
                        }
                    }))
            return

        if (not self._userRequestInterceptionEnabled
                and self._protocolRequestInterceptionEnabled):
            asyncio.ensure_future(
                self._client.send('Network.continueInterceptedRequest', {
                    'interceptionId': event['interceptionId'],
                }))

        if 'redirectStatusCode' in event:
            request = self._interceptionIdToRequest.get(
                event.get('interceptionId', ''))
            if not request:
                raise NetworkError('INTERNAL ERROR: failed to find request '
                                   'for interception redirect.')
            self._handleRequestRedirect(request,
                                        event.get('redirectStatusCode', 0),
                                        event.get('redirectHeaders', {}))
            self._handleRequestStart(request._requestId,
                                     event.get('interceptionId', ''),
                                     event.get('redirectUrl', ''),
                                     event.get('resourceType', ''),
                                     event.get('request', {}))
            return

        requestHash = generateRequestHash(event['request'])
        requestId = self._requestHashToRequestIds.firstValue(requestHash)
        if requestId:
            self._requestHashToRequestIds.delete(requestHash, requestId)
            self._handleRequestStart(requestId, event['interceptionId'],
                                     event['request']['url'],
                                     event['resourceType'], event['request'])
        else:
            self._requestHashToInterceptionIds.set(requestHash,
                                                   event['interceptionId'])
            self._handleRequestStart(None, event['interceptionId'],
                                     event['request']['url'],
                                     event['resourceType'], event['request'])

    def _handleRequestRedirect(self, request: 'Request', redirectStatus: int,
                               redirectHeaders: Dict) -> None:
        response = Response(self._client, request, redirectStatus,
                            redirectHeaders)
        request._response = response
        self._requestIdToRequest.pop(request._requestId, None)
        self._interceptionIdToRequest.pop(request._interceptionId, None)
        self._attemptedAuthentications.discard(request._interceptionId)
        self.emit(NetworkManager.Events.Response, response)
        self.emit(NetworkManager.Events.RequestFinished, request)

    def _handleRequestStart(self, requestId: Optional[str],
                            interceptionId: str, url: str, resourceType: str,
                            requestPayload: Dict) -> None:
        request = Request(self._client, requestId, interceptionId,
                          self._userRequestInterceptionEnabled, url,
                          resourceType, requestPayload)
        if requestId:
            self._requestIdToRequest[requestId] = request
        if interceptionId:
            self._interceptionIdToRequest[interceptionId] = request
        self.emit(NetworkManager.Events.Request, request)

    def _onRequestWillBeSent(self, event: dict) -> None:
        if self._protocolRequestInterceptionEnabled:
            if event.get('redirectResponse'):
                return
            requestHash = generateRequestHash(event['request'])
            interceptionId = self._requestHashToInterceptionIds.firstValue(
                requestHash)
            request = self._interceptionIdToRequest.get(interceptionId)
            if request:
                request._requestId = event['requestId']
                self._requestIdToRequest[event['requestId']] = request
                self._requestHashToInterceptionIds.delete(
                    requestHash, interceptionId)
            else:
                self._requestHashToRequestIds.set(requestHash,
                                                  event['requestId'])
            return
        if event.get('redirectResponse'):
            request = self._requestIdToRequest[event['requestId']]
            redirectResponse = event.get('redirectResponse', {})
            self._handleRequestRedirect(
                request,
                redirectResponse.get('status'),
                redirectResponse.get('headers'),
            )
        self._handleRequestStart(
            event.get('requestId', ''),
            '',
            event.get('request', {}).get('url', ''),
            event.get('type', ''),
            event.get('request', {}),
        )

    def _onResponseReceived(self, event: dict) -> None:
        request = self._requestIdToRequest.get(event['requestId'])
        # FileUpload sends a response without a matching request.
        if not request:
            return
        _resp = event.get('response', {})
        response = Response(self._client, request, _resp.get('status', 0),
                            _resp.get('headers', {}))
        request._response = response
        self.emit(NetworkManager.Events.Response, response)

    def _onLoadingFinished(self, event: dict) -> None:
        request = self._requestIdToRequest.get(event.get('requestId', ''))
        # For certain requestIds we never receive requestWillBeSent event.
        # @see https://crbug.com/750469
        if not request:
            return
        request._completePromiseFulfill()
        self._requestIdToRequest.pop(request._requestId, None)
        self._interceptionIdToRequest.pop(request._interceptionId, None)
        self._attemptedAuthentications.discard(request._interceptionId)
        self.emit(NetworkManager.Events.RequestFinished, request)

    def _onLoadingFailed(self, event: dict) -> None:
        request = self._requestIdToRequest.get(event['requestId'])
        # For certain requestIds we never receive requestWillBeSent event.
        # @see https://crbug.com/750469
        if not request:
            return
        request._failureText = event.get('errorText')
        request._completePromiseFulfill()
        self._requestIdToRequest.pop(request._requestId, None)
        self._interceptionIdToRequest.pop(request._interceptionId, None)
        self._attemptedAuthentications.discard(request._interceptionId)
        self.emit(NetworkManager.Events.RequestFailed, request)
예제 #3
0
class NetworkManager(EventEmitter):

    Events = {
        'Request': 'request',
        'Response': 'response',
        'RequestFailed': 'requestfailed',
        'RequestFinished': 'requestfinished',
    }

    def __init__(self, client):
        super().__init__()
        self._client = client
        self._request_id_to_request = {}
        self._interception_id_to_request = {}
        self._extra_http_headers = {}

        self._request_interception_enabled = False
        self._request_hash_to_request_ids = Multimap()
        self._request_hash_to_interceptions = Multimap()

        self._client.on('Network.requestWillBeSent',
                        self._on_request_will_be_sent)
        self._client.on('Network.requestIntercepted',
                        self._on_request_intercepted)
        self._client.on('Network.responseReceived', self._on_response_received)
        self._client.on('Network.loadingFinished', self._on_loading_finished)
        self._client.on('Network.loadingFailed', self._on_loading_failed)

    async def set_extra_http_headers(self, extra_http_headers):
        self._extra_http_headers = extra_http_headers
        await self._client.send('Network.setExtraHTTPHeaders',
                                extra_http_headers)

    def extra_http_headers(self):
        return self._extra_http_headers

    async def set_user_agent(self, user_agent):
        return await self._client.send('Network.setUserAgentOverride',
                                       {'userAgent': user_agent})

    async def set_request_interception_enabled(self, value):
        await self._client.send('Network.setRequestInterceptionEnabled',
                                {'enabled': not not value})
        self._request_interception_enabled = value

    def _on_request_intercepted(self, event):
        event['request']['url'] = remove_url_hash(event['request']['url'])

        if 'redirectStatusCode' in event and event['redirectStatusCode']:
            request = self._interception_id_to_request.get(
                event['interceptionId'], None)
            if not request:
                raise Exception('INTERNAL ERROR: failed to find '
                                'request for interception redirect.')
            self._handle_request_redirect(request, event['redirectStatusCode'],
                                          event['redirectHeaders'])
            self._handle_request_start(request['_requestId'],
                                       event['interceptionId'],
                                       event['redirectUrl'], event['request'])
            return
        request_hash = generate_request_hash(event['request'])
        self._request_hash_to_interceptions.set(request_hash, event)
        self._maybe_resolve_interception(request_hash)

    def _handle_request_redirect(self, request, redirect_status,
                                 redirect_headers):
        response = Response(self._client, request, redirect_status,
                            redirect_headers)
        request._response = response
        try:
            del self._request_id_to_request[request._request_id]
            del self._interception_id_to_request[request._interception_id]
        except KeyError:
            pass
        self.emit(NetworkManager.Events['Response'], response)
        self.emit(NetworkManager.Events['RequestFinished'], request)

    def _handle_request_start(self, request_id, interception_id, url,
                              request_payload):
        request = Request(self._client, request_id, interception_id, url,
                          request_payload)
        self._request_id_to_request[request_id] = request
        self._interception_id_to_request[interception_id] = request
        self.emit(NetworkManager.Events['Request'], request)

    def _on_request_will_be_sent(self, event):
        if self._request_interception_enabled and \
                not event['request']['url'].startswith('data:'):
            if 'redirectResponse' in event and event['redirectResponse']:
                return
            request_hash = generate_request_hash(event['request'])
            self._request_hash_to_request_ids.set(request_hash,
                                                  event['requestId'])
            self._maybe_resolve_interception(request_hash)
            return
        if 'redirectResponse' in event and event['redirectResponse']:
            request = self._request_id_to_request.get(event['requestId'])
            self._handle_request_redirect(request,
                                          event['redirectResponse']['status'],
                                          event['redirectResponse']['headers'])
        self._handle_request_start(event['requestId'], None,
                                   event['request']['url'], event['request'])

    def _maybe_resolve_interception(self, request_hash):
        request_id = self._request_hash_to_request_ids.first_value(
            request_hash)
        interception = self._request_hash_to_interceptions.first_value(
            request_hash)
        if not request_id or not interception:
            return
        self._request_hash_to_request_ids.delete(request_hash, request_id)
        self._request_hash_to_interceptions.delete(request_hash, interception)
        self._handle_request_start(request_id, interception['interceptionId'],
                                   interception['request']['url'],
                                   interception['request'])

    def _on_response_received(self, event):
        request = self._request_id_to_request.get(event['requestId'], {})
        if not request:
            return
        response = Response(self._client, request, event['response']['status'],
                            event['response']['headers'])
        request._response = response
        self.emit(NetworkManager.Events['Response'], response)

    def _on_loading_finished(self, event):
        request = self._request_id_to_request.get(event['requestId'], None)
        if not request:
            return
        request._complete_promise.set_result(None)
        try:
            del self._request_id_to_request[event['requestId']]
            del self._interception_id_to_request[event['interceptionId']]
        except KeyError:
            pass
        self.emit(NetworkManager.Events['RequestFinished'], request)

    def _on_loading_failed(self, event):
        request = self._request_id_to_request.get(event['requestId'], None)
        if not request:
            return
        request._complete_promise.set_result(None)
        try:
            del self._request_id_to_request[event['requestId']]
            del self._interception_id_to_request[event['interceptionId']]
        except KeyError:
            pass
        self.emit(NetworkManager.Events['RequestFailed'], request)
예제 #4
0
class NetworkManager(EventEmitter):
    'NetworkManager class.'
    Events = SimpleNamespace(Request='request', Response='response',
                             RequestFailed='requestfailed', RequestFinished='requestfinished')

    def __init__(self, client: CDPSession, frameManager: FrameManager) -> None:
        'Make new NetworkManager.'
        super().__init__()
        self._client = client
        self._frameManager = frameManager
        self._requestIdToRequest = dict()
        self._interceptionIdToRequest = dict()
        self._extraHTTPHeaders = OrderedDict()
        self._offline = False
        self._credentials = None
        self._attemptedAuthentications = set()
        self._userRequestInterceptionEnabled = False
        self._protocolRequestInterceptionEnabled = False
        self._requestHashToRequestIds = Multimap()
        self._requestHashToInterceptionIds = Multimap()
        self._client.on('Network.requestWillBeSent', self._onRequestWillBeSent)
        self._client.on('Network.requestIntercepted',
                        self._onRequestIntercepted)
        self._client.on('Network.responseReceived', self._onResponseReceived)
        self._client.on('Network.loadingFinished', self._onLoadingFinished)
        self._client.on('Network.loadingFailed', self._onLoadingFailed)

    async def authenticate(self, credentials: Dict[(str, str)]) -> None:
        'Provide credentials for http auth.'
        self._credentials = credentials
        (await self._updateProtocolRequestInterception())

    async def setExtraHTTPHeaders(self, extraHTTPHeaders: Dict[(str, str)]) -> None:
        'Set extra http headers.'
        self._extraHTTPHeaders = OrderedDict()
        for (k, v) in extraHTTPHeaders.items():
            if (not isinstance(v, str)):
                raise TypeError(''.join(['Expected value of header "', '{}'.format(
                    k), '" to be string, but {} is found.']).format(type(v)))
            self._extraHTTPHeaders[k.lower()] = v
        (await self._client.send('Network.setExtraHTTPHeaders', {
            'headers': self._extraHTTPHeaders,
        }))

    def extraHTTPHeaders(self) -> Dict[(str, str)]:
        'Get extra http headers.'
        return dict(**self._extraHTTPHeaders)

    async def setOfflineMode(self, value: bool) -> None:
        'Change offline mode enable/disable.'
        if (self._offline == value):
            return
        self._offline = value
        (await self._client.send('Network.emulateNetworkConditions', {
            'offline': self._offline,
            'latency': 0,
            'downloadThroughput': (- 1),
            'uploadThroughput': (- 1),
        }))

    async def setUserAgent(self, userAgent: str) -> None:
        'Set user agent.'
        (await self._client.send('Network.setUserAgentOverride', {
            'userAgent': userAgent,
        }))

    async def setRequestInterception(self, value: bool) -> None:
        'Enable request intercetion.'
        self._userRequestInterceptionEnabled = value
        (await self._updateProtocolRequestInterception())

    async def _updateProtocolRequestInterception(self) -> None:
        enabled = (self._userRequestInterceptionEnabled or bool(
            self._credentials))
        if (enabled == self._protocolRequestInterceptionEnabled):
            return
        self._protocolRequestInterceptionEnabled = enabled
        patterns = ([{
            'urlPattern': '*',
        }] if enabled else [])
        (await asyncio.gather(self._client.send('Network.setCacheDisabled', {
            'cacheDisabled': enabled,
        }), self._client.send('Network.setRequestInterception', {
            'patterns': patterns,
        })))

    def _onRequestIntercepted(self, event: dict) -> None:
        if event.get('authChallenge'):
            response = 'Default'
            if (event['interceptionId'] in self._attemptedAuthentications):
                response = 'CancelAuth'
            elif self._credentials:
                response = 'ProvideCredentials'
                self._attemptedAuthentications.add(event['interceptionId'])
            username = getattr(self, '_credentials', {

            }).get('username')
            password = getattr(self, '_credentials', {

            }).get('password')
            asyncio.ensure_future(self._client.send('Network.continueInterceptedRequest', {
                'interceptionId': event['interceptionId'],
                'authChallengeResponse': {
                    'response': response,
                    'username': username,
                    'password': password,
                },
            }))
            return
        if ((not self._userRequestInterceptionEnabled) and self._protocolRequestInterceptionEnabled):
            asyncio.ensure_future(self._client.send('Network.continueInterceptedRequest', {
                'interceptionId': event['interceptionId'],
            }))
        if ('redirectURL' in event):
            request = self._interceptionIdToRequest.get(
                event.get('interceptionId', ''))
            if request:
                self._handleRequestRedirect(request, event.get('redirectStatusCode', 0), event.get('redirectHeaders', {

                }))
                self._handleRequestStart(request._requestId, event.get('interceptionId', ''), event.get('redirectUrl', ''), event.get('resourceType', ''), event.get('request', {

                }), event.get('frameId'))
            return
        requestHash = generateRequestHash(event['request'])
        requestId = self._requestHashToRequestIds.firstValue(requestHash)
        if requestId:
            self._requestHashToRequestIds.delete(requestHash, requestId)
            self._handleRequestStart(requestId, event['interceptionId'], event['request']
                                     ['url'], event['resourceType'], event['request'], event['frameId'])
        else:
            self._requestHashToInterceptionIds.set(
                requestHash, event['interceptionId'])
            self._handleRequestStart(None, event['interceptionId'], event['request']
                                     ['url'], event['resourceType'], event['request'], event['frameId'])

    def _handleRequestRedirect(self, request: 'Request', redirectStatus: int, redirectHeaders: Dict) -> None:
        response = Response(self._client, request,
                            redirectStatus, redirectHeaders)
        request._response = response
        self._requestIdToRequest.pop(request._requestId, None)
        self._interceptionIdToRequest.pop(request._interceptionId, None)
        self._attemptedAuthentications.discard(request._interceptionId)
        self.emit(NetworkManager.Events.Response, response)
        self.emit(NetworkManager.Events.RequestFinished, request)

    def _handleRequestStart(self, requestId: Optional[str], interceptionId: str, url: str, resourceType: str, requestPayload: Dict, frameId: Optional[str]) -> None:
        frame = None
        if (frameId and (self._frameManager is not None)):
            frame = self._frameManager.frame(frameId)
        request = Request(self._client, requestId, interceptionId,
                          self._userRequestInterceptionEnabled, url, resourceType, requestPayload, frame)
        if requestId:
            self._requestIdToRequest[requestId] = request
        if interceptionId:
            self._interceptionIdToRequest[interceptionId] = request
        self.emit(NetworkManager.Events.Request, request)

    def _onRequestWillBeSent(self, event: dict) -> None:
        if self._protocolRequestInterceptionEnabled:
            if event.get('redirectResponse'):
                return
            requestHash = generateRequestHash(event['request'])
            interceptionId = self._requestHashToInterceptionIds.firstValue(
                requestHash)
            request = self._interceptionIdToRequest.get(interceptionId)
            if request:
                request._requestId = event['requestId']
                self._requestIdToRequest[event['requestId']] = request
                self._requestHashToInterceptionIds.delete(
                    requestHash, interceptionId)
            else:
                self._requestHashToRequestIds.set(
                    requestHash, event['requestId'])
            return
        if event.get('redirectResponse'):
            request = self._requestIdToRequest[event['requestId']]
            if request:
                redirectResponse = event.get('redirectResponse', {

                })
                self._handleRequestRedirect(request, redirectResponse.get(
                    'status'), redirectResponse.get('headers'))
        self._handleRequestStart(event.get('requestId', ''), '', event.get('request', {

        }).get('url', ''), event.get('type', ''), event.get('request', {

        }), event.get('frameId'))

    def _onResponseReceived(self, event: dict) -> None:
        request = self._requestIdToRequest.get(event['requestId'])
        if (not request):
            return
        _resp = event.get('response', {

        })
        response = Response(self._client, request, _resp.get('status', 0), _resp.get('headers', {

        }))
        request._response = response
        self.emit(NetworkManager.Events.Response, response)

    def _onLoadingFinished(self, event: dict) -> None:
        request = self._requestIdToRequest.get(event.get('requestId', ''))
        if (not request):
            return
        request._completePromiseFulfill()
        self._requestIdToRequest.pop(request._requestId, None)
        self._interceptionIdToRequest.pop(request._interceptionId, None)
        self._attemptedAuthentications.discard(request._interceptionId)
        self.emit(NetworkManager.Events.RequestFinished, request)

    def _onLoadingFailed(self, event: dict) -> None:
        request = self._requestIdToRequest.get(event['requestId'])
        if (not request):
            return
        request._failureText = event.get('errorText')
        request._completePromiseFulfill()
        self._requestIdToRequest.pop(request._requestId, None)
        self._interceptionIdToRequest.pop(request._interceptionId, None)
        self._attemptedAuthentications.discard(request._interceptionId)
        self.emit(NetworkManager.Events.RequestFailed, request)
예제 #5
0
class NetworkManager(EventEmitter):
    """NetworkManager class."""

    Events = SimpleNamespace(
        Request='request',
        Response='response',
        RequestFailed='requestfailed',
        RequestFinished='requestfinished',
    )

    def __init__(self, client: Session) -> None:
        """Make new NetworkManager."""
        super().__init__()
        self._client = client
        self._requestIdToRequest: Dict[str, Request] = dict()
        self._interceptionIdToRequest: Dict[str, Request] = dict()
        self._extraHTTPHeaders: OrderedDict[str, str] = OrderedDict()
        self._requestInterceptionEnabled = False
        self._requestHashToRequestIds = Multimap()
        self._requestHashToInterceptions = Multimap()

        self._client.on('Network.requestWillBeSent', self._onRequestWillBeSent)
        self._client.on('Network.requestIntercepted',
                        self._onRequestIntercepted)  # noqa: E501
        self._client.on('Network.responseReceived', self._onResponseReceived)
        self._client.on('Network.loadingFinished', self._onLoadingFinished)
        self._client.on('Network.loadingFailed', self._onLoadingFailed)

    async def setExtraHTTPHeaders(self, extraHTTPHeaders: Dict[str,
                                                               str]) -> None:
        """Set extra http headers."""
        self._extraHTTPHeaders = OrderedDict()
        headers = OrderedDict()  # type: Dict[str, str]
        for k, v in extraHTTPHeaders.items():
            self._extraHTTPHeaders[k] = v
            headers[k] = v
        await self._client.send('Network.setExtraHTTPHeaders',
                                {'headers': headers})

    def extraHTTPHeaders(self) -> Dict[str, str]:
        """Get extra http headers."""
        return dict(**self._extraHTTPHeaders)

    async def setUserAgent(self, userAgent: str) -> Any:
        """Set user agent."""
        return await self._client.send('Network.setUserAgentOverride',
                                       {'userAgent': userAgent})

    async def setRequestInterceptionEnabled(self, value: bool) -> None:
        """Enable request intercetion."""
        await self._client.send('Network.setRequestInterceptionEnabled',
                                {'enabled': bool(value)})
        self._requestInterceptionEnabled = value

    def _onRequestIntercepted(self, event: dict) -> None:
        event['request']['url'] = removeURLHash(event['request'].get('url'))

        if event.get('redirectStatusCode'):
            request = self._interceptionIdToRequest[event['interceptionId']]
            if not request:
                raise NetworkError('INTERNAL ERROR: failed to find request '
                                   'for interception redirect.')
            self._handleRequestRedirect(request,
                                        event.get('redirectStatusCode', 0),
                                        event.get('redirectHeaders', {}))
            self._handleRequestStart(request._requestId,
                                     event.get('interceptionId', ''),
                                     event.get('redirectUrl', ''),
                                     event.get('resourceType', ''),
                                     event.get('request', {}))
            return
        requestHash = generateRequestHash(event['request'])
        self._requestHashToInterceptions.set(requestHash, event)
        self._maybeResolveInterception(requestHash)

    def _handleRequestRedirect(self, request: 'Request', redirectStatus: int,
                               redirectHeaders: dict) -> None:
        response = Response(self._client, request, redirectStatus,
                            redirectHeaders)
        request._response = response
        self._requestIdToRequest.pop(request._requestId, None)
        self._interceptionIdToRequest.pop(request._interceptionId, None)
        self.emit(NetworkManager.Events.Response, response)
        self.emit(NetworkManager.Events.RequestFinished, request)

    def _handleRequestStart(self, requestId: str, interceptionId: str,
                            url: str, resourceType: str,
                            requestPayload: dict) -> None:
        request = Request(self._client, requestId, interceptionId, url,
                          resourceType, requestPayload)
        self._requestIdToRequest[requestId] = request
        self._interceptionIdToRequest[interceptionId] = request
        self.emit(NetworkManager.Events.Request, request)

    def _onRequestWillBeSent(self, event: dict) -> None:
        if (self._requestInterceptionEnabled
                and not event['request'].get('url', '').startswith('data:')):
            if event.get('redirectResponse'):
                return
            requestHash = generateRequestHash(event['request'])
            self._requestHashToRequestIds.set(requestHash,
                                              event.get('requestId', ''))
            self._maybeResolveInterception(requestHash)
            return
        if event.get('redirectResponse'):
            request = self._requestIdToRequest.get(event['requestId'])
            if request is not None:
                redirectResponse = event.get('redirectResponse', {})
                self._handleRequestRedirect(
                    request,
                    redirectResponse.get('status'),
                    redirectResponse.get('headers'),
                )
        self._handleRequestStart(
            event.get('requestId', ''),
            '',
            event.get('request', {}).get('url', ''),
            event.get('request', {}).get('resourceType', ''),
            event.get('request', {}),
        )

    def _maybeResolveInterception(self, requestHash: str) -> None:
        requestId = self._requestHashToRequestIds.firstValue(requestHash)
        interception = self._requestHashToInterceptions.firstValue(requestHash)
        if not requestId or not interception:
            return
        self._requestHashToRequestIds.delete(requestHash, requestId)
        self._requestHashToInterceptions.delete(requestHash, interception)
        request_obj = interception.get('request', {})
        self._handleRequestStart(
            requestId,
            interception.get('interceptionId', ''),
            request_obj.get('url', ''),
            request_obj.get('resourceType'),
            request_obj,
        )

    def _onResponseReceived(self, event: dict) -> None:
        request = self._requestIdToRequest.get(event['requestId'])
        # FileUpload sends a response without a matching request.
        if not request:
            return
        _resp = event.get('response', {})
        response = Response(self._client, request, _resp.get('status', 0),
                            _resp.get('headers', {}))
        request._response = response
        self.emit(NetworkManager.Events.Response, response)

    def _onLoadingFinished(self, event: dict) -> None:
        requestId = event.get('requestId', '')
        interceptionId = event.get('interceptionId', '')
        request = self._requestIdToRequest.get(requestId)
        # For certain requestIds we never receive requestWillBeSent event.
        # @see https://crbug.com/750469
        if not request:
            return
        request._completePromiseFulfill()
        self._requestIdToRequest.pop(requestId, None)
        self._interceptionIdToRequest.pop(interceptionId, None)
        self.emit(NetworkManager.Events.RequestFinished, request)

    def _onLoadingFailed(self, event: dict) -> None:
        request = self._requestIdToRequest.get(event['requestId'])
        # For certain requestIds we never receive requestWillBeSent event.
        # @see https://crbug.com/750469
        if not request:
            return
        request._completePromiseFulfill()
        self._requestIdToRequest.pop(event['requestId'], None)
        self._interceptionIdToRequest.pop(event['interceptionId'], None)
        self.emit(NetworkManager.Events.RequestFailed, request)
예제 #6
0
class NetworkManager(EventEmitter):
    """NetworkManager class."""

    Events = SimpleNamespace(
        Request='request',
        Response='response',
        RequestFailed='requestfailed',
        RequestFinished='requestfinished',
    )

    def __init__(self, client: CDPSession, frameManager: FrameManager) -> None:
        """Make new NetworkManager."""
        super().__init__()
        self._client = client
        self._frameManager = frameManager
        self._requestIdToRequest: Dict[Optional[str], Request] = dict()
        self._interceptionIdToRequest: Dict[Optional[str], Request] = dict()
        self._extraHTTPHeaders: OrderedDict[str, str] = OrderedDict()
        self._offline: bool = False
        self._credentials: Optional[Dict[str, str]] = None
        self._attemptedAuthentications: Set[str] = set()
        self._userRequestInterceptionEnabled = False
        self._protocolRequestInterceptionEnabled = False
        self._requestHashToRequestIds = Multimap()
        self._requestHashToInterceptionIds = Multimap()

        self._client.on('Network.requestWillBeSent', self._onRequestWillBeSent)
        self._client.on('Network.requestIntercepted', self._onRequestIntercepted)  # noqa: E501
        self._client.on('Network.responseReceived', self._onResponseReceived)
        self._client.on('Network.loadingFinished', self._onLoadingFinished)
        self._client.on('Network.loadingFailed', self._onLoadingFailed)

    async def authenticate(self, credentials: Dict[str, str]) -> None:
        """Provide credentials for http auth."""
        self._credentials = credentials
        await self._updateProtocolRequestInterception()

    async def setExtraHTTPHeaders(self, extraHTTPHeaders: Dict[str, str]
                                  ) -> None:
        """Set extra http headers."""
        self._extraHTTPHeaders = OrderedDict()
        for k, v in extraHTTPHeaders.items():
            if not isinstance(v, str):
                raise TypeError(
                    f'Expected value of header "{k}" to be string, '
                    'but {} is found.'.format(type(v)))
            self._extraHTTPHeaders[k.lower()] = v
        await self._client.send('Network.setExtraHTTPHeaders',
                                {'headers': self._extraHTTPHeaders})

    def extraHTTPHeaders(self) -> Dict[str, str]:
        """Get extra http headers."""
        return dict(**self._extraHTTPHeaders)

    async def setOfflineMode(self, value: bool) -> None:
        """Change offline mode enable/disable."""
        if self._offline == value:
            return
        self._offline = value
        await self._client.send('Network.emulateNetworkConditions', {
            'offline': self._offline,
            'latency': 0,
            'downloadThroughput': -1,
            'uploadThroughput': -1,
        })

    async def setUserAgent(self, userAgent: str) -> None:
        """Set user agent."""
        await self._client.send('Network.setUserAgentOverride',
                                {'userAgent': userAgent})

    async def setRequestInterception(self, value: bool) -> None:
        """Enable request intercetion."""
        self._userRequestInterceptionEnabled = value
        await self._updateProtocolRequestInterception()

    async def _updateProtocolRequestInterception(self) -> None:
        enabled = (self._userRequestInterceptionEnabled or
                   bool(self._credentials))
        if enabled == self._protocolRequestInterceptionEnabled:
            return
        self._protocolRequestInterceptionEnabled = enabled
        patterns = [{'urlPattern': '*'}] if enabled else []
        await asyncio.gather(
            self._client.send(
                'Network.setCacheDisabled',
                {'cacheDisabled': enabled},
            ),
            self._client.send(
                'Network.setRequestInterception',
                {'patterns': patterns},
            )
        )

    def _onRequestIntercepted(self, event: dict) -> None:  # noqa: C901
        if event.get('authChallenge'):
            response = 'Default'
            if event['interceptionId'] in self._attemptedAuthentications:
                response = 'CancelAuth'
            elif self._credentials:
                response = 'ProvideCredentials'
                self._attemptedAuthentications.add(event['interceptionId'])
            username = getattr(self, '_credentials', {}).get('username')
            password = getattr(self, '_credentials', {}).get('password')
            asyncio.ensure_future(self._client.send(
                'Network.continueInterceptedRequest', {
                    'interceptionId': event['interceptionId'],
                    'authChallengeResponse': {
                        'response': response,
                        'username': username,
                        'password': password,
                    }
                }
            ))
            return

        if (not self._userRequestInterceptionEnabled and
                self._protocolRequestInterceptionEnabled):
            asyncio.ensure_future(self._client.send(
                'Network.continueInterceptedRequest', {
                    'interceptionId': event['interceptionId'],
                }
            ))

        if 'redirectURL' in event:
            request = self._interceptionIdToRequest.get(
                event.get('interceptionId', ''))
            if request:
                self._handleRequestRedirect(request,
                                            event.get('redirectStatusCode', 0),
                                            event.get('redirectHeaders', {}))
                self._handleRequestStart(request._requestId,
                                         event.get('interceptionId', ''),
                                         event.get('redirectUrl', ''),
                                         event.get('resourceType', ''),
                                         event.get('request', {}),
                                         event.get('frameId'))
            return

        requestHash = generateRequestHash(event['request'])
        requestId = self._requestHashToRequestIds.firstValue(requestHash)
        if requestId:
            self._requestHashToRequestIds.delete(requestHash, requestId)
            self._handleRequestStart(
                requestId, event['interceptionId'], event['request']['url'],
                event['resourceType'], event['request'], event['frameId']
            )
        else:
            self._requestHashToInterceptionIds.set(
                requestHash, event['interceptionId'])
            self._handleRequestStart(
                None, event['interceptionId'], event['request']['url'],
                event['resourceType'], event['request'], event['frameId']
            )

    def _handleRequestRedirect(self, request: 'Request', redirectStatus: int,
                               redirectHeaders: Dict) -> None:
        response = Response(
            self._client, request, redirectStatus, redirectHeaders)
        request._response = response
        self._requestIdToRequest.pop(request._requestId, None)
        self._interceptionIdToRequest.pop(request._interceptionId, None)
        self._attemptedAuthentications.discard(request._interceptionId)
        self.emit(NetworkManager.Events.Response, response)
        self.emit(NetworkManager.Events.RequestFinished, request)

    def _handleRequestStart(self, requestId: Optional[str],
                            interceptionId: str, url: str, resourceType: str,
                            requestPayload: Dict, frameId: Optional[str]
                            ) -> None:
        frame = None
        if frameId and self._frameManager is not None:
            frame = self._frameManager.frame(frameId)

        request = Request(self._client, requestId, interceptionId,
                          self._userRequestInterceptionEnabled, url,
                          resourceType, requestPayload, frame)
        if requestId:
            self._requestIdToRequest[requestId] = request
        if interceptionId:
            self._interceptionIdToRequest[interceptionId] = request
        self.emit(NetworkManager.Events.Request, request)

    def _onRequestWillBeSent(self, event: dict) -> None:
        if self._protocolRequestInterceptionEnabled:
            if event.get('redirectResponse'):
                return
            requestHash = generateRequestHash(event['request'])
            interceptionId = self._requestHashToInterceptionIds.firstValue(
                requestHash)
            request = self._interceptionIdToRequest.get(interceptionId)
            if request:
                request._requestId = event['requestId']
                self._requestIdToRequest[event['requestId']] = request
                self._requestHashToInterceptionIds.delete(
                    requestHash, interceptionId)
            else:
                self._requestHashToRequestIds.set(
                    requestHash, event['requestId'])
            return
        if event.get('redirectResponse'):
            request = self._requestIdToRequest[event['requestId']]
            if request:
                redirectResponse = event.get('redirectResponse', {})
                self._handleRequestRedirect(
                    request,
                    redirectResponse.get('status'),
                    redirectResponse.get('headers'),
                )
        self._handleRequestStart(
            event.get('requestId', ''), '',
            event.get('request', {}).get('url', ''),
            event.get('type', ''),
            event.get('request', {}),
            event.get('frameId'),
        )

    def _onResponseReceived(self, event: dict) -> None:
        request = self._requestIdToRequest.get(event['requestId'])
        # FileUpload sends a response without a matching request.
        if not request:
            return
        _resp = event.get('response', {})
        response = Response(self._client, request,
                            _resp.get('status', 0),
                            _resp.get('headers', {}))
        request._response = response
        self.emit(NetworkManager.Events.Response, response)

    def _onLoadingFinished(self, event: dict) -> None:
        request = self._requestIdToRequest.get(event.get('requestId', ''))
        # For certain requestIds we never receive requestWillBeSent event.
        # @see https://crbug.com/750469
        if not request:
            return
        request._completePromiseFulfill()
        self._requestIdToRequest.pop(request._requestId, None)
        self._interceptionIdToRequest.pop(request._interceptionId, None)
        self._attemptedAuthentications.discard(request._interceptionId)
        self.emit(NetworkManager.Events.RequestFinished, request)

    def _onLoadingFailed(self, event: dict) -> None:
        request = self._requestIdToRequest.get(event['requestId'])
        # For certain requestIds we never receive requestWillBeSent event.
        # @see https://crbug.com/750469
        if not request:
            return
        request._failureText = event.get('errorText')
        request._completePromiseFulfill()
        self._requestIdToRequest.pop(request._requestId, None)
        self._interceptionIdToRequest.pop(request._interceptionId, None)
        self._attemptedAuthentications.discard(request._interceptionId)
        self.emit(NetworkManager.Events.RequestFailed, request)