自研抓包服务实现（4）—mitmproxy实现Rewrite

抓包服务实现（4）—mitmproxy实现Rewrite

Charles代理有些典型功能

• Map Remote 指向域名
• Map Local 指向返回body
• Rewrite 改写请求字段
• Start Throttlling 开始限速
• Clear 清除
• Stop Recording 停止记录
• Filter 过滤
• Breakpoint 断点
• Repeat 重复
• Copy 复制

实现Rewrite需要哪些？

1. 改变配置能够及时生效；不改变配置，之前的配置依旧生效
2. 需要一个总开关，一键启动map
3. 需要知道当前做了啥代理
4. 能对某个请求的request、response的字段更改
5. Map Local覆盖Rewrite

如何实现 Rewrite？

就上述几点，我们想到了db，恰好我们在之前使用了mongoDB。

1. 我们希望是在mitmproxy的flow执行之前能得到Rewrite配置，从而每次跑flow都能够实时根据Rewrite配置改变而生效

import pymongo
from mitmproxy.script import concurrent

class Counter:

    def __init__(self):
        client = pymongo.MongoClient('mongodb://localhost:27017/')
        mydb = client["mitmproxy"]
        self.temporary_rewrite_col = mydb['temporary_rewrite']
    
    def get_temporary_rewrite_data(self):
        new_list=[]
        temporary_rewrite_data = self.temporary_rewrite_col.find({})
        for temporary_rewrite in temporary_rewrite_data:
            new_list.append(temporary_rewrite)
        return new_list
        
    @concurrent
     def request(self,flow):
        new_dict=self.get_switch_settings()
        rewrite_list=self.get_temporary_rewrite_data()
        
        if new_dict['switch_rewrite']=='true':
            for rewrite in rewrite_list:
                if flow.request.url.startswith(rewrite['path']):
                    if rewrite['rewrite_way']=='request_query':
                        flow.request.query=ModifyJson(flow.request.query).modify_key(rewrite['field'],rewrite['field_value'])
                    if rewrite['rewrite_way']=='request_body':
                        flow.request.body=ModifyJson(flow.request.body).modify_key(rewrite['field'],rewrite['field_value'])
                    if rewrite['rewrite_way']=='request_header':
                        flow.request.request_headers=ModifyJson(flow.request.request_headers).modify_key(rewrite['field'],rewrite['field_value'])

        return flow.request

addons=[Counter()]

1）先去查询temporary_rewrite表的配置
2）遍历配置和flow的url进行匹配，若匹配成功则判断是否需要更新request的query、body、header

2. 创建一个开关配置，这样可以实现关闭开关时，不进行map，但暂存配置依然保存；下次开关打开，暂存配置依然生效

import pymongo
from mitmproxy.script import concurrent

class Counter:

    def __init__(self):
        client = pymongo.MongoClient('mongodb://localhost:27017/')
        mydb = client["mitmproxy"]
        self.temporary_rewrite_col = mydb['temporary_rewrite']
        self.switch_settings_col=mydb['switch_settings']
    
    def get_temporary_rewrite_data(self):
        new_list=[]
        temporary_rewrite_data = self.temporary_rewrite_col.find({})
        for temporary_rewrite in temporary_rewrite_data:
            new_list.append(temporary_rewrite)
        return new_list
        
    def get_switch_settings(self):
        switch_settings_data=self.switch_settings_col.find({})
        new_dict={}
        for switch_settings in switch_settings_data:
            if switch_settings['switch']=='switch_local':
                new_dict['switch_local']=switch_settings['status']
            elif switch_settings['switch']=='switch_remote':
                new_dict['switch_remote']=switch_settings['status']
            else:
                new_dict['switch_rewrite']=switch_settings['status']
            
        return new_dict
        
    @concurrent
    def request(self,flow):
        Flow_request=flow.request
        rewrite_list=self.get_temporary_rewrite_data()
        if new_dict['switch_rewrite']=='true':
            for rewrite in rewrite_list:
                if flow.request.url.startswith(rewrite['path']):
                    if rewrite['rewrite_way']=='request_query':
                        flow.request.query=ModifyJson(flow.request.query).modify_key(rewrite['field'],rewrite['field_value'])
                    if rewrite['rewrite_way']=='request_body':
                        flow.request.body=ModifyJson(flow.request.body).modify_key(rewrite['field'],rewrite['field_value'])
                    if rewrite['rewrite_way']=='request_header':
                        flow.request.request_headers=ModifyJson(flow.request.request_headers).modify_key(rewrite['field'],rewrite['field_value'])

        return flow
    @concurrent
    def response(self, flow):
        Flow_request=flow.request
        Flow_response=flow.response

        new_dict=self.get_switch_settings()
        rewrite_list=self.get_temporary_rewrite_data()

        if new_dict['switch_rewrite']=='true':
            for rewrite in rewrite_list:
                if flow.request.url.startswith(rewrite['path']):
                    if rewrite['rewrite_way']=='response_body':
                        flow.response.text=json.dumps(ModifyJson(json.loads(flow.response.text)).modify_key(rewrite['field'],rewrite['field_value']))
                    

