feat(jobs): restart scans

.vscode/settings.json

@@ -2,5 +2,5 @@
     "editor.codeActionsOnSave": {
         "source.fixAll.eslint": true
     },
-    "jest.autoRun": {}
+    "jest.autoRun": {},
 }

docker-compose.yml

@@ -7,7 +7,7 @@ services:
       target: dev
       context: ./
 
-    image: ninjeneer/vuln-request-service
+    image: ninjeneer/vuln-scanner-request-service
     ports:
       - 3000:3000
     volumes:
@@ -21,14 +21,26 @@ services:
       target: dev
       context: ./
 
-    image: ninjeneer/vuln-report-service
+    image: ninjeneer/vuln-scanner-report-service
     ports:
       - 3001:3000
     volumes:
       - .:/app
     networks:
       - vulnscanner
 
+  jobs-service:
+    build:
+      dockerfile: ./jobs.dockerfile
+      target: dev
+      context: ./
+
+    image: ninjeneer/vuln-scanner-jobs-service
+    volumes:
+      - .:/app
+    networks:
+      - vulnscanner
+
 networks:
   vulnscanner:
     name: vulnscanner

jobs.dockerfile

@@ -0,0 +1,21 @@
+FROM node:alpine as builder
+WORKDIR /app
+COPY --chown=node:node package.json yarn.lock ./
+RUN yarn install --frozen-lockfile 
+COPY --chown=node:node . .
+RUN yarn build
+
+
+# Dev stage
+FROM builder as dev
+WORKDIR /app
+EXPOSE 3000
+CMD ["yarn", "dev:jobs"]
+
+
+# Prod stage
+FROM node:alpine as prod
+WORKDIR /app
+COPY --from=builder /app/node_modules ./dist/node_modules
+COPY --from=builder /app/dist/src ./dist/src
+CMD ["node", "dist/src/services/jobs/index.js"]

package.json

@@ -4,10 +4,9 @@
   "description": "",
   "main": "src/index.ts",
   "scripts": {
-    "dev:docker": "docker-compose -f docker-compose.yml up",
-    "dev:docker-down": "docker-compose -f docker-compose.yml down",
     "dev:requests": "nodemon src/services/requests/server.ts",
     "dev:reports": "nodemon src/services/reports/server.ts",
+    "dev:jobs": "nodemon src/services/jobs/index.ts",
     "build": "swc src -d dist/src",
     "test": "jest test"
   },
@@ -20,6 +19,7 @@
     "@swc/register": "^0.1.10",
     "@types/jest": "^29.2.4",
     "@types/jsonwebtoken": "^9.0.1",
+    "@types/node-cron": "^3.0.7",
     "@types/uuid": "^9.0.0",
     "jest": "^29.3.1",
     "nodemon": "^2.0.20",
@@ -33,6 +33,7 @@
     "fastify": "^4.10.2",
     "jsonwebtoken": "^9.0.0",
     "mongoose": "^7.0.3",
+    "node-cron": "^3.0.2",
     "uuid": "^9.0.0",
     "zod": "^3.19.1"
   }

src/models/probe.ts

