From 7c3c31e4461828c763c876077e0b59d58067593d Mon Sep 17 00:00:00 2001 From: Paul Tavares Date: Tue, 25 Apr 2023 10:31:14 -0400 Subject: [PATCH] Extracted agent file management to its own service --- .../common/agent_downloads_service.ts | 161 ++++++++++++++++++ .../endpoint/common/endpoint_host_services.ts | 128 ++------------ 2 files changed, 179 insertions(+), 110 deletions(-) create mode 100644 x-pack/plugins/security_solution/scripts/endpoint/common/agent_downloads_service.ts diff --git a/x-pack/plugins/security_solution/scripts/endpoint/common/agent_downloads_service.ts b/x-pack/plugins/security_solution/scripts/endpoint/common/agent_downloads_service.ts new file mode 100644 index 0000000000000..ed5d0296d61a4 --- /dev/null +++ b/x-pack/plugins/security_solution/scripts/endpoint/common/agent_downloads_service.ts @@ -0,0 +1,161 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import { mkdir, readdir, stat, unlink } from 'fs/promises'; +import { join } from 'path'; +import fs from 'fs'; +import nodeFetch from 'node-fetch'; +import { finished } from 'stream/promises'; +import { SettingsStorage } from './settings_storage'; + +export interface DownloadedAgentInfo { + filename: string; + directory: string; + fullFilePath: string; +} + +interface AgentDownloadStorageSettings { + /** + * Last time a cleanup was ran. Date in ISO format + */ + lastCleanup: string; + + /** + * The max file age in milliseconds. Defaults to 2 days + */ + maxFileAge: number; +} + +/** + * Class for managing Agent Downloads on the local disk + * @private + */ +class AgentDownloadStorage extends SettingsStorage { + private downloadsFolderExists = false; + private readonly downloadsDirName = 'agent_download_storage'; + private readonly downloadsDirFullPath: string; + + constructor() { + super('agent_download_storage_settings.json', { + defaultSettings: { + maxFileAge: 1.728e8, // 2 days + lastCleanup: new Date().toISOString(), + }, + }); + + this.downloadsDirFullPath = this.buildPath(this.downloadsDirName); + } + + protected async ensureExists(): Promise { + await super.ensureExists(); + + if (!this.downloadsFolderExists) { + await mkdir(this.downloadsDirFullPath, { recursive: true }); + this.downloadsFolderExists = true; + } + } + + public getPathsForUrl(agentDownloadUrl: string): DownloadedAgentInfo { + const filename = agentDownloadUrl.replace(/^https?:\/\//gi, '').replace(/\//g, '#'); + const directory = this.downloadsDirFullPath; + const fullFilePath = this.buildPath(join(this.downloadsDirName, filename)); + + return { + filename, + directory, + fullFilePath, + }; + } + + public async downloadAndStore(agentDownloadUrl: string): Promise { + // TODO: should we add "retry" attempts to file downloads? + + await this.ensureExists(); + + const newDownloadInfo = this.getPathsForUrl(agentDownloadUrl); + + // If download is already present on disk, then just return that info. No need to re-download it + if (fs.existsSync(newDownloadInfo.fullFilePath)) { + return newDownloadInfo; + } + + try { + const outputStream = fs.createWriteStream(newDownloadInfo.fullFilePath); + const { body } = await nodeFetch(agentDownloadUrl); + + await finished(body.pipe(outputStream)); + } catch (e) { + // Try to clean up download case it failed halfway through + await unlink(newDownloadInfo.fullFilePath); + + throw e; + } + + return newDownloadInfo; + } + + public async cleanupDownloads(): Promise<{ deleted: string[] }> { + const settings = await this.get(); + const maxAgeDate = new Date(); + const response: { deleted: string[] } = { deleted: [] }; + + maxAgeDate.setMilliseconds(settings.maxFileAge * -1); // `* -1` to set time back + + // If cleanup already happen within the file age, then nothing to do. Exit. + if (settings.lastCleanup > maxAgeDate.toISOString()) { + return response; + } + + await this.save({ + ...settings, + lastCleanup: new Date().toISOString(), + }); + + const deleteFilePromises: Array> = []; + const allFiles = await readdir(this.downloadsDirFullPath); + + for (const fileName of allFiles) { + const filePath = join(this.downloadsDirFullPath, fileName); + const fileStats = await stat(filePath); + + if (fileStats.isFile() && fileStats.birthtime < maxAgeDate) { + deleteFilePromises.push(unlink(filePath)); + response.deleted.push(filePath); + } + } + + await Promise.allSettled(deleteFilePromises); + + return response; + } +} + +const agentDownloadsClient = new AgentDownloadStorage(); + +/** + * Downloads the agent file provided via the input URL to a local folder on disk. If the file + * already exists on disk, then no download is actually done - the information about the cached + * version is returned instead + * @param agentDownloadUrl + */ +export const downloadAndStoreAgent = async ( + agentDownloadUrl: string +): Promise => { + const downloadedAgent = await agentDownloadsClient.downloadAndStore(agentDownloadUrl); + + return { + url: agentDownloadUrl, + ...downloadedAgent, + }; +}; + +/** + * Cleans up the old agent downloads on disk. + */ +export const cleanupDownloads = async (): ReturnType => { + return agentDownloadsClient.cleanupDownloads(); +}; diff --git a/x-pack/plugins/security_solution/scripts/endpoint/common/endpoint_host_services.ts b/x-pack/plugins/security_solution/scripts/endpoint/common/endpoint_host_services.ts index c107ed325a84a..5b249ee238436 100644 --- a/x-pack/plugins/security_solution/scripts/endpoint/common/endpoint_host_services.ts +++ b/x-pack/plugins/security_solution/scripts/endpoint/common/endpoint_host_services.ts @@ -10,11 +10,8 @@ import type { KbnClient } from '@kbn/test'; import type { ToolingLog } from '@kbn/tooling-log'; import execa from 'execa'; import assert from 'assert'; -import { mkdir, unlink, readdir, stat } from 'fs/promises'; -import { join } from 'path'; -import fs from 'fs'; -import { finished } from 'stream/promises'; -import nodeFetch from 'node-fetch'; +import type { DownloadedAgentInfo } from './agent_downloads_service'; +import { cleanupDownloads, downloadAndStoreAgent } from './agent_downloads_service'; import { fetchAgentPolicyEnrollmentKey, fetchFleetServerUrl, @@ -22,7 +19,6 @@ import { unEnrollFleetAgent, waitForHostToEnroll, } from './fleet_services'; -import { SettingsStorage } from './settings_storage'; export interface CreateAndEnrollEndpointHostOptions extends Pick { @@ -60,6 +56,10 @@ export const createAndEnrollEndpointHost = async ({ useClosestVersionMatch = false, useCache = true, }: CreateAndEnrollEndpointHostOptions): Promise => { + let cacheCleanupPromise: ReturnType = Promise.resolve({ + deleted: [], + }); + const [vm, agentDownload, fleetServerUrl, enrollmentToken] = await Promise.all([ createMultipassVm({ vmName: hostname ?? `test-host-${Math.random().toString().substring(2, 6)}`, @@ -73,9 +73,9 @@ export const createAndEnrollEndpointHost = async ({ cache?: DownloadedAgentInfo; }>((url) => { if (useCache) { - const agentDownloadClient = new AgentDownloadStorage(); + cacheCleanupPromise = cleanupDownloads(); - return agentDownloadClient.downloadAndStore(url).then((cache) => { + return downloadAndStoreAgent(url).then((cache) => { return { url, cache, @@ -112,6 +112,16 @@ export const createAndEnrollEndpointHost = async ({ vmName: vm.vmName, }); + await cacheCleanupPromise.then((results) => { + if (results.deleted.length > 0) { + log.verbose(`Agent Downloads cache directory was cleaned up and the following ${ + results.deleted.length + } were deleted: +${results.deleted.join('\n')} +`); + } + }); + return { hostname: vm.vmName, agentId, @@ -254,105 +264,3 @@ const enrollHostWithFleet = async ({ agentId: agent.id, }; }; - -interface DownloadedAgentInfo { - filename: string; - directory: string; - fullFilePath: string; -} - -class AgentDownloadStorage extends SettingsStorage { - private downloadsFolderExists = false; - private readonly downloadsDirName = 'agent_download_storage'; - private readonly downloadsDirFullPath: string; - - constructor() { - super('agent_download_storage_settings.json'); - - this.downloadsDirFullPath = this.buildPath(this.downloadsDirName); - } - - protected async ensureExists(): Promise { - await super.ensureExists(); - - if (!this.downloadsFolderExists) { - await mkdir(this.downloadsDirFullPath, { recursive: true }); - this.downloadsFolderExists = true; - } - } - - public getPathsForUrl(agentDownloadUrl: string): DownloadedAgentInfo { - const filename = agentDownloadUrl.replace(/^https?:\/\//gi, '').replace(/\//g, '#'); - const directory = this.downloadsDirFullPath; - const fullFilePath = this.buildPath(join(this.downloadsDirName, filename)); - - return { - filename, - directory, - fullFilePath, - }; - } - - public async downloadAndStore(agentDownloadUrl: string): Promise { - // TODO: should we add "retry" attempts to file downloads? - - await this.ensureExists(); - - const newDownloadInfo = this.getPathsForUrl(agentDownloadUrl); - - // If download is already present on disk, then just return that info. No need to re-download it - if (fs.existsSync(newDownloadInfo.fullFilePath)) { - return newDownloadInfo; - } - - try { - const outputStream = fs.createWriteStream(newDownloadInfo.fullFilePath); - const { body } = await nodeFetch(agentDownloadUrl); - - await finished(body.pipe(outputStream)); - } catch (e) { - // Try to clean up download case it failed halfway through - await unlink(newDownloadInfo.fullFilePath); - - throw e; - } - - return newDownloadInfo; - } - - public async cleanupDownloads(): Promise { - const deletedFiles: string[] = []; - const deleteFilePromises: Array> = []; - const maxAgeDate = new Date(); - - maxAgeDate.setMilliseconds(-1.728e8); // -2 days - - const allFiles = await readdir(this.downloadsDirFullPath); - - for (const fileName of allFiles) { - const filePath = join(this.downloadsDirFullPath, fileName); - const fileStats = await stat(filePath); - - if (fileStats.isFile() && fileStats.birthtime < maxAgeDate) { - deleteFilePromises.push(unlink(filePath)); - deletedFiles.push(filePath); - } - } - - await Promise.allSettled(deleteFilePromises); - - return deletedFiles; - } -} - -export const downloadAgent = async ( - agentDownloadUrl: string -): Promise => { - const agentDownloadClient = new AgentDownloadStorage(); - const downloadedAgent = await agentDownloadClient.downloadAndStore(agentDownloadUrl); - - return { - url: agentDownloadUrl, - ...downloadedAgent, - }; -};