[Security Solution][Endpoint] Cleanup and improvements to run_endpoint_agent.js CLI tool

x-pack/plugins/security_solution/public/management/cypress/support/data_loaders.ts

@@ -195,7 +195,12 @@ export const dataLoadersForRealEndpoints = (
       options: Omit<CreateAndEnrollEndpointHostOptions, 'log' | 'kbnClient'>
     ): Promise<CreateAndEnrollEndpointHostResponse> => {
       const { kbnClient, log } = await stackServicesPromise;
-      return createAndEnrollEndpointHost({ ...options, log, kbnClient }).then((newHost) => {
+      return createAndEnrollEndpointHost({
+        useClosestVersionMatch: true,
+        ...options,
+        log,
+        kbnClient,
+      }).then((newHost) => {
         return waitForEndpointToStreamData(kbnClient, newHost.agentId, 120000).then(() => {
           return newHost;
         });

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<AgentDownloadStorageSettings> {
+  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<void> {
+    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<DownloadedAgentInfo> {
+    // 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<Promise<unknown>> = [];
+    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<DownloadedAgentInfo & { url: string }> => {
+  const downloadedAgent = await agentDownloadsClient.downloadAndStore(agentDownloadUrl);
+
+  return {
+    url: agentDownloadUrl,
+    ...downloadedAgent,
+  };
+};
+
+/**
+ * Cleans up the old agent downloads on disk.
+ */
+export const cleanupDownloads = async (): ReturnType<AgentDownloadStorage['cleanupDownloads']> => {
+  return agentDownloadsClient.cleanupDownloads();
+};

x-pack/plugins/security_solution/scripts/endpoint/common/endpoint_host_services.ts

@@ -10,6 +10,8 @@ import type { KbnClient } from '@kbn/test';
 import type { ToolingLog } from '@kbn/tooling-log';
 import execa from 'execa';
 import assert from 'assert';
+import type { DownloadedAgentInfo } from './agent_downloads_service';
+import { cleanupDownloads, downloadAndStoreAgent } from './agent_downloads_service';
 import {
   fetchAgentPolicyEnrollmentKey,
   fetchFleetServerUrl,
@@ -28,6 +30,10 @@ export interface CreateAndEnrollEndpointHostOptions
   version?: string;
   /** The name for the host. Will also be the name of the VM */
   hostname?: string;
+  /** If `version` should be exact, or if this is `true`, then the closest version will be used. Defaults to `false` */
+  useClosestVersionMatch?: boolean;
+  /** If the local cache of agent downloads should be used. Defaults to `true` */
+  useCache?: boolean;
 }
 
 export interface CreateAndEnrollEndpointHostResponse {
@@ -47,24 +53,48 @@ export const createAndEnrollEndpointHost = async ({
   memory,
   hostname,
   version = kibanaPackageJson.version,
+  useClosestVersionMatch = false,
+  useCache = true,
 }: CreateAndEnrollEndpointHostOptions): Promise<CreateAndEnrollEndpointHostResponse> => {
-  const [vm, agentDownloadUrl, fleetServerUrl, enrollmentToken] = await Promise.all([
+  let cacheCleanupPromise: ReturnType<typeof cleanupDownloads> = Promise.resolve({
+    deleted: [],
+  });
+
+  const [vm, agentDownload, fleetServerUrl, enrollmentToken] = await Promise.all([
     createMultipassVm({
       vmName: hostname ?? `test-host-${Math.random().toString().substring(2, 6)}`,
       disk,
       cpus,
       memory,
     }),
 
-    getAgentDownloadUrl(version, true, log),
+    getAgentDownloadUrl(version, useClosestVersionMatch, log).then<{
+      url: string;
+      cache?: DownloadedAgentInfo;
+    }>((url) => {
+      if (useCache) {
+        cacheCleanupPromise = cleanupDownloads();
+
+        return downloadAndStoreAgent(url).then((cache) => {
+          return {
+            url,
+            cache,
+          };
+        });
+      }
+
+      return { url };
+    }),
 
     fetchFleetServerUrl(kbnClient),
 
     fetchAgentPolicyEnrollmentKey(kbnClient, agentPolicyId),
   ]);
 
+  log.verbose(await execa('multipass', ['info', vm.vmName]));
+
   // Some validations before we proceed
-  assert(agentDownloadUrl, 'Missing agent download URL');
+  assert(agentDownload.url, 'Missing agent download URL');
   assert(fleetServerUrl, 'Fleet server URL not set');
   assert(enrollmentToken, `No enrollment token for agent policy id [${agentPolicyId}]`);
 
@@ -76,11 +106,22 @@ export const createAndEnrollEndpointHost = async ({
     kbnClient,
     log,
     fleetServerUrl,
-    agentDownloadUrl,
+    agentDownloadUrl: agentDownload.url,
+    cachedAgentDownload: agentDownload.cache,
     enrollmentToken,
     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,
@@ -143,6 +184,7 @@ interface EnrollHostWithFleetOptions {
   log: ToolingLog;
   vmName: string;
   agentDownloadUrl: string;
+  cachedAgentDownload?: DownloadedAgentInfo;
   fleetServerUrl: string;
   enrollmentToken: string;
 }
@@ -153,16 +195,35 @@ const enrollHostWithFleet = async ({
   vmName,
   fleetServerUrl,
   agentDownloadUrl,
+  cachedAgentDownload,
   enrollmentToken,
 }: EnrollHostWithFleetOptions): Promise<{ agentId: string }> => {
   const agentDownloadedFile = agentDownloadUrl.substring(agentDownloadUrl.lastIndexOf('/') + 1);
   const vmDirName = agentDownloadedFile.replace(/\.tar\.gz$/, '');
 
-  await execa.command(
-    `multipass exec ${vmName} -- curl -L ${agentDownloadUrl} -o ${agentDownloadedFile}`
-  );
-  await execa.command(`multipass exec ${vmName} -- tar -zxf ${agentDownloadedFile}`);
-  await execa.command(`multipass exec ${vmName} -- rm -f ${agentDownloadedFile}`);
+  if (cachedAgentDownload) {
+    log.verbose(
+      `Installing agent on host using cached download from [${cachedAgentDownload.fullFilePath}]`
+    );
+
+    // mount local folder on VM
+    await execa.command(
+      `multipass mount ${cachedAgentDownload.directory} ${vmName}:~/_agent_downloads`
+    );
+    await execa.command(
+      `multipass exec ${vmName} -- tar -zxf _agent_downloads/${cachedAgentDownload.filename}`
+    );
+    await execa.command(`multipass unmount ${vmName}:~/_agent_downloads`);
+  } else {
+    log.verbose(`downloading and installing agent from URL [${agentDownloadUrl}]`);
+
+    // download into VM
+    await execa.command(
+      `multipass exec ${vmName} -- curl -L ${agentDownloadUrl} -o ${agentDownloadedFile}`
+    );
+    await execa.command(`multipass exec ${vmName} -- tar -zxf ${agentDownloadedFile}`);
+    await execa.command(`multipass exec ${vmName} -- rm -f ${agentDownloadedFile}`);
+  }
 
   const agentInstallArguments = [
     'exec',

x-pack/plugins/security_solution/scripts/endpoint/common/settings_storage.ts

@@ -28,10 +28,8 @@ export class SettingsStorage<TSettingsDef extends object = object> {
   private dirExists: boolean = false;
 
   constructor(fileName: string, options: SettingStorageOptions<TSettingsDef> = {}) {
-    const {
-      directory = join(homedir(), '.kibanaSecuritySolutionCliTools'),
-      defaultSettings = {} as TSettingsDef,
-    } = options;
+    const { directory = SettingsStorage.getDirectory(), defaultSettings = {} as TSettingsDef } =
+      options;
 
     this.options = {
       directory,
@@ -41,7 +39,17 @@ export class SettingsStorage<TSettingsDef extends object = object> {
     this.settingsFileFullPath = join(this.options.directory, fileName);
   }
 
-  private async ensureExists(): Promise<void> {
+  /** Returns the default path to the directory where settings are saved to. */
+  public static getDirectory(): string {
+    return join(homedir(), '.kibanaSecuritySolutionCliTools');
+  }
+
+  /** Build a path using the root directory of where settings are saved */
+  protected buildPath(path: string): string {
+    return join(this.options.directory, path);
+  }
+
+  protected async ensureExists(): Promise<void> {
     if (!this.dirExists) {
       await mkdir(this.options.directory, { recursive: true });
       this.dirExists = true;