Merge pull request #1628 from gchq/feature/BAI-1459-create-a-new-file-scan-connector-for-modelscan

Feature/bai 1459 create a new file scan connector for modelscan

Author: PE39806

Diff for: backend/config/default.cjs

@@ -148,6 +148,12 @@ module.exports = {
       host: '127.0.0.1',
       port: 3310,
     },
+
+    modelscan: {
+      protocol: 'http',
+      host: '127.0.0.1',
+      port: 3311,
+    },
   },
 
   // These settings are PUBLIC and shared with the UI

Diff for: backend/config/docker_compose.cjs

@@ -65,11 +65,15 @@ module.exports = {
     clamdscan: {
       host: 'clamd',
     },
+
+    modelscan: {
+      host: 'modelscan',
+    },
   },
 
   connectors: {
     fileScanners: {
-      kinds: ['clamAV'],
+      kinds: ['clamAV', 'modelScan'],
     },
   },
 }

Diff for: backend/src/clients/modelScan.ts

@@ -0,0 +1,93 @@
+import fetch, { Response } from 'node-fetch'
+
+import config from '../utils/config.js'
+import { BadReq, InternalError } from '../utils/error.js'
+
+interface ModelScanInfoResponse {
+  apiName: string
+  apiVersion: string
+  scannerName: string
+  modelscanVersion: string
+}
+
+interface ModelScanResponse {
+  summary: {
+    total_issues: number
+    total_issues_by_severity: {
+      LOW: number
+      MEDIUM: number
+      HIGH: number
+      CRITICAL: number
+    }
+    input_path: string
+    absolute_path: string
+    modelscan_version: string
+    timestamp: string
+    scanned: {
+      total_scanned: number
+      scanned_files: string[]
+    }
+    skipped: {
+      total_skipped: number
+      skipped_files: string[]
+    }
+  }
+  issues: [
+    {
+      description: string
+      operator: string
+      module: string
+      source: string
+      scanner: string
+      severity: string
+    },
+  ]
+  // TODO: currently unknown what this might look like
+  errors: object[]
+}
+
+export async function getModelScanInfo() {
+  const url = `${config.avScanning.modelscan.protocol}://${config.avScanning.modelscan.host}:${config.avScanning.modelscan.port}`
+  let res: Response
+
+  try {
+    res = await fetch(`${url}/info`, {
+      method: 'GET',
+      headers: { 'Content-Type': 'application/json' },
+    })
+  } catch (err) {
+    throw InternalError('Unable to communicate with the ModelScan service.', { err })
+  }
+  if (!res.ok) {
+    throw BadReq('Unrecognised response returned by the ModelScan service.')
+  }
+
+  return (await res.json()) as ModelScanInfoResponse
+}
+
+export async function scanFile(file: Blob, file_name: string) {
+  const url = `${config.avScanning.modelscan.protocol}://${config.avScanning.modelscan.host}:${config.avScanning.modelscan.port}`
+  let res: Response
+
+  try {
+    const formData = new FormData()
+    formData.append('in_file', file, file_name)
+
+    res = await fetch(`${url}/scan/file`, {
+      method: 'POST',
+      headers: {
+        accept: 'application/json',
+      },
+      body: formData,
+    })
+  } catch (err) {
+    throw InternalError('Unable to communicate with the ModelScan service.', { err })
+  }
+  if (!res.ok) {
+    throw BadReq('Unrecognised response returned by the ModelScan service.', {
+      body: JSON.stringify(await res.json()),
+    })
+  }
+
+  return (await res.json()) as ModelScanResponse
+}

Diff for: backend/src/connectors/fileScanning/clamAv.ts

@@ -25,7 +25,7 @@ export class ClamAvFileScanningConnector extends BaseFileScanningConnector {
       av = await new NodeClam().init({ clamdscan: config.avScanning.clamdscan })
     } catch (error) {
       throw ConfigurationError('Could not scan file as Clam AV is not running.', {
-        clamAvConfig: config.avScanning,
+        clamAvConfig: config.avScanning.clamdscan,
       })
     }
   }
@@ -35,7 +35,7 @@ export class ClamAvFileScanningConnector extends BaseFileScanningConnector {
       throw ConfigurationError(
         'Clam AV does not look like it is running. Check that it has been correctly initialised by calling the init function.',
         {
-          clamAvConfig: config.avScanning,
+          clamAvConfig: config.avScanning.clamdscan,
         },
       )
     }

Diff for: backend/src/connectors/fileScanning/index.ts

@@ -2,10 +2,12 @@ import config from '../../utils/config.js'
 import { ConfigurationError } from '../../utils/error.js'
 import { BaseFileScanningConnector } from './Base.js'
 import { ClamAvFileScanningConnector } from './clamAv.js'
+import { ModelScanFileScanningConnector } from './modelScan.js'
 import { FileScanningWrapper } from './wrapper.js'
 
 export const FileScanKind = {
   ClamAv: 'clamAV',
+  ModelScan: 'modelScan',
 } as const
 export type FileScanKindKeys = (typeof FileScanKind)[keyof typeof FileScanKind]
 
