ahooks#useCreation TypeScript Examples

YakScriptRunner: React.FC<YakScriptRunnerProp> = (props) => {
    const token= useCreation(()=>randomString(40),[]);
    const [infoState, { reset, setXtermRef }] = useHoldingIPCRStream(
        "exec-script-immediately", 
        "exec-yak-script",
        token,
        () => {
            setFinished(true)
        }, 
        () => {
            ipcRenderer.invoke("exec-yak-script", {
                Params: props.params,
                YakScriptId: props.script.Id,
            }, token)
        }
    )
    
    const [finished, setFinished] = useState(false);

    useEffect(()=>{
        return () => reset()
    },[])

    return <div style={{width: "100%", height: "100%"}}>
        <PluginResultUI
            script={props.script} debugMode={props.debugMode}
            results={infoState.messageState} statusCards={infoState.statusState}
            featureType={infoState.featureTypeState}
            progress={infoState.processState} feature={infoState.featureMessageState}
            loading={!finished} onXtermRef={ref => setXtermRef(ref)}
        />
    </div>
}

BatchExecutorResultByFilter: React.FC<BatchExecutorResultByFilterProp> = (props) => {
    const [activeTask, setActiveTask] = useState<BatchTask[]>([]);
    const allPluginTasks = useRef<Map<string, ExecBatchYakScriptResult[]>>(new Map<string, ExecBatchYakScriptResult[]>())
    const [allTasks, setAllTasks] = useState<BatchTask[]>([]);
    const [hitTasks, setHitTasks] = useState<ExecResultLog[]>([]);
    const [taskLog, setTaskLog, getTaskLog] = useGetState<TaskResultLog[]>([]);
    const allTasksMap = useCreation<Map<string, BatchTask>>(() => {
        return new Map<string, BatchTask>()
    }, [])
    const [jsonRisks, setJsonRisks] = useState<Risk[]>([]);

    const [tableContentHeight, setTableContentHeight] = useState<number>(0);
    const [activeKey, setActiveKey] = useState<string>("executing")

    useEffect(() => {
        if (props.executing && (!!allPluginTasks) && (!!allPluginTasks.current)) allPluginTasks.current.clear()
    }, [props.executing])

    // 转换task内的result数据
    const convertTask = (task: BatchTask) => {
        // @ts-ignore
        const results: ExecResult[] = task.Results.filter((item) => !!item.Result).map((item) => item.Result)

        const messages: ExecResultMessage[] = []
        for (let item of results) {
            if (!item.IsMessage) continue

            try {
                const raw = item.Message
                const obj: ExecResultMessage = JSON.parse(Buffer.from(raw).toString("utf8"))
                messages.push(obj)
            } catch (e) {
                console.error(e)
            }
        }

        return messages
    }

    useEffect(() => {
        const update = () => {
            const result: BatchTask[] = [];
            let hitResult: ExecResultLog[] = [];
            allTasksMap.forEach(value => {
                if (value.Results[value.Results.length - 1]?.Status === "end") {
                    result.push(value)
                    if (value.Results.length !== 0) {
                        const arr: ExecResultLog[] =
                            (convertTask(value)
                                .filter((e) => e.type === "log")
                                .map((i) => i.content) as ExecResultLog[])
                                .filter((i) => (i?.level || "").toLowerCase() === "json-risk")
                        if (arr.length > 0) {
                            hitResult = hitResult.concat(...arr)
                        }
                    }
                }
            })
            setAllTasks(result)
            setHitTasks(hitResult)
        }
        update()
        const id = setInterval(update, 3000)
        return () => {
            clearInterval(id)
        }
    }, [])

    useEffect(() => {
        let index = 0
        const activeTask = new Map<string, ExecBatchYakScriptResult[]>();
        ipcRenderer.on(`${props.token}-error`, async (e, exception) => {
            if (`${exception}`.includes("Cancelled on client")) {
                return
            }
            console.info("call exception")
            console.info(exception)
        })

        ipcRenderer.on(`${props.token}-data`, async (e, data: ExecBatchYakScriptResult) => {
            // 处理进度信息
            if (data.ProgressMessage) {
                if (!!props.setPercent) {
                    props.setPercent(data.ProgressPercent || 0)
                }
                return
            }

            // 处理其他任务信息
            const taskId: string = data.TaskId || "";
            if (taskId === "") return

            // 缓存内容
            let activeResult = activeTask.get(taskId);
            if (!activeResult) activeResult = []
            activeResult.push(data)
            activeTask.set(taskId, activeResult)
            // 缓存全部
            let allresult = allPluginTasks.current.get(taskId);
            if (!allresult) allresult = []
            allresult.push(data)
            allPluginTasks.current.set(taskId, allresult)

            if (data.Result && data.Result.IsMessage) {
                const info: TaskResultLog = JSON.parse(new Buffer(data.Result.Message).toString()).content
                if (info) {
                    info.key = index
                    index += 1
                    const arr: TaskResultLog[] = [...getTaskLog()]
                    if (arr.length >= 20) arr.shift()
                    arr.push(info)
                    setTaskLog([...arr])
                }
            }

            // 设置状态
            if (data.Status === "end") {
                activeTask.delete(taskId)
                return
            }

            // 看一下输出结果
            // if (data.Result && data.Result.IsMessage) {
            //     console.info(321,new Buffer(data.Result.Message).toString())
            // }
        })

        let cached = "";
        const syncActiveTask = () => {
            if (activeTask.size <= 0) setActiveTask([]);
            if (activeTask.size <= 0 && allPluginTasks.current.size <= 0) return

            const result: BatchTask[] = [];
            const tasks: string[] = [];
            activeTask.forEach(value => {
                if (value.length <= 0) return

                const first = value[0];
                const task = {
                    Target: first.Target || "",
                    ExtraParam: first.ExtraParams || [],
                    PoC: first.PoC,
                    TaskId: first.TaskId,
                    CreatedAt: first.Timestamp,
                } as BatchTask;
                task.Results = value;
                result.push(task)
                tasks.push(`${value.length}` + task.TaskId)
            })
            const allResult: BatchTask[] = [];
            allPluginTasks.current.forEach(value => {
                if (value.length <= 0) return

                const task = {
                    Target: value[0].Target || "",
                    ExtraParam: value[0].ExtraParams || [],
                    PoC: value[0].PoC,
                    TaskId: value[0].TaskId,
                    CreatedAt: value[0].Timestamp,
                } as BatchTask;
                task.Results = value;
                allResult.push(task)
            })

            const oldAllResult: BatchTask[] = []
            allTasksMap.forEach(value => oldAllResult.push(value))
            if (JSON.stringify(allResult) !== JSON.stringify(oldAllResult)) {
                allResult.forEach((value) => allTasksMap.set(value.TaskId, value))
            }

            const tasksRaw = tasks.sort().join("|")
            if (tasksRaw !== cached) {
                cached = tasksRaw
                setActiveTask(result)
            }
        }

        let id = setInterval(syncActiveTask, 300);
        return () => {
            ipcRenderer.removeAllListeners(`${props.token}-data`)
            ipcRenderer.removeAllListeners(`${props.token}-end`)
            ipcRenderer.removeAllListeners(`${props.token}-error`)
            allTasksMap.clear()
            setTaskLog([])
            setAllTasks([])
            setActiveKey("executing")
            clearInterval(id);
        }
    }, [props.token])

    useEffect(() => {
        if (hitTasks.length <= 0) {
            return
        }
        setJsonRisks(hitTasks.map(i => {
            try {
                return JSON.parse(i.data)
            } catch (e) {
                return undefined
            }
        }).filter(i => !!i))
    }, [hitTasks])

    return <div className="batch-executor-result">
        <div className="result-notice-body">
            <div className="notice-body">
                <div className="notice-body-header notice-font-in-progress">正在执行任务</div>
                <div className="notice-body-counter">{activeTask.length}</div>
            </div>
            <Divider type="vertical" className="notice-divider"/>
            <div className="notice-body">
                <div className="notice-body-header notice-font-completed">已完成任务</div>
                <div className="notice-body-counter">{allTasks.length}</div>
            </div>
            <Divider type="vertical" className="notice-divider"/>
            <div className="notice-body">
                <div className="notice-body-header notice-font-vuln">命中风险/漏洞</div>
                <div className="notice-body-counter">{jsonRisks.length}</div>
            </div>
        </div>

        <Divider style={{margin: 4}}/>

        <div className="result-table-body">
            <Tabs className="div-width-height-100 yakit-layout-tabs" activeKey={activeKey} onChange={setActiveKey}>
                <Tabs.TabPane tab="任务日志" key={"executing"}>
                    <div className="div-width-height-100" style={{overflow: "hidden"}}>
                        <Timeline className="body-time-line" pending={props.executing} reverse={true}>
                            {taskLog.map(item => {
                                return <Timeline.Item key={item.key}>
                                    <YakitLogFormatter data={item.data} level={item.level}
                                                       timestamp={item.timestamp} onlyTime={true} isCollapsed={true}/>
                                </Timeline.Item>
                            })}
                        </Timeline>
                    </div>
                </Tabs.TabPane>
                <Tabs.TabPane tab="命中风险与漏洞" key={"hitTable"}>
                    <div style={{width: "100%", height: "100%"}}>
                        <ReactResizeDetector
                            onResize={(width, height) => {
                                if (!width || !height) return
                                setTableContentHeight(height - 4)
                            }}
                            handleWidth={true} handleHeight={true} refreshMode={"debounce"} refreshRate={50}/>
                        <TableResizableColumn
                            virtualized={true}
                            sortFilter={() => {
                            }}
                            autoHeight={tableContentHeight <= 0}
                            height={tableContentHeight}
                            data={jsonRisks}
                            wordWrap={true}
                            renderEmpty={() => {
                                return <Empty className="table-empty" description="数据加载中"/>
                            }}
                            columns={[
                                {
                                    dataKey: "TitleVerbose",
                                    width: 400,
                                    resizable: true,
                                    headRender: () => "标题",
                                    cellRender: ({rowData, dataKey, ...props}: any) => {
                                        return (
                                            <div
                                                className="div-font-ellipsis"
                                                style={{width: "100%"}}
                                                title={rowData?.TitleVerbose || rowData.Title}
                                            >
                                                {rowData?.TitleVerbose || rowData.Title}
                                            </div>
                                        )
                                    }
                                },
                                {
                                    dataKey: "RiskTypeVerbose",
                                    width: 130,
                                    headRender: () => "类型",
                                    cellRender: ({rowData, dataKey, ...props}: any) => {
                                        return rowData?.RiskTypeVerbose || rowData.RiskType
                                    }
                                },
                                {
                                    dataKey: "Severity",
                                    width: 90,
                                    headRender: () => "等级",
                                    cellRender: ({rowData, dataKey, ...props}: any) => {
                                        const title = TitleColor.filter((item) => item.key.includes(rowData.Severity || ""))[0]
                                        return (
                                            <span className={title?.value || "title-default"}>
                                                {title ? title.name : rowData.Severity || "-"}
                                            </span>
                                        )
                                    }
                                },
                                {
                                    dataKey: "IP",
                                    width: 140,
                                    headRender: () => "IP",
                                    cellRender: ({rowData, dataKey, ...props}: any) => {
                                        return rowData?.IP || "-"
                                    }
                                },
                                {
                                    dataKey: "ReverseToken",
                                    headRender: () => "Token",
                                    cellRender: ({rowData, dataKey, ...props}: any) => {
                                        return rowData?.ReverseToken || "-"
                                    }
                                },
                                {
                                    dataKey: "operate",
                                    width: 90,
                                    fixed: "right",
                                    headRender: () => "操作",
                                    cellRender: ({rowData}: any) => {
                                        return (
                                            <a
                                                onClick={(e) => {
                                                    showModal({
                                                        width: "80%",
                                                        title: "详情",
                                                        content: (
                                                            <div style={{overflow: "auto"}}>
                                                                <RiskDetails info={rowData} isShowTime={false}/>
                                                            </div>
                                                        )
                                                    })
                                                }}
                                            >详情</a>
                                        )
                                    }
                                }
                            ].map(item => {
                                item["verticalAlign"] = "middle"
                                return item
                            })}
                        />
                    </div>
                </Tabs.TabPane>
            </Tabs>
        </div>
    </div>
}