@@ -3,6 +3,7 @@ export type Probe = {
     status: ProbeStatus;
     scanId: string;
     name: string
+    settings: Record<string, any>
 }
 
 export enum ProbeStatus {

src/services/jobs/index.ts

@@ -0,0 +1,3 @@
+import initializeJobs from './scansStarter'
+
+initializeJobs()

src/services/jobs/scansStarter.ts

@@ -0,0 +1,55 @@
+import { getScan, getScansWithProbes, listenScans } from "../../storage/scan.storage";
+import cron from 'node-cron'
+import { restartScan } from "../requests/scanService";
+
+const cronMapping: Record<string, cron.ScheduledTask> = {}
+
+// Keep in memory scans periodicities to avoid restarting job on scan update
+const periodicityMapping = {}
+
+const init = async () => {
+    console.log(`[JOBS][SCAN][INIT] Initializating scan jobs...`)
+    const scans = await getScansWithProbes()
+    for (const scan of scans) {
+        if (scan.periodicity !== 'ONCE') {
+            const scanTask = cron.schedule(scan.periodicity, () => restartScan(scan))
+            cronMapping[scan.id] = scanTask
+            cronMapping[scan.id].start()
+        }
+        periodicityMapping[scan.id] = scan.periodicity
+    }
+    console.log(`[JOBS][SCAN][INIT] Initialized ${Object.values(cronMapping).length} scans`)
+
+    listenScans(async (change) => {
+        if (change.eventType === 'UPDATE' || change.eventType === 'INSERT') {
+            const isUpdate = change.eventType === 'UPDATE'
+            const changeScan = change.new
+            if (isUpdate && changeScan.periodicity === periodicityMapping[changeScan.id]) {
+                // The periodicity has not been edited
+                return
+            }
+
+            if (changeScan.periodicity === 'ONCE') {
+                periodicityMapping[changeScan.id] = changeScan.periodicity
+                if (isUpdate) {
+                    cronMapping[changeScan.id].stop()
+                    delete cronMapping[changeScan.id]
+                }
+                return
+            }
+
+            console.log(`[JOBS][SCAN][${change.eventType.toUpperCase()}] ${isUpdate ? 'Restarting' : 'Starting'} scan ${change.new.id} job...`)
+            const scan = await getScan(change.new.id)
+            if (isUpdate && cronMapping[scan.id]) {
+                cronMapping[scan.id].stop()
+            }
+            cronMapping[scan.id] = cron.schedule(scan.periodicity, () => restartScan(scan))
+            cronMapping[scan.id].start()
+            periodicityMapping[scan.id] = scan.periodicity
+            console.log({ old: periodicityMapping[scan.id], new: scan.periodicity })
+            console.log(`[JOBS][SCAN][${change.eventType.toUpperCase()}] ${isUpdate ? 'Restarted' : 'Started'} scan ${change.new.id} job`)
+        }
+    })
+}
+
+export default init

src/services/reports/probeResultService.ts

@@ -30,12 +30,14 @@ const onProbeResult = async (probeId: string, resultId: string): Promise<boolean
         throw new ProbeDoesNotExist();
     }
     await updateProbe(probeId, { status: ProbeStatus.FINISHED });
-    await createProbeResult({ probeId, resultId })
-
+
     const scan = await getScan(probe.scanId);
     if (!scan) {
         throw new ScanDoesNotExist(probe.scanId);
     }
+
+    await createProbeResult({ probeId, resultId, reportId: scan.currentReportId })
+
     if (isScanFinished(scan)) {
         console.log(`[RESULT][${scan.id}] Scan ${scan.id} is finished`)
 
@@ -51,7 +53,8 @@ const onProbeResult = async (probeId: string, resultId: string): Promise<boolean
             status: ScanStatus.FINISHED,
             notification: true,
             lastReportId: savedReport.id,
-            currentReportId: null
+            currentReportId: null,
+            lastRun: new Date()
         })
         console.log(`[REPORT][${scan.id}] Created report !`)
     }

src/services/reports/reportService.ts

@@ -1,7 +1,7 @@
 import { NoProbeResultsForReport, ReportDoesNotExist, ScanDoesNotExist } from "../../exceptions/exceptions";
 import { ProbeResult } from "../../models/probe";
 import { Report } from "../../models/report";
-import { getProbeResultsByCurrentReportId, getProbeResultsByLastReportId } from "../../storage/probe.storage";
+import { getProbeResultsByReportId } from "../../storage/probe.storage";
 import { getScan } from "../../storage/scan.storage";
 import { getResultsByIds } from "../../storage/mongo/mongoProbe.storage";
 import * as reportMongoStorage from "../../storage/mongo/mongoReport.storage"
@@ -20,8 +20,7 @@ export const buildReport = async (reportId: string, isRebuild?: boolean): Promis
         throw new ReportDoesNotExist(reportId);
     }
 
-    const getProbesMethod = isRebuild ? getProbeResultsByLastReportId : getProbeResultsByCurrentReportId
-    const probeResults = await getProbesMethod(reportId);
+    const probeResults = await getProbeResultsByReportId(reportId);
     if (!probeResults) {
         throw new NoProbeResultsForReport(reportId);
     }

src/services/requests/scanService.ts

