recording_processors.py

# --------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
#
# Code generated by Microsoft (R) AutoRest Code Generator.
# Changes may cause incorrect behavior and will be lost if the code is
# regenerated.
# --------------------------------------------------------------------------

import re
from azure.cli.testsdk.scenario_tests import RecordingProcessor
from .utils import is_text_payload
try:
    from urllib.parse import urlparse
except ImportError:
    from urlparse import urlparse


class URIIdentityReplacer(RecordingProcessor):
    """Replace the identity in request uri"""

    def process_request(self, request):
        resource = (urlparse(request.uri).netloc).split('.')[0]
        request.uri = re.sub('/phoneNumbers/[%2B\d]+', '/phoneNumbers/sanitized', request.uri)
        request.uri = re.sub('/identities/([^/?]+)', '/identities/sanitized', request.uri)
        request.uri = re.sub('/chat/threads/([^/?]+)', '/chat/threads/sanitized', request.uri)
        request.uri = re.sub('/chat/threads/([^/?]+)/messages/([^/?]+)', '/chat/threads/sanitized/messages/sanitized', request.uri)
        request.uri = re.sub(resource, 'sanitized', request.uri)
        return request
    
    def process_response(self, response):
        if 'url' in response:
            response['url'] = re.sub('/phoneNumbers/[%2B\d]+', '/phoneNumbers/sanitized', response['url'])
            response['url'] = re.sub('/identities/([^/?]+)', '/identities/sanitized', response['url'])
            response['url'] = re.sub('/chat/threads/([^/?]+)', '/chat/threads/sanitized', response['url'])
            response['url'] = re.sub('/chat/threads/([^/?]+)/messages/([^/?]+)', '/chat/threads/sanitized/messages/sanitized', response['url'])
        return response


class PhoneNumberResponseReplacerProcessor(RecordingProcessor):

    def __init__(self, keys=None, replacement="sanitized"):
        self._keys = keys if keys else []
        self._replacement = replacement

    def process_response(self, response):
        import json
        try:
            body = json.loads(response['body']['string'])
            if 'phoneNumbers' in body:
                for item in body["phoneNumbers"]:
                    if isinstance(item, str):
                        body["phoneNumbers"] = [self._replacement]
                        break
                    if "phoneNumber" in item:
                        item['phoneNumber'] = self._replacement
                    if "id" in item:
                        item['id'] = self._replacement
            response['body']['string'] = json.dumps(body)
            response['url'] = self._replacement
            return response
        except (KeyError, ValueError, TypeError):
            return response


class SMSResponseReplacerProcessor(RecordingProcessor):

    def __init__(self, keys=None, replacement="sanitized"):
        self._keys = keys if keys else []
        self._replacement = replacement

    def process_request(self, request):
        import json
        try:
            if request.body is None:
                return request
            body = json.loads(request.body.decode())
            if 'smsRecipients' in body:
                for item in body["smsRecipients"]:
                    if isinstance(item, str):
                        body["smsRecipients"] = [self._replacement]
                        break
                    if "to" in item:
                        item['to'] = self._replacement
                    if "repeatabilityRequestId" in item:
                        item['repeatabilityRequestId'] = self._replacement
                    if "repeatabilityFirstSent" in item:
                        item['repeatabilityFirstSent'] = self._replacement
            
            request.body = (json.dumps(body)).encode()
        except (KeyError, ValueError, TypeError):
            return request
            
        return request

    def process_response(self, response):
        import json
        try:
            body = json.loads(response['body']['string'])
            if 'value' in body:
                for item in body["value"]:
                    if isinstance(item, str):
                        body["value"] = [self._replacement]
                        break
                    if "to" in item:
                        item['to'] = self._replacement
                    if "messageId" in item:
                        item['messageId'] = self._replacement
            response['body']['string'] = json.dumps(body)
            response['url'] = self._replacement
            return response
        except (KeyError, ValueError, TypeError):
            return response


class BodyReplacerProcessor(RecordingProcessor):
    """Sanitize the sensitive info inside request or response bodies"""

    def __init__(self, keys=None, replacement="sanitized"):
        self._replacement = replacement
        self._keys = keys if keys else []

    def process_request(self, request):
        if is_text_payload(request) and request.body:
            request.body = self._replace_keys(request.body.decode()).encode()

        return request

    def process_response(self, response):
        if is_text_payload(response) and response['body']['string']:
            response['body']['string'] = self._replace_keys(response['body']['string'])

        return response
    
    def _replace_keys(self, body):
        def _replace_recursively(dictionary):
            for key in dictionary:
                value = dictionary[key]
                if key in self._keys:
                    dictionary[key] = self._replacement
                elif isinstance(value, dict):
                    _replace_recursively(value)

        import json
        try:
            body = json.loads(body)
            _replace_recursively(body)

        except (KeyError, ValueError):
            return body

        return json.dumps(body)