Song/fix/mitm page (#1902)

* fix: mitm部分调整 * zancun
yaklang · Aug 12, 2024 · 7da17bd · 7da17bd
1 parent 4ff5d0f
commit 7da17bd
Show file tree

Hide file tree

Showing 8 changed files with 158 additions and 69 deletions.
diff --git a/app/renderer/src/main/src/components/HTTPFlowTable/HTTPFlowTable.tsx b/app/renderer/src/main/src/components/HTTPFlowTable/HTTPFlowTable.tsx
@@ -554,7 +554,7 @@ const defSort: SortProps = {
     orderBy: "id"
 }
 
-const SourceType = [
+export const SourceType = [
     {text: "MITM", value: "mitm"},
     {text: "插件", value: "scan"},
     {
@@ -3393,24 +3393,40 @@ export const HTTPFlowTable = React.memo<HTTPFlowTableProp>((props) => {
         )
     })
 
+    useEffect(() => {
+        if (props.params?.SourceType !== undefined) {
+            let selectTypeList = props.params?.SourceType.split(",") || [""]
+            const newParams = {...params, SourceType: selectTypeList.join(",")}
+            setParams(newParams)
+            setTimeout(() => {
+                updateData()
+            }, 10)
+        }
+    }, [props.params?.SourceType])
+
     const onHasParamsJumpHistory = useMemoizedFn((mitmHasParamsNames: string) => {
-        const arr = params.SourceType?.split(",") || []
-        arr.push("scan")
-        const selectTypeList = [...new Set(arr)]
+        const mitmHasParamsNamesArr = mitmHasParamsNames.split(",")
+        let selectTypeList = params.SourceType?.split(",") || [""]
+        if (mitmHasParamsNamesArr[0] !== "" || !mitmHasParamsNamesArr.length) {
+            selectTypeList = ["scan"]
+        } else {
+            selectTypeList = selectTypeList.filter((item) => item !== "scan")
+            if (selectTypeList[0] === "" || !selectTypeList.length) {
+                selectTypeList = ["mitm"]
+            }
+        }
         const newParams = {...params, SourceType: selectTypeList.join(",")}
         getHTTPFlowsFieldGroup(true, (t: FiltersItemProps[]) => {
             let tagsList: string[] = []
-            mitmHasParamsNames.split(",").forEach((item) => {
+            mitmHasParamsNamesArr.forEach((item) => {
                 if (t.findIndex((ele) => ele.label === item) !== -1) {
                     tagsList.push(item)
                 }
             })
             setTagsFilter(tagsList)
             newParams.Tags = tagsList
             setParams(newParams)
-            setTimeout(() => {
-                updateData()
-            }, 10)
+            emiter.emit("onHistorySourceTypeToMitm", newParams.SourceType)
         })
     })
 
@@ -3467,39 +3483,45 @@ export const HTTPFlowTable = React.memo<HTTPFlowTableProp>((props) => {
                                     style["http-history-table-row"]
                                 )}
                             >
-                                <div
-                                    style={onlyShowSearch ? {marginLeft: 8} : {}}
-                                    className={classNames(style["http-history-table-flex"])}
-                                >
-                                    {!onlyShowSearch &&
-                                        SourceType.map((tag) => (
-                                            <YakitCheckableTag
-                                                key={tag.value}
-                                                checked={!!params.SourceType?.split(",").includes(tag.value)}
-                                                onChange={(checked) => {
-                                                    if (checked) {
-                                                        const selectTypeList = [
-                                                            ...(params.SourceType?.split(",") || []),
-                                                            tag.value
-                                                        ]
-                                                        setParams({...params, SourceType: selectTypeList.join(",")})
-                                                    } else {
-                                                        const selectTypeList = (
-                                                            params.SourceType?.split(",") || []
-                                                        ).filter((ele) => ele !== tag.value)
-                                                        setParams({...params, SourceType: selectTypeList.join(",")})
-                                                    }
-                                                    setTimeout(() => {
-                                                        updateData()
-                                                    }, 10)
-                                                }}
-                                            >
-                                                {tag.text}
-                                            </YakitCheckableTag>
-                                        ))}
-                                </div>
                                 {pageType === "History" && (
                                     <>
+                                        <div
+                                            style={onlyShowSearch ? {marginLeft: 8} : {}}
+                                            className={classNames(style["http-history-table-flex"])}
+                                        >
+                                            {!onlyShowSearch &&
+                                                SourceType.map((tag) => (
+                                                    <YakitCheckableTag
+                                                        key={tag.value}
+                                                        checked={!!params.SourceType?.split(",").includes(tag.value)}
+                                                        onChange={(checked) => {
+                                                            if (checked) {
+                                                                const selectTypeList = [
+                                                                    ...(params.SourceType?.split(",") || []),
+                                                                    tag.value
+                                                                ]
+                                                                setParams({
+                                                                    ...params,
+                                                                    SourceType: selectTypeList.join(",")
+                                                                })
+                                                            } else {
+                                                                const selectTypeList = (
+                                                                    params.SourceType?.split(",") || []
+                                                                ).filter((ele) => ele !== tag.value)
+                                                                setParams({
+                                                                    ...params,
+                                                                    SourceType: selectTypeList.join(",")
+                                                                })
+                                                            }
+                                                            setTimeout(() => {
+                                                                updateData()
+                                                            }, 10)
+                                                        }}
+                                                    >
+                                                        {tag.text}
+                                                    </YakitCheckableTag>
+                                                ))}
+                                        </div>
                                         <div className={style["http-history-table-flex"]}>
                                             {shieldData?.data.length > 0 && (
                                                 <div style={{marginRight: 16}}>

diff --git a/app/renderer/src/main/src/components/HTTPHistory.tsx b/app/renderer/src/main/src/components/HTTPHistory.tsx
@@ -27,8 +27,8 @@ export interface HTTPPacketFuzzable {
 export type HTTPHistorySourcePageType = "MITM" | "History"
 
 export interface HTTPHistoryProp extends HTTPPacketFuzzable {
-    websocket?: boolean
     pageType?: HTTPHistorySourcePageType
+    params?: YakQueryHTTPFlowRequest
 }
 
 type tabKeys = "web-tree"
@@ -40,7 +40,7 @@ interface HTTPHistoryTabsItem {
 }
 
 export const HTTPHistory: React.FC<HTTPHistoryProp> = (props) => {
-    const {pageType, downstreamProxyStr} = props
+    const {pageType, downstreamProxyStr, params} = props
     const ref = useRef(null)
     const [inViewport] = useInViewport(ref)
     const {isRefreshHistory, setIsRefreshHistory} = useStore()
@@ -281,13 +281,7 @@ export const HTTPHistory: React.FC<HTTPHistoryProp> = (props) => {
                             >
                                 <HTTPFlowTable
                                     noHeader={true}
-                                    params={
-                                        props?.websocket
-                                            ? ({
-                                                  OnlyWebsocket: true
-                                              } as YakQueryHTTPFlowRequest)
-                                            : undefined
-                                    }
+                                    params={params}
                                     searchURL={searchURL}
                                     includeInUrl={includeInUrl}
                                     // tableHeight={200}

diff --git a/app/renderer/src/main/src/pages/mitm/MITMPage.tsx b/app/renderer/src/main/src/pages/mitm/MITMPage.tsx
@@ -36,7 +36,7 @@ import {YakitSelect} from "@/components/yakitUI/YakitSelect/YakitSelect"
 import emiter from "@/utils/eventBus/eventBus"
 import {YakitRoute} from "@/enums/yakitRoute"
 import {apiDownloadPluginOther, apiQueryYakScript} from "../plugins/utils"
-import {setRemoteValue} from "@/utils/kv"
+import {getRemoteValue, setRemoteValue} from "@/utils/kv"
 import {MITMConsts} from "./MITMConsts"
 import {onSetRemoteValuesBase} from "@/components/yakitUI/utils"
 import {CacheDropDownGV} from "@/yakitGV"
@@ -75,7 +75,7 @@ export const CONST_DEFAULT_ENABLE_INITIAL_PLUGIN = "CONST_DEFAULT_ENABLE_INITIAL
 export const MITMPage: React.FC<MITMPageProp> = (props) => {
     // 整体的劫持状态
     const [status, setStatus, getStatus] = useGetState<"idle" | "hijacked" | "hijacking">("idle")
-    const [isHasParams, setIsHasParams] = useState<boolean>(true) // mitm插件类型是否带参数
+    const [isHasParams, setIsHasParams] = useState<boolean>(false) // mitm插件类型是否带参数
     // 通过启动表单的内容
     const [addr, setAddr] = useState("")
     const [host, setHost] = useState("127.0.0.1")
@@ -313,7 +313,7 @@ export const MITMPage: React.FC<MITMPageProp> = (props) => {
                         logs={[]}
                         statusCards={[]}
                         downstreamProxyStr={downstreamProxyStr}
-                        isHasParams={isHasParams}
+                        isHasParams={false}
                         onIsHasParams={setIsHasParams}
                         showPluginHistoryList={showPluginHistoryList}
                         setShowPluginHistoryList={setShowPluginHistoryList}
@@ -458,8 +458,19 @@ interface MITMServerProps {
     setShowPluginHistoryList: (l: string[]) => void
 }
 export const MITMServer: React.FC<MITMServerProps> = React.memo((props) => {
-    const {visible, setVisible, status, setStatus, logs, statusCards, downstreamProxyStr, isHasParams, onIsHasParams, showPluginHistoryList, setShowPluginHistoryList} =
-        props
+    const {
+        visible,
+        setVisible,
+        status,
+        setStatus,
+        logs,
+        statusCards,
+        downstreamProxyStr,
+        isHasParams,
+        onIsHasParams = () => {},
+        showPluginHistoryList,
+        setShowPluginHistoryList
+    } = props
 
     const [openTabsFlag, setOpenTabsFlag] = useState<boolean>(true)
 
@@ -488,7 +499,33 @@ export const MITMServer: React.FC<MITMServerProps> = React.memo((props) => {
 
     const [loadedPluginLen, setLoadedPluginLen] = useState<number>(0)
 
-
+    useDebounceEffect(
+        () => {
+            if (status === "idle") {
+                const CHECK_CACHE_LIST_DATA = "CHECK_CACHE_LIST_DATA"
+                getRemoteValue(CHECK_CACHE_LIST_DATA)
+                    .then((data: string) => {
+                        getRemoteValue(CONST_DEFAULT_ENABLE_INITIAL_PLUGIN).then((is) => {
+                            if (!!data && !!is) {
+                                const cacheData: string[] = JSON.parse(data)
+                                if (cacheData.length) {
+                                    onIsHasParams(false) 
+                                }
+                            } else {
+                                if (noParamsCheckList.length) {
+                                    onIsHasParams(false)
+                                } else {
+                                    onIsHasParams(true)
+                                }
+                            }
+                        })
+                    })
+
+            }
+        },
+        [status, noParamsCheckList],
+        {wait: 300}
+    )
 
     const onSubmitYakScriptId = useMemoizedFn((id: number, params: YakExecutorParam[]) => {
         info(`加载 MITM 插件[${id}]`)

diff --git a/app/renderer/src/main/src/pages/mitm/MITMServerHijacking/MITMHijackedContent.tsx b/app/renderer/src/main/src/pages/mitm/MITMServerHijacking/MITMHijackedContent.tsx
@@ -103,6 +103,8 @@ const MITMHijackedContent: React.FC<MITMHijackedContentProps> = React.memo((prop
     const [beautifyOpen, setBeautifyOpen] = useState<boolean>(false)
     const [currentBeautifyPacket, setCurrentBeautifyPacket] = useState<string>("")
 
+    const [sourceType, setSourceType] = useState<string>("mitm")
+
     const getMITMFilter = useMemoizedFn(() => {
         ipcRenderer
             .invoke("mitm-get-filter")
@@ -293,6 +295,7 @@ const MITMHijackedContent: React.FC<MITMHijackedContentProps> = React.memo((prop
                 forward()
             }
             setShowPluginHistoryList([])
+            setSourceType("mitm")
         } catch (e) {
             console.info(e)
         }
@@ -436,7 +439,7 @@ const MITMHijackedContent: React.FC<MITMHijackedContentProps> = React.memo((prop
                     />
                 )
             case "log":
-                return <MITMLogHeardExtra />
+                return <MITMLogHeardExtra sourceType={sourceType} onSetSourceType={setSourceType} />
             default:
                 break
         }
@@ -484,7 +487,11 @@ const MITMHijackedContent: React.FC<MITMHijackedContentProps> = React.memo((prop
             case "log":
                 return (
                     <>
-                        <HTTPHistory pageType='MITM' downstreamProxyStr={downstreamProxyStr} />
+                        <HTTPHistory
+                            pageType='MITM'
+                            downstreamProxyStr={downstreamProxyStr}
+                            params={{SourceType: sourceType}}
+                        />
                     </>
                 )
             // 被动日志

diff --git a/app/renderer/src/main/src/pages/mitm/MITMServerHijacking/MITMLog.tsx b/app/renderer/src/main/src/pages/mitm/MITMServerHijacking/MITMLog.tsx
@@ -1,21 +1,20 @@
 import React, {useEffect, useMemo, useRef, useState} from "react"
 import emiter from "@/utils/eventBus/eventBus"
 import styles from "./MITMServerHijacking.module.scss"
-import {
-    HTTPFlowShield,
-    ShieldData
-} from "@/components/HTTPFlowTable/HTTPFlowTable"
-import {
-    useMemoizedFn,
-} from "ahooks"
+import {HTTPFlowShield, ShieldData, SourceType} from "@/components/HTTPFlowTable/HTTPFlowTable"
+import {useMemoizedFn} from "ahooks"
 import {yakitFailed, yakitNotify} from "@/utils/notification"
 import {YakitDropdownMenu} from "@/components/yakitUI/YakitDropdownMenu/YakitDropdownMenu"
 import {YakitButton} from "@/components/yakitUI/YakitButton/YakitButton"
+import {YakitCheckableTag} from "@/components/yakitUI/YakitTag/YakitCheckableTag"
 
 const {ipcRenderer} = window.require("electron")
 interface MITMLogHeardExtraProps {
+    sourceType: string
+    onSetSourceType: (s: string) => void
 }
 export const MITMLogHeardExtra: React.FC<MITMLogHeardExtraProps> = React.memo((props) => {
+    const {sourceType, onSetSourceType} = props
     // 屏蔽数据
     const [shieldData, setShieldData] = useState<ShieldData>({
         data: []
@@ -28,13 +27,13 @@ export const MITMLogHeardExtra: React.FC<MITMLogHeardExtraProps> = React.memo((p
         }
     }, [])
 
-    const onGetMITMShieldData = useMemoizedFn((str:string)=>{
+    const onGetMITMShieldData = useMemoizedFn((str: string) => {
         const value = JSON.parse(str)
         setShieldData(value)
     })
 
-    const cancleFilter = useMemoizedFn((value)=>{
-        emiter.emit("cancleMitmFilterEvent",JSON.stringify(value))
+    const cancleFilter = useMemoizedFn((value) => {
+        emiter.emit("cancleMitmFilterEvent", JSON.stringify(value))
     })
 
     const cleanMitmLogTableData = useMemoizedFn((params: {DeleteAll: boolean; Filter?: {}}) => {
@@ -47,12 +46,41 @@ export const MITMLogHeardExtra: React.FC<MITMLogHeardExtraProps> = React.memo((p
                 yakitNotify("error", `历史记录删除失败: ${e}`)
             })
             .finally(() => {
-                emiter.emit("onDeleteToUpdate",JSON.stringify({sourcePage:"MITM"}))
+                emiter.emit("onDeleteToUpdate", JSON.stringify({sourcePage: "MITM"}))
             })
     })
 
+    const onHistorySourceTypeToMitm = useMemoizedFn((sourceType: string) => {
+        onSetSourceType(sourceType)
+    })
+    useEffect(() => {
+        emiter.on("onHistorySourceTypeToMitm", onHistorySourceTypeToMitm)
+        return () => {
+            emiter.off("onHistorySourceTypeToMitm", onHistorySourceTypeToMitm)
+        }
+    }, [])
+
     return (
         <div className={styles["mitm-log-heard"]}>
+            <div>
+                {SourceType.map((tag) => (
+                    <YakitCheckableTag
+                        key={tag.value}
+                        checked={!!sourceType.split(",").includes(tag.value)}
+                        onChange={(checked) => {
+                            if (checked) {
+                                const selectTypeList = [...(sourceType.split(",") || []), tag.value]
+                                onSetSourceType(selectTypeList.join(","))
+                            } else {
+                                const selectTypeList = (sourceType.split(",") || []).filter((ele) => ele !== tag.value)
+                                onSetSourceType(selectTypeList.join(","))
+                            }
+                        }}
+                    >
+                        {tag.text}
+                    </YakitCheckableTag>
+                ))}
+            </div>
             <YakitDropdownMenu
                 menu={{
                     data: [

diff --git a/app/renderer/src/main/src/pages/mitm/MITMServerHijacking/MITMServerHijacking.module.scss b/app/renderer/src/main/src/pages/mitm/MITMServerHijacking/MITMServerHijacking.module.scss
@@ -181,7 +181,7 @@
 .mitm-log-heard {
     display: flex;
     align-items: center;
-    justify-content: flex-end;
+    justify-content: space-between;
     flex: 1;
 }
 

diff --git a/app/renderer/src/main/src/pages/mitm/MITMServerHijacking/MITMServerHijacking.tsx b/app/renderer/src/main/src/pages/mitm/MITMServerHijacking/MITMServerHijacking.tsx
@@ -140,7 +140,7 @@ export const MITMServerHijacking: React.FC<MITMServerHijackingProp> = (props) =>
             ipcRenderer
                 .invoke("mitm-stop-call")
                 .then(() => {
-                    onIsHasParams(true)
+                    onIsHasParams(false)
                     setShowPluginHistoryList([])
                     setStatus("idle")
                     resolve("ok")