[EDR Workflows] Improve agent downloader

x-pack/plugins/security_solution/scripts/endpoint/agent_downloader_cli/agent_downloader.test.ts

@@ -0,0 +1,203 @@
+/*
+ * 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 { getAgentDownloadUrl, getAgentFileName } from '../common/fleet_services';
+import { downloadAndStoreAgent } from '../common/agent_downloads_service';
+import type { ToolingLog } from '@kbn/tooling-log';
+import { agentDownloaderRunner } from './agent_downloader';
+import type { RunContext } from '@kbn/dev-cli-runner';
+
+jest.mock('../common/fleet_services');
+jest.mock('../common/agent_downloads_service');
+
+describe('agentDownloaderRunner', () => {
+  let log: ToolingLog;
+
+  beforeEach(() => {
+    log = {
+      info: jest.fn(),
+      error: jest.fn(),
+    } as unknown as ToolingLog;
+
+    jest.clearAllMocks();
+  });
+
+  const version = '8.15.0';
+  let closestMatch = false;
+  const url = 'http://example.com/agent.tar.gz';
+  const fileName = 'elastic-agent-8.15.0.tar.gz';
+
+  it('downloads and stores the specified version', async () => {
+    (getAgentDownloadUrl as jest.Mock).mockResolvedValue({ url });
+    (getAgentFileName as jest.Mock).mockReturnValue('elastic-agent-8.15.0');
+    (downloadAndStoreAgent as jest.Mock).mockResolvedValue(undefined);
+
+    await agentDownloaderRunner({
+      flags: { version, closestMatch },
+      log,
+    } as unknown as RunContext);
+
+    expect(getAgentDownloadUrl).toHaveBeenCalledWith(version, closestMatch, log);
+    expect(getAgentFileName).toHaveBeenCalledWith(version);
+    expect(downloadAndStoreAgent).toHaveBeenCalledWith(url, fileName);
+    expect(log.info).toHaveBeenCalledWith('Successfully downloaded and stored version 8.15.0');
+  });
+
+  it('logs an error if the download fails', async () => {
+    (getAgentDownloadUrl as jest.Mock).mockResolvedValue({ url });
+    (getAgentFileName as jest.Mock).mockReturnValue('elastic-agent-8.15.0');
+    (downloadAndStoreAgent as jest.Mock).mockRejectedValue(new Error('Download failed'));
+
+    await agentDownloaderRunner({
+      flags: { version, closestMatch },
+      log,
+    } as unknown as RunContext);
+
+    expect(getAgentDownloadUrl).toHaveBeenCalledWith(version, closestMatch, log);
+    expect(getAgentFileName).toHaveBeenCalledWith(version);
+    expect(downloadAndStoreAgent).toHaveBeenCalledWith(url, fileName);
+    expect(log.error).toHaveBeenCalledWith(
+      'Failed to download or store version 8.15.0: Download failed'
+    );
+  });
+
+  it('downloads and stores the previous patch version if the specified version fails', async () => {
+    const fallbackVersion = '8.15.0';
+    const fallbackFileName = 'elastic-agent-8.15.0.tar.gz';
+
+    (getAgentDownloadUrl as jest.Mock)
+      .mockResolvedValueOnce({ url })
+      .mockResolvedValueOnce({ url });
+    (getAgentFileName as jest.Mock)
+      .mockReturnValueOnce('elastic-agent-8.15.1')
+      .mockReturnValueOnce('elastic-agent-8.15.0');
+    (downloadAndStoreAgent as jest.Mock)
+      .mockRejectedValueOnce(new Error('Download failed'))
+      .mockResolvedValueOnce(undefined);
+
+    await agentDownloaderRunner({
+      flags: { version: '8.15.1', closestMatch },
+      log,
+    } as unknown as RunContext);
+
+    expect(getAgentDownloadUrl).toHaveBeenCalledWith('8.15.1', closestMatch, log);
+    expect(getAgentDownloadUrl).toHaveBeenCalledWith(fallbackVersion, closestMatch, log);
+    expect(getAgentFileName).toHaveBeenCalledWith('8.15.1');
+    expect(getAgentFileName).toHaveBeenCalledWith(fallbackVersion);
+    expect(downloadAndStoreAgent).toHaveBeenCalledWith(url, 'elastic-agent-8.15.1.tar.gz');
+    expect(downloadAndStoreAgent).toHaveBeenCalledWith(url, fallbackFileName);
+    expect(log.error).toHaveBeenCalledWith(
+      'Failed to download or store version 8.15.1: Download failed'
+    );
+    expect(log.info).toHaveBeenCalledWith('Successfully downloaded and stored version 8.15.0');
+  });
+
+  it('logs an error if all downloads fail', async () => {
+    (getAgentDownloadUrl as jest.Mock).mockResolvedValue({ url });
+    (getAgentFileName as jest.Mock)
+      .mockReturnValueOnce('elastic-agent-8.15.1')
+      .mockReturnValueOnce('elastic-agent-8.15.0');
+    (downloadAndStoreAgent as jest.Mock)
+      .mockRejectedValueOnce(new Error('Download failed'))
+      .mockRejectedValueOnce(new Error('Download failed'));
+
+    await agentDownloaderRunner({
+      flags: { version: '8.15.1', closestMatch },
+      log,
+    } as unknown as RunContext);
+
+    expect(getAgentDownloadUrl).toHaveBeenCalledWith('8.15.1', closestMatch, log);
+    expect(getAgentDownloadUrl).toHaveBeenCalledWith('8.15.0', closestMatch, log);
+    expect(getAgentFileName).toHaveBeenCalledWith('8.15.1');
+    expect(getAgentFileName).toHaveBeenCalledWith('8.15.0');
+    expect(downloadAndStoreAgent).toHaveBeenCalledWith(url, 'elastic-agent-8.15.1.tar.gz');
+    expect(downloadAndStoreAgent).toHaveBeenCalledWith(url, 'elastic-agent-8.15.0.tar.gz');
+    expect(log.error).toHaveBeenCalledWith(
+      'Failed to download or store version 8.15.1: Download failed'
+    );
+    expect(log.error).toHaveBeenCalledWith(
+      'Failed to download or store version 8.15.0: Download failed'
+    );
+  });
+
+  it('does not attempt fallback when patch version is 0', async () => {
+    (getAgentDownloadUrl as jest.Mock).mockResolvedValue({ url });
+    (getAgentFileName as jest.Mock).mockReturnValue('elastic-agent-8.15.0');
+    (downloadAndStoreAgent as jest.Mock).mockResolvedValue(undefined);
+
+    await agentDownloaderRunner({
+      flags: { version: '8.15.0', closestMatch },
+      log,
+    } as unknown as RunContext);
+
+    expect(getAgentDownloadUrl).toHaveBeenCalledTimes(1); // Only one call for 8.15.0
+    expect(getAgentFileName).toHaveBeenCalledTimes(1);
+    expect(downloadAndStoreAgent).toHaveBeenCalledWith(url, fileName);
+    expect(log.info).toHaveBeenCalledWith('Successfully downloaded and stored version 8.15.0');
+  });
+
+  it('logs an error for an invalid version format', async () => {
+    const invalidVersion = '7.x.x';
+
+    await expect(
+      agentDownloaderRunner({
+        flags: { version: invalidVersion, closestMatch },
+        log,
+      } as unknown as RunContext)
+    ).rejects.toThrow('Invalid version format');
+  });
+
+  it('passes the closestMatch flag correctly', async () => {
+    closestMatch = true;
+
+    (getAgentDownloadUrl as jest.Mock).mockResolvedValue({ url });
+    (getAgentFileName as jest.Mock).mockReturnValue('elastic-agent-8.15.0');
+    (downloadAndStoreAgent as jest.Mock).mockResolvedValue(undefined);
+
+    await agentDownloaderRunner({
+      flags: { version, closestMatch },
+      log,
+    } as unknown as RunContext);
+
+    expect(getAgentDownloadUrl).toHaveBeenCalledWith(version, closestMatch, log);
+  });
+
+  it('throws an error when version is not provided', async () => {
+    await expect(
+      agentDownloaderRunner({
+        flags: { closestMatch },
+        log,
+      } as unknown as RunContext)
+    ).rejects.toThrow('version argument is required');
+  });
+
+  it('logs the correct messages when both version and fallback version are processed', async () => {
+    const primaryVersion = '8.15.1';
+
+    (getAgentDownloadUrl as jest.Mock)
+      .mockResolvedValueOnce({ url })
+      .mockResolvedValueOnce({ url });
+
+    (getAgentFileName as jest.Mock)
+      .mockReturnValueOnce('elastic-agent-8.15.1')
+      .mockReturnValueOnce('elastic-agent-8.15.0');
+
+    (downloadAndStoreAgent as jest.Mock)
+      .mockRejectedValueOnce(new Error('Download failed')) // Fail on primary
+      .mockResolvedValueOnce(undefined); // Success on fallback
+
+    await agentDownloaderRunner({
+      flags: { version: primaryVersion, closestMatch },
+      log,
+    } as unknown as RunContext);
+
+    expect(log.error).toHaveBeenCalledWith(
+      'Failed to download or store version 8.15.1: Download failed'
+    );
+    expect(log.info).toHaveBeenCalledWith('Successfully downloaded and stored version 8.15.0');
+  });
+});

x-pack/plugins/security_solution/scripts/endpoint/agent_downloader_cli/agent_downloader.ts

@@ -8,24 +8,72 @@
 import { ok } from 'assert';
 import type { RunFn } from '@kbn/dev-cli-runner';
 import type { ToolingLog } from '@kbn/tooling-log';
+import semver from 'semver';
 import { getAgentDownloadUrl, getAgentFileName } from '../common/fleet_services';
 import { downloadAndStoreAgent } from '../common/agent_downloads_service';
 
+// Decrement the patch version by 1 and preserve pre-release tag (if any)
+const decrementPatchVersion = (version: string): string | null => {
+  const parsedVersion = semver.parse(version);
+  if (!parsedVersion) {
+    return null;
+  }
+  const newPatchVersion = parsedVersion.patch - 1;
+  // Create a new version string with the decremented patch - removing any possible pre-release tag
+  const newVersion = `${parsedVersion.major}.${parsedVersion.minor}.${newPatchVersion}`;
+  return semver.valid(newVersion) ? newVersion : null;
+};
+
+// Generate a list of versions to attempt downloading, including a fallback to the previous patch (GA)
+const getVersionsToDownload = (version: string): string[] => {
+  const parsedVersion = semver.parse(version);
+  if (!parsedVersion) return [];
+  // If patch version is 0, return only the current version.
+  if (parsedVersion.patch === 0) {
+    return [version];
+  }
+
+  const decrementedVersion = decrementPatchVersion(version);
+  return decrementedVersion ? [version, decrementedVersion] : [version];
+};
+
+// Download and store the Elastic Agent for the specified version(s)
 const downloadAndStoreElasticAgent = async (
   version: string,
   closestMatch: boolean,
   log: ToolingLog
-) => {
-  const downloadUrlResponse = await getAgentDownloadUrl(version, closestMatch, log);
-  const fileNameNoExtension = getAgentFileName(version);
-  const agentFile = `${fileNameNoExtension}.tar.gz`;
-  await downloadAndStoreAgent(downloadUrlResponse.url, agentFile);
+): Promise<void> => {
+  const versionsToDownload = getVersionsToDownload(version);
+
+  // Although we have a list of versions to try downloading, we only need to download one, and will return as soon as it succeeds.
+  for (const versionToDownload of versionsToDownload) {
+    try {
+      const { url } = await getAgentDownloadUrl(versionToDownload, closestMatch, log);
+      const fileName = `${getAgentFileName(versionToDownload)}.tar.gz`;
+
+      await downloadAndStoreAgent(url, fileName);
+      log.info(`Successfully downloaded and stored version ${versionToDownload}`);
+      return; // Exit once successful
+    } catch (error) {
+      log.error(`Failed to download or store version ${versionToDownload}: ${error.message}`);
+    }
+  }
+
+  log.error(`Failed to download agent for any available version: ${versionsToDownload.join(', ')}`);
 };
 
 export const agentDownloaderRunner: RunFn = async (cliContext) => {
-  ok(cliContext.flags.version, 'version argument is required');
+  const { version } = cliContext.flags;
+
+  ok(version, 'version argument is required');
+
+  // Validate version format
+  if (!semver.valid(version as string)) {
+    throw new Error('Invalid version format');
+  }
+
   await downloadAndStoreElasticAgent(
-    cliContext.flags.version as string,
+    version as string,
     cliContext.flags.closestMatch as boolean,
     cliContext.log
   );

x-pack/plugins/security_solution/scripts/endpoint/common/agent_downloads_service.test.ts

@@ -0,0 +1,77 @@
+/*
+ * 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.
+ */
+// Adjust path if needed
+
+import { downloadAndStoreAgent, isAgentDownloadFromDiskAvailable } from './agent_downloads_service';
+import fs from 'fs';
+import nodeFetch from 'node-fetch';
+import { finished } from 'stream/promises';
+
+jest.mock('fs');
+jest.mock('node-fetch');
+jest.mock('stream/promises', () => ({
+  finished: jest.fn(),
+}));
+jest.mock('../../../common/endpoint/data_loaders/utils', () => ({
+  createToolingLogger: jest.fn(() => ({
+    debug: jest.fn(),
+    info: jest.fn(),
+    error: jest.fn(),
+  })),
+}));
+
+describe('AgentDownloadStorage', () => {
+  const url = 'http://example.com/agent.tar.gz';
+  const fileName = 'elastic-agent-7.10.0.tar.gz';
+  beforeEach(() => {
+    jest.clearAllMocks(); // Ensure no previous test state affects the current one
+  });
+
+  it('downloads and stores the agent if not cached', async () => {
+    (fs.existsSync as unknown as jest.Mock).mockReturnValue(false);
+    (fs.createWriteStream as unknown as jest.Mock).mockReturnValue({
+      on: jest.fn(),
+      end: jest.fn(),
+    });
+    (nodeFetch as unknown as jest.Mock).mockResolvedValue({ body: { pipe: jest.fn() } });
+    (finished as unknown as jest.Mock).mockResolvedValue(undefined);
+
+    const result = await downloadAndStoreAgent(url, fileName);
+
+    expect(result).toEqual({
+      url,
+      filename: fileName,
+      directory: expect.any(String),
+      fullFilePath: expect.stringContaining(fileName), // Dynamically match the file path
+    });
+  });
+
+  it('reuses cached agent if available', async () => {
+    (fs.existsSync as unknown as jest.Mock).mockReturnValue(true);
+
+    const result = await downloadAndStoreAgent(url, fileName);
+
+    expect(result).toEqual({
+      url,
+      filename: fileName,
+      directory: expect.any(String),
+      fullFilePath: expect.stringContaining(fileName), // Dynamically match the path
+    });
+  });
+
+  it('checks if agent download is available from disk', () => {
+    (fs.existsSync as unknown as jest.Mock).mockReturnValue(true);
+
+    const result = isAgentDownloadFromDiskAvailable(fileName);
+
+    expect(result).toEqual({
+      filename: fileName,
+      directory: expect.any(String),
+      fullFilePath: expect.stringContaining(fileName), // Dynamically match the path
+    });
+  });
+});