@@ -26,6 +28,15 @@ export function runFileScanners(cache = true) {
           throw ConfigurationError('Could not configure or initialise Clam AV')
         }
         break
+      case FileScanKind.ModelScan:
+        try {
+          const scanner = new ModelScanFileScanningConnector()
+          await scanner.ping()
+          fileScanConnectors.push(scanner)
+        } catch (error) {
+          throw ConfigurationError('Could not configure or initialise ModelScan')
+        }
+        break
       default:
         throw ConfigurationError(`'${fileScanner}' is not a valid file scanning kind.`, {
           validKinds: Object.values(FileScanKind),

Diff for: backend/src/connectors/fileScanning/modelScan.ts

@@ -0,0 +1,76 @@
+import { Response } from 'node-fetch'
+import { Readable } from 'stream'
+
+import { getModelScanInfo, scanFile } from '../../clients/modelScan.js'
+import { getObjectStream } from '../../clients/s3.js'
+import { FileInterfaceDoc, ScanState } from '../../models/File.js'
+import log from '../../services/log.js'
+import config from '../../utils/config.js'
+import { ConfigurationError } from '../../utils/error.js'
+import { BaseFileScanningConnector, FileScanResult } from './Base.js'
+
+export const modelScanToolName = 'ModelScan'
+
+export class ModelScanFileScanningConnector extends BaseFileScanningConnector {
+  constructor() {
+    super()
+  }
+
+  info() {
+    return [modelScanToolName]
+  }
+
+  async ping() {
+    try {
+      // discard the results as we only want to know if the endpoint is reachable
+      await getModelScanInfo()
+    } catch (error) {
+      throw ConfigurationError(
+        'ModelScan does not look like it is running. Check that the service configuration is correct.',
+        {
+          modelScanConfig: config.avScanning.modelscan,
+        },
+      )
+    }
+  }
+
+  async scan(file: FileInterfaceDoc): Promise<FileScanResult[]> {
+    this.ping()
+
+    const s3Stream = (await getObjectStream(file.bucket, file.path)).Body as Readable
+    try {
+      // TODO: see if it's possible to directly send the Readable stream rather than a blob
+      const fileBlob = await new Response(s3Stream).blob()
+      const scanResults = await scanFile(fileBlob, file.name)
+
+      const issues = scanResults.summary.total_issues
+      const isInfected = issues > 0
+      const viruses: string[] = []
+      if (isInfected) {
+        for (const issue of scanResults.issues) {
+          viruses.push(`${issue.severity}: ${issue.description}. ${issue.scanner}`)
+        }
+      }
+      log.info(
+        { modelId: file.modelId, fileId: file._id, name: file.name, result: { isInfected, viruses } },
+        'Scan complete.',
+      )
+      return [
+        {
+          toolName: modelScanToolName,
+          state: ScanState.Complete,
+          isInfected,
+          viruses,
+        },
+      ]
+    } catch (error) {
+      log.error({ error, modelId: file.modelId, fileId: file._id, name: file.name }, 'Scan errored.')
+      return [
+        {
+          toolName: modelScanToolName,
+          state: ScanState.Error,
+        },
+      ]
+    }
+  }
+}

Diff for: backend/src/services/file.ts

@@ -59,9 +59,9 @@ async function updateFileWithResults(_id: Schema.Types.ObjectId, results: FileSc
     )
     if (updateExistingResult.modifiedCount === 0) {
       await FileModel.updateOne(
-        { _id },
+        { _id, avScan: { $exists: true } },
         {
-          $set: { avScan: { toolName: result.toolName, state: result.state } },
+          $push: { avScan: { toolName: result.toolName, state: result.state } },
         },
       )
     }

Diff for: backend/src/utils/config.ts

@@ -140,6 +140,12 @@ export interface Config {
       host: string
       port: number
     }
+
+    modelscan: {
+      protocol: string
+      host: string
+      port: number
+    }
   }
 
   modelMirror: {

Diff for: backend/test/clients/__snapshots__/modelScan.spec.ts.snap

@@ -0,0 +1,41 @@
+// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
+
+exports[`clients > modelScan > getModelScanInfo > success 1`] = `
+[
+  [
+    "undefined://undefined:undefined/info",
+    {
+      "headers": {
+        "Content-Type": "application/json",
+      },
+      "method": "GET",
+    },
+  ],
+]
+`;
+
+exports[`clients > modelScan > scanFile > success 1`] = `
+[
+  [
+    "undefined://undefined:undefined/scan/file",
+    {
+      "body": FormData {
+        Symbol(state): [
+          {
+            "name": "in_file",
+            "value": File {
+              Symbol(kHandle): Blob {},
+              Symbol(kLength): 0,
+              Symbol(kType): "application/x-hdf5",
+            },
+          },
+        ],
+      },
+      "headers": {
+        "accept": "application/json",
+      },
+      "method": "POST",
+    },
+  ],
+]
+`;