import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query'; import { deleteC2Config, executeC2, getC2Config, getC2Tasks, importC2, listCallbackHistory, listCallbacks, putC2Config, testC2Config, } from '@/api/c2'; import type { C2ConfigInput, C2ExecuteInput, C2ImportInput, C2TasksResponse, } from '@/api/types'; function c2ConfigKey(engagementId: number) { return ['c2-config', engagementId] as const; } function c2CallbacksKey(engagementId: number) { return ['c2-callbacks', engagementId] as const; } function c2TasksKey(simulationId: number) { return ['c2-tasks', simulationId] as const; } function c2HistoryKey(engagementId: number, callbackDisplayId: number, page: number, pageSize: number) { return ['c2-history', engagementId, callbackDisplayId, page, pageSize] as const; } function simulationKey(id: number) { return ['simulations', id] as const; } export function useC2Config(engagementId: number | undefined, options?: { enabled?: boolean }) { const enabled = typeof engagementId === 'number' && !Number.isNaN(engagementId) && (options?.enabled !== false); return useQuery({ queryKey: engagementId ? c2ConfigKey(engagementId) : ['c2-config', 'none'], queryFn: () => getC2Config(engagementId as number), enabled, }); } export function useUpdateC2Config(engagementId: number) { const qc = useQueryClient(); return useMutation({ mutationFn: (input: C2ConfigInput) => putC2Config(engagementId, input), onSuccess: () => qc.invalidateQueries({ queryKey: c2ConfigKey(engagementId) }), }); } export function useDeleteC2Config(engagementId: number) { const qc = useQueryClient(); return useMutation({ mutationFn: () => deleteC2Config(engagementId), onSuccess: () => qc.invalidateQueries({ queryKey: c2ConfigKey(engagementId) }), }); } export function useTestC2Config(engagementId: number) { return useMutation({ mutationFn: () => testC2Config(engagementId), }); } export function useC2Callbacks(engagementId: number | undefined, options?: { enabled?: boolean }) { const enabled = typeof engagementId === 'number' && !Number.isNaN(engagementId) && (options?.enabled !== false); return useQuery({ queryKey: engagementId ? c2CallbacksKey(engagementId) : ['c2-callbacks', 'none'], queryFn: () => listCallbacks(engagementId as number), enabled, }); } export function useC2Tasks(simulationId: number | undefined, options?: { enabled?: boolean }) { const enabled = typeof simulationId === 'number' && !Number.isNaN(simulationId) && (options?.enabled !== false); return useQuery({ queryKey: simulationId ? c2TasksKey(simulationId) : ['c2-tasks', 'none'], queryFn: () => getC2Tasks(simulationId as number), enabled, // Poll every 2500 ms while any task is incomplete; stop when all done. refetchInterval: (query: { state: { data?: C2TasksResponse } }) => query.state.data?.tasks?.some((t) => !t.completed) ? 2500 : false, }); } export function useExecuteC2(simulationId: number, engagementId: number) { const qc = useQueryClient(); return useMutation({ mutationFn: (input: C2ExecuteInput) => executeC2(simulationId, input), onSuccess: () => { qc.invalidateQueries({ queryKey: simulationKey(simulationId) }); qc.invalidateQueries({ queryKey: ['engagements', engagementId, 'simulations'] }); qc.invalidateQueries({ queryKey: c2TasksKey(simulationId) }); }, }); } export function useC2CallbackHistory( engagementId: number | undefined, callbackDisplayId: number | null, options?: { page?: number; pageSize?: number; enabled?: boolean }, ) { const page = options?.page ?? 1; const pageSize = options?.pageSize ?? 25; const enabled = typeof engagementId === 'number' && !Number.isNaN(engagementId) && callbackDisplayId !== null && (options?.enabled !== false); return useQuery({ queryKey: engagementId && callbackDisplayId !== null ? c2HistoryKey(engagementId, callbackDisplayId, page, pageSize) : ['c2-history', 'none'], queryFn: () => listCallbackHistory(engagementId as number, callbackDisplayId as number, { page, pageSize, }), enabled, }); } export function useImportC2(simulationId: number) { const qc = useQueryClient(); return useMutation({ mutationFn: (input: C2ImportInput) => importC2(simulationId, input), onSuccess: () => { qc.invalidateQueries({ queryKey: c2TasksKey(simulationId) }); qc.invalidateQueries({ queryKey: simulationKey(simulationId) }); }, }); }