[Security Solution][Endpoint] Cleanup and improvements to run_endpoint_agent.js CLI tool (#155730)

## Summary

As a follow up to PR #155455 : 

- `run_endpoint_agent` CLI: Removed internal code that handles the
creation of a VM and the enrollment of the agent with fleet and replaced
it with use of methods now found in `endpoint_host_services`
- created new service for managing agent download cache on local system
- enhanced `run_endpoint_agent` to use cached version of agent download,
and thus increase the performance of this 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,8 +53,14 @@ 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,
@@ -56,15 +68,33 @@ export const createAndEnrollEndpointHost = async ({
       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;

x-pack/plugins/security_solution/scripts/endpoint/endpoint_agent_runner/elastic_endpoint.ts

@@ -7,7 +7,6 @@
 
 import { userInfo } from 'os';
 import execa from 'execa';
-import nodeFetch from 'node-fetch';
 import {
   AGENT_POLICY_SAVED_OBJECT_TYPE,
   packagePolicyRouteService,
@@ -15,35 +14,14 @@ import {
   type UpdatePackagePolicy,
 } from '@kbn/fleet-plugin/common';
 import chalk from 'chalk';
+import { createAndEnrollEndpointHost } from '../common/endpoint_host_services';
 import { getEndpointPackageInfo } from '../../../common/endpoint/utils/package';
 import { indexFleetEndpointPolicy } from '../../../common/endpoint/data_loaders/index_fleet_endpoint_policy';
-import {
-  fetchAgentPolicyEnrollmentKey,
-  fetchAgentPolicyList,
-  fetchFleetServerUrl,
-  waitForHostToEnroll,
-} from '../common/fleet_services';
+import { fetchAgentPolicyList } from '../common/fleet_services';
 import { getRuntimeServices } from './runtime';
 import { type PolicyData, ProtectionModes } from '../../../common/endpoint/types';
 import { dump } from './utils';
 
-interface ElasticArtifactSearchResponse {
-  manifest: {
-    'last-update-time': string;
-    'seconds-since-last-update': number;
-  };
-  packages: {
-    [packageFileName: string]: {
-      architecture: string;
-      os: string[];
-      type: string;
-      asc_url: string;
-      sha_url: string;
-      url: string;
-    };
-  };
-}
-
 export const enrollEndpointHost = async (): Promise<string | undefined> => {
   let vmName;
   const {
@@ -61,81 +39,26 @@ export const enrollEndpointHost = async (): Promise<string | undefined> => {
     const policyId: string = policy || (await getOrCreateAgentPolicyId());
 
     if (!policyId) {
-      throw new Error(`No valid policy id provide or unable to create it`);
+      throw new Error(`No valid policy id provided or unable to create it`);
     }
 
     if (!version) {
       throw new Error(`No 'version' specified`);
     }
 
-    const [fleetServerHostUrl, enrollmentToken] = await Promise.all([
-      fetchFleetServerUrl(kbnClient),
-      fetchAgentPolicyEnrollmentKey(kbnClient, policyId),
-    ]);
-
-    if (!fleetServerHostUrl) {
-      throw new Error(`Fleet setting does not have a Fleet Server host defined!`);
-    }
-
-    if (!enrollmentToken) {
-      throw new Error(`No API enrollment key found for policy id [${policyId}]`);
-    }
-
     vmName = `${username}-dev-${uniqueId}`;
 
     log.info(`Creating VM named: ${vmName}`);
 
-    await execa.command(`multipass launch --name ${vmName} --disk 8G`);
-
-    log.verbose(await execa('multipass', ['info', vmName]));
-
-    const agentDownloadUrl = await getAgentDownloadUrl(version);
-    const agentDownloadedFile = agentDownloadUrl.substring(agentDownloadUrl.lastIndexOf('/') + 1);
-    const vmDirName = agentDownloadedFile.replace(/\.tar\.gz$/, '');
-
-    log.info(`Downloading and installing agent`);
-    log.verbose(`Agent download:\n    ${agentDownloadUrl}`);
-
-    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',
-
-      vmName,
-
-      '--working-directory',
-      `/home/ubuntu/${vmDirName}`,
-
-      '--',
-
-      'sudo',
-
-      './elastic-agent',
-
-      'install',
-
-      '--insecure',
-
-      '--force',
-
-      '--url',
-      fleetServerHostUrl,
-
-      '--enrollment-token',
-      enrollmentToken,
-    ];
-
-    log.info(`Enrolling elastic agent with Fleet`);
-    log.verbose(`Command: multipass ${agentInstallArguments.join(' ')}`);
-
-    await execa(`multipass`, agentInstallArguments);
-
-    log.info(`Waiting for Agent to check-in with Fleet`);
-    await waitForHostToEnroll(kbnClient, vmName);
+    await createAndEnrollEndpointHost({
+      kbnClient,
+      log,
+      hostname: vmName,
+      agentPolicyId: policyId,
+      version,
+      useClosestVersionMatch: false,
+      disk: '8G',
+    });
 
     log.info(`VM created using Multipass.
     VM Name: ${vmName}
@@ -155,36 +78,6 @@ export const enrollEndpointHost = async (): Promise<string | undefined> => {
   return vmName;
 };
 
-const getAgentDownloadUrl = async (version: string): Promise<string> => {
-  const { log } = getRuntimeServices();
-  const downloadArch =
-    { arm64: 'arm64', x64: 'x86_64' }[process.arch] ?? `UNSUPPORTED_ARCHITECTURE_${process.arch}`;
-  const agentFile = `elastic-agent-${version}-linux-${downloadArch}.tar.gz`;
-  const artifactSearchUrl = `https://artifacts-api.elastic.co/v1/search/${version}/${agentFile}`;
-
-  log.verbose(`Retrieving elastic agent download URL from:\n    ${artifactSearchUrl}`);
-
-  const searchResult: ElasticArtifactSearchResponse = await nodeFetch(artifactSearchUrl).then(
-    (response) => {
-      if (!response.ok) {
-        throw new Error(
-          `Failed to search elastic's artifact repository: ${response.statusText} (HTTP ${response.status})`
-        );
-      }
-
-      return response.json();
-    }
-  );
-
-  log.verbose(searchResult);
-
-  if (!searchResult.packages[agentFile]) {
-    throw new Error(`Unable to find an Agent download URL for version [${version}]`);
-  }
-
-  return searchResult.packages[agentFile].url;
-};
-
 const getOrCreateAgentPolicyId = async (): Promise<string> => {
   const { kbnClient, log } = getRuntimeServices();
   const username = userInfo().username.toLowerCase();