addons=[Counter()]

1）我们先去switch_settings表里查询开关状态
2）若为开启状态，则再执行rewrite规则
3）若为关闭状态，则不执行rewrite规则

3. 创建一个接口去查询temporary_rewrite表里的数据
4. 按照上一篇文章那样调用response API

from mitmproxy import proxy, options
from mitmproxy.tools.dump import DumpMaster
from mitmproxy.script import concurrent
from mitmproxy import flowfilter
from mitmproxy import ctx, http
import time
import pymongo
import json


class AddHeader:

    def __init__(self):
        client = pymongo.MongoClient('mongodb://localhost:27017/')
        mydb = client["mitmproxy"]
        self.mycol = mydb["all_capture"]
        self.temporary_rewrite_col = mydb['temporary_rewrite']
        self.switch_settings_col=mydb['switch_settings']
    
    def get_temporary_rewrite_data(self):
        new_list=[]
        temporary_rewrite_data = self.temporary_rewrite_col.find({})
        for temporary_rewrite in temporary_rewrite_data:
            new_list.append(temporary_rewrite)
        return new_list
    
    def get_switch_settings(self):
        switch_settings_data=self.switch_settings_col.find({})
        new_dict={}
        for switch_settings in switch_settings_data:
            if switch_settings['switch']=='switch_local':
                new_dict['switch_local']=switch_settings['status']
            elif switch_settings['switch']=='switch_remote':
                new_dict['switch_remote']=switch_settings['status']
            else:
                new_dict['switch_rewrite']=switch_settings['status']
            
        return new_dict
 
    @concurrent
    def request(self,flow):
        new_dict=self.get_switch_settings()
        rewrite_list=self.get_temporary_rewrite_data()

        if new_dict['switch_rewrite']=='true':
            for rewrite in rewrite_list:
                if flow.request.url.startswith(rewrite['path']):
                    if rewrite['rewrite_way']=='request_query':
                        flow.request.query=ModifyJson(flow.request.query).modify_key(rewrite['field'],rewrite['field_value'])
                    if rewrite['rewrite_way']=='request_body':
                        flow.request.body=ModifyJson(flow.request.body).modify_key(rewrite['field'],rewrite['field_value'])
                    if rewrite['rewrite_way']=='request_header':
                        flow.request.request_headers=ModifyJson(flow.request.request_headers).modify_key(rewrite['field'],rewrite['field_value'])

        return flow.request

    def response(self, flow):
        Flow_request=flow.request
        Flow_response=flow.response

        new_dict=self.get_switch_settings()
        local_list=self.get_temporary_local_data()
        rewrite_list=self.get_temporary_rewrite_data()

        if new_dict['switch_rewrite']=='true':
            for rewrite in rewrite_list:
                if flow.request.url.startswith(rewrite['path']):
                    if rewrite['rewrite_way']=='response_body':
                        flow.response.text=json.dumps(ModifyJson(json.loads(flow.response.text)).modify_key(rewrite['field'],rewrite['field_value']))
                    
        if new_dict['switch_local']=='true':
            for local in local_list:
                url=local['protocol_origin']+'://'+local['host_origin']+local['path_origin']
                if Flow_request.url.startswith(url):
                    Flow_response.text=local['local_response']
        
        response_data={
            'client_ip':flow.client_conn.address[0][7:],
            'request_headers':Flow_request.headers,
            'host':Flow_request.host,
            'url':Flow_request.url,
            'path':Flow_request.path,
            'body':Flow_request.text,
            'query':Flow_request.query,
            'method':Flow_request.method,
            'protocol':Flow_request.scheme,
            'timestamp_start':int(round(Flow_request.timestamp_start*1000)),
            'time_start':time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(int(Flow_request.timestamp_start))),
            'timestamp_end':int(round(Flow_request.timestamp_end*1000)),
            'duration':str(int(round(Flow_request.timestamp_end*1000))-int(round(Flow_request.timestamp_start*1000)))+' ms',
            'response_headers':Flow_response.headers,
            'status':Flow_response.status_code,
            'response':Flow_response.text,
            'size':str(len(Flow_response.raw_content))+' B'
        }
        self.mycol.insert_one(response_data.copy())
        


def start():
    myaddon = AddHeader()
    opts = options.Options(listen_port=8090)
    opts.add_option("http2", bool, False,"")
    pconf = proxy.config.ProxyConfig(opts)
    m = DumpMaster(opts)
    m.server = proxy.server.ProxyServer(pconf)
    m.addons.add(myaddon)

    try:
        m.run()
    except KeyboardInterrupt:
        m.shutdown()


if __name__ == '__main__':
    start()

至此实现了mitmproxy的Rewrite功能，基本功能已实现，这是完成了mitmproxy和mongoDB的交互。