Merge pull request #1310 from yaklang/song/feat/temporaryProject

Song/feat/temporary project

app/protos/grpc.proto

@@ -166,6 +166,8 @@ service Yak {
   rpc YakScriptRiskTypeList(Empty) returns (YakScriptRiskTypeListResponse);
   rpc SaveNewYakScript(SaveNewYakScriptRequest) returns (YakScript);
   rpc SaveYakScriptToOnline(SaveYakScriptToOnlineRequest) returns (stream SaveYakScriptToOnlineResponse);
+  rpc ExportLocalYakScript(ExportLocalYakScriptRequest) returns (ExportLocalYakScriptResponse);
+  rpc ImportYakScript(ImportYakScriptRequest) returns (stream ImportYakScriptResult);
 
   // HTTPFlow
   rpc GetHTTPFlowByHash(GetHTTPFlowByHashRequest) returns (HTTPFlow);
@@ -597,7 +599,7 @@ service Yak {
 
   rpc GetSpaceEngineStatus(GetSpaceEngineStatusRequest) returns(SpaceEngineStatus);
   rpc FetchPortAssetFromSpaceEngine(FetchPortAssetFromSpaceEngineRequest) returns (stream ExecResult);
-
+  
   rpc YakVersionAtLeast(YakVersionAtLeastRequest) returns (GeneralResponse);
 }
 
@@ -1018,7 +1020,7 @@ message AuthInfo {
     string AuthType = 3;
     string Host = 4;
     bool Forbidden = 5;
-  }
+}
 
 message ThirdPartyApplicationConfig {
     // zoomeye / hunter / shodan / fofa / github / token
@@ -1128,6 +1130,7 @@ message DebugPluginRequest {
 
   string Input = 3;
   HTTPRequestBuilderParams HTTPRequestTemplate = 4;
+  repeated KVPair ExecParams = 5;
 }
 
 message HTTPRequestBuilderResult {
@@ -1529,6 +1532,7 @@ message ProjectDescription {
   int64 UpdateAt = 9;
   string FolderName = 10;
   string ChildFolderName = 11;
+  string FileSize = 12;
 }
 
 message GetProjectsResponse {
@@ -1853,6 +1857,7 @@ message StaticAnalyzeErrorResult {
 message StaticAnalyzeErrorResponse {
   repeated StaticAnalyzeErrorResult Result = 1;
 }
+
 message SavePayloadProgress {
   double Progress = 1; // 进度
 
@@ -3223,7 +3228,6 @@ message GetYakVMBuildInMethodCompletionResponse {
 }
 
 
-
 message RenameRequest {
   string Name = 1;
   string NewName = 2;
@@ -3467,6 +3471,31 @@ message ToOnlineResult {
   string ScriptName = 1;
 }
 
+message ExportLocalYakScriptRequest {
+  string OutputDir = 1;
+  string OutputPluginDir = 2;
+  repeated int64 YakScriptIds = 3;
+  string Keywords = 4;
+  string Type = 5;
+  string UserName = 6;
+  string Tags = 7;
+}
+
+message ExportLocalYakScriptResponse {
+  string OutputDir = 1;
+}
+
+message ImportYakScriptRequest {
+  repeated string Dirs = 1;
+}
+
+message ImportYakScriptResult {
+  // 进度
+  double Progress = 1;
+  string Message = 2;
+  string MessageType = 3;
+}
+
 message GetYakScriptTagsAndTypeResponse {
   repeated TagsAndType Type = 1;
   repeated TagsAndType Tag = 2;
@@ -3960,9 +3989,9 @@ message QueryHTTPFlowRequest {
 }
 
 message ExportHTTPFlowsRequest {
-  QueryHTTPFlowRequest ExportWhere = 1;
-  repeated int64 Ids = 2;
-  repeated string FieldName = 3;
+    QueryHTTPFlowRequest ExportWhere = 1;
+    repeated int64 Ids = 2;
+    repeated string FieldName = 3;
 }
 
 message DeleteHTTPFlowRequest {

app/renderer/src/main/src/NewApp.tsx

@@ -1,6 +1,6 @@
 import React, {useRef, useEffect, useState, Suspense, lazy} from "react"
 // by types
-import {failed,warn} from "./utils/notification"
+import {failed, warn} from "./utils/notification"
 import {useHotkeys} from "react-hotkeys-hook"
 import {getCompletions} from "./utils/monacoSpec/yakCompletionSchema"
 import {showModal} from "./utils/showModal"
@@ -16,8 +16,10 @@ import {isCommunityEdition} from "@/utils/envfile"
 import {LocalGV, RemoteGV} from "./yakitGV"
 import {YakitModal} from "./components/yakitUI/YakitModal/YakitModal"
 import styles from "./app.module.scss"
-import { coordinate } from "./pages/globalVariable"
-import { remoteOperation } from "./pages/dynamicControl/DynamicControl"
+import {coordinate} from "./pages/globalVariable"
+import {remoteOperation} from "./pages/dynamicControl/DynamicControl"
+import {useTemporaryProjectStore} from "./store/temporaryProject"
+import {ProjectDescription} from "./pages/softwareSettings/ProjectManage"
 
 /** 快捷键目录 */
 const InterceptKeyword = [
@@ -56,7 +58,7 @@ const {ipcRenderer} = window.require("electron")
 interface OnlineProfileProps {
     BaseUrl: string
     Password?: string
-    IsCompany?:boolean
+    IsCompany?: boolean
 }
 
 function NewApp() {
@@ -235,15 +237,14 @@ function NewApp() {
     // 退出时 确保渲染进程各类事项已经处理完毕
     const {dynamicStatus} = yakitDynamicStatus()
     useEffect(() => {
-        ipcRenderer.on("close-windows-renderer", async(e, res: any) => {
+        ipcRenderer.on("close-windows-renderer", async (e, res: any) => {
             // 通知应用退出
-            if(dynamicStatus.isDynamicStatus){
+            if (dynamicStatus.isDynamicStatus) {
                 warn("远程控制关闭中...")
-                await remoteOperation(false,dynamicStatus)
-                ipcRenderer.invoke("app-exit") 
-            }
-            else{
-               ipcRenderer.invoke("app-exit") 
+                await remoteOperation(false, dynamicStatus)
+                ipcRenderer.invoke("app-exit")
+            } else {
+                ipcRenderer.invoke("app-exit")
             }
         })
         return () => {

app/renderer/src/main/src/components/layout/FuncDomain.tsx

@@ -73,6 +73,7 @@ import classNames from "classnames"
 import styles from "./funcDomain.module.scss"
 import yakitImg from "../../assets/yakit.jpg"
 import {onImportPlugin} from "@/pages/fuzzer/components/ShareImport"
+import { useTemporaryProjectStore } from "@/store/temporaryProject"
 
 const {ipcRenderer} = window.require("electron")
 const {Dragger} = Upload
@@ -1010,6 +1011,8 @@ const UIOpSetting: React.FC<UIOpSettingProp> = React.memo((props) => {
     const [available, setAvailable] = useState(false) // cve数据库是否可用
     const [isDiffUpdate, setIsDiffUpdate] = useState(false)
     const {dynamicStatus} = yakitDynamicStatus()
+    const {temporaryProjectId, setTemporaryProjectId} = useTemporaryProjectStore()
+
     useEffect(() => {
         onIsCVEDatabaseReady()
     }, [])
@@ -1119,7 +1122,7 @@ const UIOpSetting: React.FC<UIOpSettingProp> = React.memo((props) => {
                 addToTab("**beta-codec")
                 return
             case "invalidCache":
-                invalidCacheAndUserData()
+                invalidCacheAndUserData(temporaryProjectId, setTemporaryProjectId)
                 return
             case "pcapfix":
                 showPcapPermission()

app/renderer/src/main/src/components/layout/MacUIOp.tsx

@@ -1,14 +1,22 @@
-import React, {useEffect, useState} from "react"
+import React, {useEffect, useRef, useState} from "react"
 import {MacUIOpCloseSvgIcon, MacUIOpMaxSvgIcon, MacUIOpMinSvgIcon, MacUIOpRestoreSvgIcon} from "./icons"
 import {useMemoizedFn} from "ahooks"
 import classNames from "classnames"
 import styles from "./uiOperate.module.scss"
 import {YakitHint} from "../yakitUI/YakitHint/YakitHint"
 import {useRunNodeStore} from "@/store/runNode"
+import {useTemporaryProjectStore} from "@/store/temporaryProject"
+import {TemporaryProjectPop} from "./WinUIOp"
+import emiter from "@/utils/eventBus/eventBus"
+import { yakitFailed } from "@/utils/notification"
+import { isEnpriTraceAgent } from "@/utils/envfile"
 
 const {ipcRenderer} = window.require("electron")
 
-export interface MacUIOpProp {}
+export interface MacUIOpProp {
+    currentProjectId: string // 当前项目id
+    pageChildrenShow: boolean
+}
 
 export const MacUIOp: React.FC<MacUIOpProp> = React.memo((props) => {
     const [show, setShow] = useState<boolean>(false)
@@ -31,21 +39,69 @@ export const MacUIOp: React.FC<MacUIOpProp> = React.memo((props) => {
         }
     }, [])
 
-    /**
-     * 运行节点
-     */
-    const {runNodeList} = useRunNodeStore()
+    const {runNodeList, clearRunNodeList} = useRunNodeStore()
     const [closeRunNodeItemVerifyVisible, setCloseRunNodeItemVerifyVisible] = useState<boolean>(false)
+    const {temporaryProjectId, temporaryProjectNoPromptFlag, setTemporaryProjectId} = useTemporaryProjectStore()
+    const lastTemporaryProjectIdRef = useRef<string>("")
+    const [closeTemporaryProjectVisible, setCloseTemporaryProjectVisible] = useState<boolean>(false)
+    const lastTemporaryProjectNoPromptRef = useRef<boolean>(false)
+    const temporaryProjectPopRef = useRef<any>(null)
 
-    const handleCloseSoft = () => {
-        // 如果运行节点存在
-        if (Array.from(runNodeList).length) {
-            setCloseRunNodeItemVerifyVisible(true)
-            return
+    useEffect(() => {
+        lastTemporaryProjectNoPromptRef.current = temporaryProjectNoPromptFlag
+    }, [temporaryProjectNoPromptFlag])
+
+    useEffect(() => {
+        lastTemporaryProjectIdRef.current = temporaryProjectId
+    }, [temporaryProjectId])
+
+    const handleCloseSoft = async () => {
+        if (props.pageChildrenShow) {
+            // 如果运行节点存在
+            if (Array.from(runNodeList).length) {
+                setCloseRunNodeItemVerifyVisible(true)
+                return
+            }
+            // 如果打开得是临时项目
+            if (
+                !isEnpriTraceAgent() &&
+                lastTemporaryProjectIdRef.current === props.currentProjectId &&
+                !lastTemporaryProjectNoPromptRef.current
+            ) {
+                setCloseTemporaryProjectVisible(true)
+                return
+            } else {
+                await handleTemporaryProject()
+            }
+        } else {
+            if (Array.from(runNodeList).length) {
+                handleKillAllRunNode()
+            }
         }
         operate("close")
     }
 
+    const handleTemporaryProject = async () => {
+        if (temporaryProjectId) {
+            await ipcRenderer.invoke("DeleteProject", {Id: +temporaryProjectId, IsDeleteLocal: true})
+            setTemporaryProjectId("")
+            emiter.emit("onFeachGetCurrentProject")
+        }
+    }
+
+    const handleKillAllRunNode = async () => {
+        let promises: (() => Promise<any>)[] = []
+        Array.from(runNodeList).forEach(([key, pid]) => {
+            promises.push(() => ipcRenderer.invoke("kill-run-node", {pid}))
+        })
+        try {
+            await Promise.all(promises.map((promiseFunc) => promiseFunc()))
+            clearRunNodeList()
+        } catch (error) {
+            yakitFailed(error + "")
+        }
+    }
+
     return (
         <div
             className={styles["mac-ui-op-wrapper"]}
@@ -93,13 +149,30 @@ export const MacUIOp: React.FC<MacUIOpProp> = React.memo((props) => {
                     visible={closeRunNodeItemVerifyVisible}
                     title='是否确认关闭节点'
                     content='关闭Yakit会默认关掉所有启用的节点'
-                    onOk={() => {
-                        operate("close")
+                    onOk={async () => {
+                        await handleKillAllRunNode()
+                        setCloseRunNodeItemVerifyVisible(false)
+                        handleCloseSoft()
                     }}
                     onCancel={() => {
                         setCloseRunNodeItemVerifyVisible(false)
                     }}
                 />
+                {/* 退出临时项目确认弹框 */}
+                {closeTemporaryProjectVisible && (
+                    <TemporaryProjectPop
+                        ref={temporaryProjectPopRef}
+                        onOk={async () => {
+                            await handleTemporaryProject()
+                            setCloseTemporaryProjectVisible(false)
+                            lastTemporaryProjectIdRef.current = ""
+                            handleCloseSoft()
+                        }}
+                        onCancel={() => {
+                            setCloseTemporaryProjectVisible(false)
+                        }}
+                    />
+                )}
             </div>
         </div>
     )