BatchExecutorResultUI: React.FC<BatchExecutorResultUIProp> = (props) => {
    const [total, setTotal] = useState(0);
    const [finished, setFinished] = useState(0);
    const [activeTask, setActiveTask] = useState<BatchTask[]>([]);
    const allTasksMap = useCreation<Map<string, BatchTask>>(() => {
        return new Map<string, BatchTask>()
    }, [])
    const [allTasks, setAllTasks] = useState<BatchTask[]>([]);

    const [activeKey, setActiveKey] = useState<string>('executing')
    const [checked, setChecked] = useState<boolean>(false)

    const allPluginTasks = useRef<Map<string, ExecBatchYakScriptResult[]>>(new Map<string, ExecBatchYakScriptResult[]>())

    useEffect(() => {
        if (props.executing && (!!allPluginTasks) && (!!allPluginTasks.current)) allPluginTasks.current.clear()
    }, [props.executing])

    useEffect(() => {
        const update = () => {
            const result: BatchTask[] = [];
            allTasksMap.forEach(value => {
                result.push(value)
            })
            setAllTasks(result);
        }
        update()
        const id = setInterval(update, 3000)
        return () => {
            clearInterval(id)
        }
    }, [])

    useEffect(() => {
        const activeTask = new Map<string, ExecBatchYakScriptResult[]>();
        ipcRenderer.on(`${props.token}-error`, async (e, exception) => {
            if (`${exception}`.includes("Cancelled on client")) {
                return
            }
            console.info("call exception")
            console.info(exception)
        })

        ipcRenderer.on(`${props.token}-data`, async (e, data: ExecBatchYakScriptResult) => {
            // 处理进度信息
            if (data.ProgressMessage) {
                if (!!props.setPercent) {
                    props.setPercent(data.ProgressPercent || 0)
                }
                setTotal(data.ProgressTotal || 0)
                setFinished(data.ProgressCount || 0)
                return
            }

            // 处理其他任务信息
            const taskId: string = data.TaskId || "";
            if (taskId === "") return

            // 缓存内容
            let activeResult = activeTask.get(taskId);
            if (!activeResult) activeResult = []
            activeResult.push(data)
            activeTask.set(taskId, activeResult)
            // 缓存全部
            let allresult = allPluginTasks.current.get(taskId);
            if (!allresult) allresult = []
            allresult.push(data)
            allPluginTasks.current.set(taskId, allresult)

            // 设置状态
            if (data.Status === "end") {
                activeTask.delete(taskId)
                return
            }

            // 看一下输出结果
            // if (data.Result && data.Result.IsMessage) {
            //     console.info(new Buffer(data.Result.Message).toString())
            // }
        })

        let cached = "";
        const syncActiveTask = () => {
            if (activeTask.size <= 0) setActiveTask([]);
            if (activeTask.size <= 0 && allPluginTasks.current.size <= 0) return

            const result: BatchTask[] = [];
            const tasks: string[] = [];
            activeTask.forEach(value => {
                if (value.length <= 0) return

                const first = value[0];
                const task = {
                    Target: first.Target || "",
                    ExtraParam: first.ExtraParams || [],
                    PoC: first.PoC,
                    TaskId: first.TaskId,
                    CreatedAt: first.Timestamp,
                } as BatchTask;
                task.Results = value;
                result.push(task)
                tasks.push(`${value.length}` + task.TaskId)
            })
            const allResult: BatchTask[] = [];
            allPluginTasks.current.forEach(value => {
                if (value.length <= 0) return

                const task = {
                    Target: value[0].Target || "",
                    ExtraParam: value[0].ExtraParams || [],
                    PoC: value[0].PoC,
                    TaskId: value[0].TaskId,
                    CreatedAt: value[0].Timestamp,
                } as BatchTask;
                task.Results = value;
                allResult.push(task)
            })

            const oldAllResult: BatchTask[] = []
            allTasksMap.forEach(value => oldAllResult.push(value))
            if (JSON.stringify(allResult) !== JSON.stringify(oldAllResult)) {
                allResult.forEach((value) => allTasksMap.set(value.TaskId, value))
            }

            const tasksRaw = tasks.sort().join("|")
            if (tasksRaw !== cached) {
                cached = tasksRaw
                setActiveTask(result)
            }
        }

        let id = setInterval(syncActiveTask, 1000);
        return () => {
            ipcRenderer.removeAllListeners(`${props.token}-data`)
            ipcRenderer.removeAllListeners(`${props.token}-end`)
            ipcRenderer.removeAllListeners(`${props.token}-error`)
            allTasksMap.clear()
            setAllTasks([])
            clearInterval(id);
        }
    }, [props.token])

    return <div style={{height: "100%", overflow: "auto"}}>
        <Tabs
            className="exec-result-body"
            onChange={setActiveKey}
            activeKey={activeKey}
            tabBarExtraContent={{
                right: activeKey !== "executing" && <div style={{width: 140}}>
                    <span style={{display: "inline-block", height: 22, marginRight: 5}}>只展示命中项</span>
                    <Switch checked={checked} onChange={checked => setChecked(checked)}></Switch>
                </div>
            }}
        >
            <Tabs.TabPane tab={"执行中的任务"} key={"executing"}>
                <List<BatchTask>
                    style={{height: '100%', overflow: "auto"}}
                    pagination={false}
                    dataSource={activeTask.sort((a, b) => a.TaskId.localeCompare(b.TaskId))}
                    rowKey={(item) => item.TaskId}
                    renderItem={(task: BatchTask) => {
                        const res = task.Results[task.Results.length - 1];
                        return (
                            <Card
                                bordered={false}
                                style={{marginBottom: 8, width: "100%"}}
                                size={"small"}
                                title={
                                    <Row gutter={8}>
                                        <Col span={20}>
                                            <Text
                                                ellipsis={{tooltip: true}}>{task.Target + " / " + task.PoC.ScriptName}</Text>
                                        </Col>
                                        <Col span={4}>
                                            <Timer fromTimestamp={task.CreatedAt} executing={!!props.executing}/>
                                        </Col>
                                    </Row>
                                }>
                                <div style={{width: "100%"}}>
                                    <ExecResultsViewer
                                        results={task.Results.map(i => i.Result).filter(i => !!i) as ExecResult[]}
                                        oneLine={true}
                                    />
                                </div>
                            </Card>
                        )
                    }}
                />
            </Tabs.TabPane>
            <Tabs.TabPane tab={`历史任务列表[${allTasks.length}]`} key={"results"} forceRender={true}>
                <BatchTaskViewer tasks={allTasks} checked={checked}/>
            </Tabs.TabPane>
        </Tabs>
    </div>
}