@@ -1,26 +1,25 @@
 import { CreateScanRequest, UpdateScanRequest } from "../requests/validators/scanRequest";
 import { v4 as uuidv4 } from 'uuid';
 import * as scanStorage from "../../storage/scan.storage";
-import { ScanStatus, ScanRequestResponse } from "../../models/scan";
-import { ProbeStatus } from "../../models/probe";
+import { ScanStatus, ScanRequestResponse, ScanWithProbes } from "../../models/scan";
+import { Probe, ProbeStatus } from "../../models/probe";
 import { Scan } from '../../models/scan'
 import { Report, SupabaseReport } from '../../models/report'
 import { publishProbeRequest } from "../../storage/awsSqsQueue";
-import { saveReport } from "../../storage/mongo/mongoReport.storage";
 import { createReport } from "../../storage/report.storage";
 import { ScanDoesNotExist } from "../../exceptions/exceptions";
-import { deleteProbes } from "../../storage/probe.storage";
+import { deleteProbes, updateProbesByScanId } from "../../storage/probe.storage";
 
 
 export const requestScan = async (scanRequest: CreateScanRequest): Promise<ScanRequestResponse> => {
     const newScanId = uuidv4();
     console.log(`[REQUEST][${newScanId}] Received scan request ${newScanId}`)
 
     // Assing uids to probes
-    const probes = scanRequest.probes.map((probe) => {
+    const probes: Partial<Probe>[] = scanRequest.probes.map((probe) => {
         return {
             ...probe,
-            uid: uuidv4()
+            id: uuidv4()
         }
     })
 
@@ -40,18 +39,19 @@ export const requestScan = async (scanRequest: CreateScanRequest): Promise<ScanR
 
     console.log(`[REQUEST][${newScanId}] Saving probe start data...`)
     await scanStorage.saveProbesStartData(probes.map((probe) => ({
-        id: probe.uid,
+        id: probe.id,
         status: ProbeStatus.PENDING,
         scanId: newScanId,
-        name: probe.name
+        name: probe.name,
+        settings: probe.settings
     })))
     console.log(`[REQUEST][${newScanId}] Probe start data saved !`)
 
 
     console.log(`[REQUEST][${newScanId}] Publishing request to Queue...`)
     await publishProbeRequest(probes.map((probe) => ({
         context: {
-            id: probe.uid,
+            id: probe.id,
             name: probe.name,
             target: scanRequest.target
         },
@@ -97,7 +97,8 @@ export const updateScan = async (scanId: string, scanWithProbes: UpdateScanReque
             id: uuidv4(),
             status: ProbeStatus.PENDING,
             scanId,
-            name: probe.name
+            name: probe.name,
+            settings: probe.settings
         })))
         console.log(`[REQUEST][SCAN][UPDATE][${scanId}] Probes created successfully`)
     }
@@ -115,4 +116,24 @@ export const updateScan = async (scanId: string, scanWithProbes: UpdateScanReque
     const updatedScan = await scanStorage.updateScan(scanId, payload)
     console.log(`[REQUEST][SCAN][UPDATE][${scanId}] Scan updated`)
     return updatedScan
+}
+
+export const restartScan = async (scan: ScanWithProbes): Promise<void> => {
+    console.log(`[REQUEST][SCAN][RESTART][${scan.id}] Restarting scan...`)
+    const report = await setupScanNewReport(scan)
+    await scanStorage.updateScan(scan.id, { currentReportId: report.id})
+
+    console.log(`[REQUEST][SCAN][RESTART][${scan.id}] Updating probes...`)
+    await updateProbesByScanId(scan.id, { status: ProbeStatus.PENDING })
+
+    console.log(`[REQUEST][SCAN][RESTART][${scan.id}] Publishing request to Queue...`)
+    await publishProbeRequest(scan.probes.map((probe) => ({
+        context: {
+            id: probe.id,
+            name: probe.name,
+            target: scan.target
+        },
+        settings: probe.settings
+    })));
+    console.log(`[REQUEST][SCAN][RESTART][${scan.id}] Published request to Queue !`)
 }

src/services/requests/validators/scanRequest.ts

@@ -5,7 +5,7 @@ export const createScanRequest = z.object({
     periodicity: z.string().trim(),
     probes: z.array(z.object({
         name: z.string().trim(),
-        settings: z.any()
+        settings: z.object({}).optional()
     })),
     user_id: z.string()
 })

src/storage/dto/probe.dto.ts

@@ -7,4 +7,5 @@ export type ProbeUpdatePayload = {
 export type ProbeResultPayload = {
     probeId: string;
     resultId: string;
+    reportId: string;
 }

src/storage/dto/scan.dto.ts

@@ -5,4 +5,5 @@ export type ScanUpdatePayload = {
     notification: boolean
     lastReportId?: string;
     currentReportId?: string
+    lastRun?: Date
 }

src/storage/probe.storage.ts

@@ -7,6 +7,10 @@ export const updateProbe = async (id: string, data: ProbeUpdatePayload): Promise
     await supabaseClient.from('probes').update({ status: data.status }).eq('id', id)
 }
 
+export const updateProbesByScanId = async (scanId: string, payload: Partial<Probe>) => {
+    await supabaseClient.from('probes').update(payload).eq('scanId', scanId)
+}
+
 export const createProbeResult = async (data: ProbeResultPayload): Promise<SupabaseProbeResult> => {
     return (await supabaseClient.from('probes_results').insert(data).select().single()).data
 }
@@ -24,19 +28,11 @@ export const getProbe = async (probeId: string): Promise<Probe> => {
     return (await supabaseClient.from('probes').select('*').eq('id', probeId).single()).data as Probe
 }
 
-export const getProbeResultsByCurrentReportId = async (reportId: string): Promise<SupabaseProbeResult[]> => {
-    return (await supabaseClient
-        .from('probes_results')
-        .select('*, probes!inner(*, scans!inner(*))')
-        .eq('probes.scans.currentReportId', reportId)
-    ).data as SupabaseProbeResult[]
-}
-
-export const getProbeResultsByLastReportId = async (reportId: string): Promise<SupabaseProbeResult[]> => {
+export const getProbeResultsByReportId = async (reportId: string): Promise<SupabaseProbeResult[]> => {
     return (await supabaseClient
         .from('probes_results')
-        .select('*, probes!inner(*, scans!inner(*))')
-        .eq('probes.scans.lastReportId', reportId)
+        .select('*')
+        .eq('reportId', reportId)
     ).data as SupabaseProbeResult[]
 }
 

src/storage/scan.storage.ts

@@ -1,3 +1,4 @@
+import { RealtimePostgresChangesPayload } from '@supabase/supabase-js'
 import { Probe } from '../models/probe'
 import { Scan, ScanWithProbes } from '../models/scan'
 import { ScanUpdatePayload } from './dto/scan.dto'
@@ -17,10 +18,20 @@ export const getScan = async (scanId: string): Promise<ScanWithProbes> => {
     return (await supabaseClient.from('scans').select('*, probes(*)').eq('id', scanId).single()).data as ScanWithProbes
 }
 
+export const getScansWithProbes = async (): Promise<ScanWithProbes[]> => {
+    return (await supabaseClient.from('scans').select('*, probes(*)')).data
+}
+
 export const updateScan = async (id: string, data: Partial<ScanUpdatePayload>): Promise<Scan> => {
     const res = await supabaseClient.from('scans').update(data).eq('id', id).select().single()
     if (res.error) {
         console.log(res.error)
     }
     return res.data
+}
+
+export const listenScans = (onChange: (payload: RealtimePostgresChangesPayload<Scan>) => void) => {
+    supabaseClient.channel('scans')
+        .on('postgres_changes', { event: '*', schema: 'public', table: 'scans' }, onChange)
+        .subscribe()
 }

supabase/migrations/20230422115400_scans_restart.sql

@@ -0,0 +1,9 @@
+alter table "public"."probes" add column "settings" jsonb;
+
+alter table "public"."probes_results" add column "reportId" uuid;
+
+alter table "public"."scans" add column "lastRun" timestamp with time zone;
+
+alter table "public"."probes_results" add constraint "probes_results_reportId_fkey" FOREIGN KEY ("reportId") REFERENCES reports(id) ON DELETE CASCADE not valid;
+
+alter table "public"."probes_results" validate constraint "probes_results_reportId_fkey";