From 8c3f47e1c94a419b33d01603ecade5b09fab5f88 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E7=8E=8B=E8=89=AF?= <841369634@qq.com>
Date: Fri, 19 Apr 2024 14:59:16 +0800
Subject: [PATCH] =?UTF-8?q?feature:=20requestReplace=E5=92=8CresponseRepla?=
 =?UTF-8?q?ce=E4=B8=80=E8=B5=B7=E5=AE=8C=E6=88=90=E4=B8=8B=E8=BD=BD?=
 =?UTF-8?q?=E8=AF=B7=E6=B1=82=E7=9A=84=E5=A4=84=E7=90=86=E3=80=82?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../src/lib/interceptor/impl/req/requestReplace.js  |  7 +++++++
 .../src/lib/interceptor/impl/res/responseReplace.js | 13 ++++++++++---
 2 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/packages/mitmproxy/src/lib/interceptor/impl/req/requestReplace.js b/packages/mitmproxy/src/lib/interceptor/impl/req/requestReplace.js
index 0c2681f5..9ecf6e71 100644
--- a/packages/mitmproxy/src/lib/interceptor/impl/req/requestReplace.js
+++ b/packages/mitmproxy/src/lib/interceptor/impl/req/requestReplace.js
@@ -35,6 +35,13 @@ module.exports = {
       actions += 'headers'
     }
 
+    // 替换下载文件请求的请求地址（此功能主要是为了方便拦截配置）
+    if (requestReplaceConfig.doDownload && rOptions.path.indexOf('?DS_DOWNLOAD') >= 0) {
+      rOptions.doDownload = true
+      rOptions.path = rOptions.path.replace('?DS_DOWNLOAD', '')
+      actions += (actions ? ',' : '') + 'path:remove-DS_DOWNLOAD'
+    }
+
     res.setHeader('DS-RequestReplace-Interceptor', actions)
   },
   is (interceptOpt) {
diff --git a/packages/mitmproxy/src/lib/interceptor/impl/res/responseReplace.js b/packages/mitmproxy/src/lib/interceptor/impl/res/responseReplace.js
index 25e24ac7..c81fba9d 100644
--- a/packages/mitmproxy/src/lib/interceptor/impl/res/responseReplace.js
+++ b/packages/mitmproxy/src/lib/interceptor/impl/res/responseReplace.js
@@ -57,18 +57,25 @@ module.exports = {
   priority: 201,
   replaceResponseHeaders,
   responseIntercept (context, interceptOpt, req, res, proxyReq, proxyRes, ssl, next) {
-    const { log } = context
+    const { rOptions, log } = context
 
     if (proxyRes.statusCode !== 200) {
       return
     }
 
-    const responseConfig = interceptOpt.responseReplace
+    const responseReplaceConfig = interceptOpt.responseReplace
 
     let actions = ''
 
+    // 处理文件下载请求
+    if (responseReplaceConfig.doDownload && rOptions.doDownload) {
+      const filename = (rOptions.path.match('^.*/([^/?]+)/?(\\?.*)?$') || [])[1] || 'UNKNOWN_FILENAME'
+      res.setHeader('Content-Disposition', 'attachment; filename=' + filename)
+      actions += 'download:' + filename
+    }
+
     // 替换响应头
-    if (replaceResponseHeaders(responseConfig.headers, res, proxyRes)) {
+    if (replaceResponseHeaders(responseReplaceConfig.headers, res, proxyRes)) {
       actions += 'headers'
     }