End-to-End Testing｜用 Snapshot API Local Mock Server 解决 iOS App 测试资料不稳定问题

"""
Example:
    Record: mitmdump -m reverse:https://yourapihost.com -s mock.py --set record=true --set dumper_folder=loginFlow --set config_file=config.json
    Replay: mitmdump -m reverse:https://yourapihost.com -s mock.py --set dumper_folder=loginFlow --set config_file=config.json
"""

import re
import logging
import mimetypes
import os
import json
import hashlib

from pathlib import Path
from mitmproxy import ctx
from mitmproxy import http

class MockServerHandler:

    def load(self, loader):
        self.readHistory = {}
        self.configuration = {}

        loader.add_option(
            name="dumper_folder",
            typespec=str,
            default="dump",
            help="Response Dump 目录，可以 by Test Case Name 建立",
        )

        loader.add_option(
            name="network_restricted",
            typespec=bool,
            default=True,
            help="本地没有 Mapping 资料...设置 true 会 return 404、false 会去打真实请求拿资料。",
        )

        loader.add_option(
            name="record",
            typespec=bool,
            default=False,
            help="设置 true 录制 Request's Response",
        )

        loader.add_option(
            name="config_file",
            typespec=str,
            default="",
            help="设置档案路径，范例档案在下面",
        )
    
    def configure(self, updated):
        self.loadConfig()

    def loadConfig(self):
        configFile = Path(ctx.options.config_file)
        if ctx.options.config_file == "" or not configFile.exists():
            return

        self.configuration = json.loads(open(configFile, "r").read())

    def hash(self, request):
        query = request.query
        requestPath = "-".join(request.path_components)

        ignoredQueryParameterByPaths = self.configuration.get("ignored", {}).get("paths", {}).get(request.host, {}).get(requestPath, {}).get(request.method, {}).get("queryParamters", [])
        ignoredQueryParameterGlobal = self.configuration.get("ignored", {}).get("global", {}).get("queryParamters", [])

        filteredQuery = []
        if query:
            filteredQuery = [(key, value) for key, value in query.items() if key not in ignoredQueryParameterByPaths + ignoredQueryParameterGlobal]
        
        formData = []
        if request.get_content() != None and request.get_content() != b'':
            formData = json.loads(request.get_content())
        
        # or just formData = request.urlencoded_form
        # or just formData = request.multipart_form
        # depends on your api design

        ignoredFormDataParametersByPaths = self.configuration.get("ignored", {}).get("paths", {}).get(request.host, {}).get(requestPath, {}).get(request.method, {}).get("formDataParameters", [])
        ignoredFormDataParametersGlobal = self.configuration.get("ignored", {}).get("global", {}).get("formDataParameters", [])

        filteredFormData = []
        if formData:
            filteredFormData = [(key, value) for key, value in formData.items() if key not in ignoredFormDataParametersByPaths + ignoredFormDataParametersGlobal]
        
        # Serialize the dictionary to a JSON string
        hashData = {"query":sorted(filteredQuery), "form": sorted(filteredFormData)}
        json_str = json.dumps(hashData, sort_keys=True)

        # Apply SHA-256 hash function
        hash_object = hashlib.sha256(json_str.encode())
        hash_string = hash_object.hexdigest()
        
        return hash_string

    def readFromFile(self, request):
        host = request.host
        method = request.method
        hash = self.hash(request)
        requestPath = "-".join(request.path_components)

        folder = Path(ctx.options.dumper_folder) / host / method / requestPath / hash

        if not folder.exists():
            return None

        content_type = request.headers.get("content-type", "").split(";")[0]
        ext = mimetypes.guess_extension(content_type) or ".json"


        count = self.readHistory.get(host, {}).get(method, {}).get(requestPath, {}) or 0

        filepath = folder / f"Content-{str(count)}{ext}"

        while not filepath.exists() and count > 0:
            count = count - 1
            filepath = folder / f"Content-{str(count)}{ext}"

        if self.readHistory.get(host) is None:
            self.readHistory[host] = {}
        if self.readHistory.get(host).get(method) is None:
            self.readHistory[host][method] = {}
        if self.readHistory.get(host).get(method).get(requestPath) is None:
            self.readHistory[host][method][requestPath] = {}

        if filepath.exists():
            headerFilePath = folder / f"Header-{str(count)}.json"
            if not headerFilePath.exists():
                headerFilePath = None
            
            count += 1
            self.readHistory[host][method][requestPath] = count

            return {"content": filepath, "header": headerFilePath}
        else:
            return None


    def saveToFile(self, request, response):
        host = request.host
        method = request.method
        hash = self.hash(request)
        requestPath = "-".join(request.path_components)

        iterable = self.configuration.get("ignored", {}).get("paths", {}).get(request.host, {}).get(requestPath, {}).get(request.method, {}).get("iterable", False)
        
        folder = Path(ctx.options.dumper_folder) / host / method / requestPath / hash

        # create dir if not exists
        if not folder.exists():
            os.makedirs(folder)

        content_type = response.headers.get("content-type", "").split(";")[0]
        ext = mimetypes.guess_extension(content_type) or ".json"

        repeatNumber = 0
        filepath = folder / f"Content-{str(repeatNumber)}{ext}"
        while filepath.exists() and iterable == False:
            repeatNumber += 1
            filepath = folder / f"Content-{str(repeatNumber)}{ext}"
        
        # dump to file
        with open(filepath, "wb") as f:
            f.write(response.content or b'')
            
        
        headerFilepath = folder / f"Header-{str(repeatNumber)}.json"
        with open(headerFilepath, "wb") as f:
            responseDict = dict(response.headers.items())
            responseDict['_status_code'] = response.status_code
            f.write(json.dumps(responseDict).encode('utf-8'))

        return {"content": filepath, "header": headerFilepath}

    def request(self, flow):
        if ctx.options.record != True:
            host = flow.request.host
            path = flow.request.path

            result = self.readFromFile(flow.request)
            if result is not None:
                content = b''
                headers = {}
                statusCode = 200

                if result.get('content') is not None:
                    content = open(result['content'], "r").read()

                if result.get('header') is not None:
                    headers = json.loads(open(result['header'], "r").read())
                    statusCode = headers['_status_code']
                    del headers['_status_code']

                
                headers['_responseFromMitmproxy'] = '1'
                flow.response = http.Response.make(statusCode, content, headers)
                logging.info("Fullfill response from local with "+str(result['content']))
                return

            if ctx.options.network_restricted == True:
                flow.response = http.Response.make(404, b'', {'_responseFromMitmproxy': '1'})
        
    def response(self, flow):
        if ctx.options.record == True and flow.response.headers.get('_responseFromMitmproxy') != '1':
            result = self.saveToFile(flow.request, flow.response)
            logging.info("Save response to local with "+str(result['content']))

addons = [MockServerHandler()]

#...
    def response(self, flow):
        setCookies = flow.response.headers.get_all("set-cookie")
        # setCookies = ['ad=0; Domain=.xxx.com; expires=Wed, 23 Aug 2023 04:59:07 GMT; Max-Age=1800; Path=/', 'sessionid=xxxx; Secure; HttpOnly; Domain=.xxx.com; expires=Wed, 23 Aug 2023 04:59:07 GMT; Max-Age=1800; Path=/']
        
        # OR Replace Cookie Domain From .xxx.com To 127.0.0.1
        setCookies = [re.sub(r"\s*\.xxx\.com\s*", "127.0.0.1", s) for s in setCookies]

        # AND 移除安全性相关限制
        setCookies = [re.sub(r";\s*Secure\s*", "", s) for s in setCookies]
        setCookies = [re.sub(r";\s*HttpOnly;\s*", "", s) for s in setCookies]

        flow.response.headers.set_all("Set-Cookie", setCookies)

        #...