Refactor mainAction: extract config-audit persistence and cleanup builder

src/commands/main-action.test.ts

@@ -4,6 +4,8 @@
 const mockMkdirSync = jest.fn();
 const mockWriteFileSync = jest.fn();
 const mockChmodSync = jest.fn();
+const mockOpenSync = jest.fn().mockReturnValue(42);
+const mockCloseSync = jest.fn();
 
 jest.mock('fs', () => {
   const actual = jest.requireActual<typeof import('fs')>('fs');
@@ -12,10 +14,12 @@ jest.mock('fs', () => {
     mkdirSync: (...args: unknown[]) => mockMkdirSync(...args),
     writeFileSync: (...args: unknown[]) => mockWriteFileSync(...args),
     chmodSync: (...args: unknown[]) => mockChmodSync(...args),
+    openSync: (...args: unknown[]) => mockOpenSync(...args),
+    closeSync: (...args: unknown[]) => mockCloseSync(...args),
   };
 });
 
-import { createMainAction } from './main-action';
+import { createMainAction, testHelpers } from './main-action';
 
 // eslint-disable-next-line @typescript-eslint/no-require-imports
 jest.mock('../logger', () => require('../test-helpers/mock-logger.test-utils').loggerMockFactory());
@@ -400,6 +404,9 @@ describe('createMainAction', () => {
       mockMkdirSync.mockReset();
       mockWriteFileSync.mockReset();
       mockChmodSync.mockReset();
+      mockOpenSync.mockReset();
+      mockOpenSync.mockReturnValue(42);
+      mockCloseSync.mockReset();
     });
 
     afterEach(() => jest.restoreAllMocks());
@@ -415,16 +422,15 @@ describe('createMainAction', () => {
       const action = createMainAction(getOptionValueSource);
       await action(['echo hi'], {});
 
-      expect(mockMkdirSync).toHaveBeenCalledWith('/tmp/awf-audit', { recursive: true, mode: 0o755 });
+      expect(mockMkdirSync).toHaveBeenCalledWith('/tmp/awf-audit', { recursive: true, mode: 0o700 });
+      expect(mockOpenSync).toHaveBeenCalledWith('/tmp/awf-audit/awf-resolved-config.json', 'wx', 0o600);
       expect(mockWriteFileSync).toHaveBeenCalledWith(
-        '/tmp/awf-audit/awf-resolved-config.json',
+        42,
         expect.stringContaining('"allowedDomains"'),
-        { mode: 0o644 },
       );
-      expect(mockChmodSync).toHaveBeenCalledWith('/tmp/awf-audit/awf-resolved-config.json', 0o644);
       // Verify secret key names are excluded from the artifact
       const written = mockWriteFileSync.mock.calls.find(
-        (c) => String(c[0]).includes('awf-resolved-config.json')
+        (c) => typeof c[0] === 'number'
       );
       expect(written).toBeDefined();
       const writtenJson = String(written![1]);
@@ -451,7 +457,7 @@ describe('createMainAction', () => {
       await action(['echo hi'], {});
 
       const written = mockWriteFileSync.mock.calls.find(
-        (c) => String(c[0]).includes('awf-resolved-config.json')
+        (c) => typeof c[0] === 'number'
       );
       expect(written).toBeDefined();
       const writtenJson = String(written![1]);
@@ -464,11 +470,68 @@ describe('createMainAction', () => {
       const action = createMainAction(getOptionValueSource);
       await action(['echo hi'], {});
 
-      expect(mockWriteFileSync).toHaveBeenCalledWith(
+      expect(mockOpenSync).toHaveBeenCalledWith(
         '/tmp/awf-test/audit/awf-resolved-config.json',
+        'wx',
+        0o600,
+      );
+      expect(mockWriteFileSync).toHaveBeenCalledWith(
+        42,
         expect.any(String),
-        { mode: 0o644 },
       );
     });
   });
+
+  describe('extracted helper functions', () => {
+    it('redactConfigForLogging removes sensitive keys and redacts agentCommand', () => {
+      const secretValue = 'secret-123';
+      const configWithSecrets = {
+        ...STUB_CONFIG,
+        agentCommand: `agent --token ${secretValue}`,
+        openaiApiKey: 'sk-secret',
+      } as unknown as import('../types').WrapperConfig;
+      mockedRedactSecrets.redactSecrets.mockImplementation((s: string) =>
+        s.replace(secretValue, '[REDACTED]')
+      );
+
+      const redacted = testHelpers.redactConfigForLogging(configWithSecrets);
+
+      expect(redacted).not.toHaveProperty('openaiApiKey');
+      expect(redacted.agentCommand).toContain('[REDACTED]');
+      expect(redacted.agentCommand).not.toContain(secretValue);
+    });
+
+    it('persistConfigAuditArtifact logs debug when writing the artifact fails', () => {
+      mockMkdirSync.mockImplementationOnce(() => {
+        throw new Error('write failed');
+      });
+
+      testHelpers.persistConfigAuditArtifact(STUB_CONFIG, { foo: 'bar' });
+
+      expect(mockedLogger.debug).toHaveBeenCalledWith(
+        expect.stringContaining('Failed to write resolved config artifact:')
+      );
+    });
+
+    it('buildCleanupFn runs cleanup using provided state getters', async () => {
+      const performCleanup = testHelpers.buildCleanupFn(
+        STUB_CONFIG,
+        () => true,
+        () => true,
+      );
+
+      await performCleanup();
+
+      expect(mockedDockerManager.preserveIptablesAudit).toHaveBeenCalledWith(
+        STUB_CONFIG.workDir,
+        STUB_CONFIG.auditDir
+      );
+      expect(mockedDockerManager.stopContainers).toHaveBeenCalledWith(
+        STUB_CONFIG.workDir,
+        STUB_CONFIG.keepContainers
+      );
+      expect(mockedHostIptables.cleanupHostIptables).toHaveBeenCalled();
+      expect(mockedDockerManager.cleanup).toHaveBeenCalled();
+    });
+  });
 });

src/commands/main-action.ts

@@ -26,6 +26,100 @@ import { validateOptions } from './validate-options';
 import { probeSplitFilesystem } from '../dind-probe';
 import { assertTopologySupported, connectTopologyContainers } from '../topology';
 import { runDindBootstrap } from '../dind-bootstrap';
+import type { WrapperConfig } from '../types';
+
+const SENSITIVE_CONFIG_KEYS = new Set([
+  'openaiApiKey',
+  'anthropicApiKey',
+  'copilotGithubToken',
+  'copilotProviderApiKey',
+  'geminiApiKey',
+  'githubToken',
+]);
+
+function redactConfigForLogging(config: WrapperConfig): Record<string, unknown> {
+  const redactedConfig: Record<string, unknown> = {};
+  for (const [key, value] of Object.entries(config)) {
+    if (SENSITIVE_CONFIG_KEYS.has(key)) continue;
+
+    if (key === 'agentCommand') {
+      redactedConfig[key] = redactSecrets(value as string);
+      continue;
+    }
+
+    if (key === 'additionalEnv' && value && typeof value === 'object') {
+      redactedConfig[key] = Object.fromEntries(
+        Object.keys(value as Record<string, string>).map((envKey) => [envKey, '[REDACTED]']),
+      );
+      continue;
+    }
+
+    redactedConfig[key] = value;
+  }
+  return redactedConfig;
+}
+
+function persistConfigAuditArtifact(
+  config: WrapperConfig,
+  redactedConfig: Record<string, unknown>,
+): void {
+  try {
+    const configArtifactDir = config.auditDir || path.join(config.workDir, 'audit');
+    fs.mkdirSync(configArtifactDir, { recursive: true, mode: 0o700 });
+    const configArtifactPath = path.join(configArtifactDir, 'awf-resolved-config.json');
+    const fd = fs.openSync(configArtifactPath, 'wx', 0o600);
+    try {
+      fs.writeFileSync(fd, JSON.stringify(redactedConfig, null, 2) + '\n');
+    } finally {
+      fs.closeSync(fd);
+    }
+  } catch (err) {
+    logger.debug(`Failed to write resolved config artifact: ${err}`);
+  }
+}
+
+function buildCleanupFn(
+  config: WrapperConfig,
+  getContainersStarted: () => boolean,
+  getHostIptablesSetup: () => boolean,
+) {
+  return async (signal?: string) => {
+    if (signal) {
+      logger.info(`Received ${signal}, cleaning up...`);
+    }
+
+    // Copy iptables audit BEFORE stopping containers (volumes are destroyed by `docker compose down -v`)
+    if (getContainersStarted()) {
+      preserveIptablesAudit(config.workDir, config.auditDir);
+      await stopContainers(config.workDir, config.keepContainers);
+    }
+
+    if (getHostIptablesSetup() && !config.keepContainers) {
+      await cleanupHostIptables();
+    }
+
+    if (!config.keepContainers) {
+      await cleanup(
+        config.workDir,
+        false,
+        config.proxyLogsDir,
+        config.auditDir,
+        config.sessionStateDir,
+        config.dockerHostPathPrefix,
+        config.imageRegistry,
+        config.imageTag,
+        config.agentImage,
+      );
+      // Note: We don't remove the firewall network here since it can be reused
+      // across multiple runs. Cleanup script will handle removal if needed.
+    } else {
+      logger.info(`Configuration files preserved at: ${config.workDir}`);
+      logger.info(`Agent logs available at: ${config.workDir}/agent-logs/`);
+      logger.info(`Squid logs available at: ${config.workDir}/squid-logs/`);
+      logger.info(`Host iptables rules preserved (--keep-containers enabled)`);
+    }
+  };
+}
 
 /**
  * Resolves the Commander option-value source for a given option name.
@@ -106,33 +200,9 @@ export function createMainAction(getOptionValueSource: OptionSourceResolver) {
 
   // Log config with redacted secrets - remove API keys entirely
   // to prevent sensitive data from flowing to logger (CodeQL sensitive data logging)
-  const redactedConfig: Record<string, unknown> = {};
-  for (const [key, value] of Object.entries(config)) {
-    if (
-      key === 'openaiApiKey' ||
-      key === 'anthropicApiKey' ||
-      key === 'copilotGithubToken' ||
-      key === 'copilotProviderApiKey' ||
-      key === 'geminiApiKey'
-    ) continue;
-    redactedConfig[key] = key === 'agentCommand' ? redactSecrets(value as string) : value;
-  }
+  const redactedConfig = redactConfigForLogging(config);
   logger.debug('Configuration:', JSON.stringify(redactedConfig, null, 2));
-
-  // Persist redacted config to audit artifact for post-run diagnostics
-  try {
-    const configArtifactDir = config.auditDir || path.join(config.workDir, 'audit');
-    fs.mkdirSync(configArtifactDir, { recursive: true, mode: 0o755 });
-    const configArtifactPath = path.join(configArtifactDir, 'awf-resolved-config.json');
-    fs.writeFileSync(
-      configArtifactPath,
-      JSON.stringify(redactedConfig, null, 2) + '\n',
-      { mode: 0o644 },
-    );
-    fs.chmodSync(configArtifactPath, 0o644);
-  } catch (err) {
-    logger.debug(`Failed to write resolved config artifact: ${err}`);
-  }
+  persistConfigAuditArtifact(config, redactedConfig);
 
   logger.info(`Allowed domains: ${config.allowedDomains.join(', ')}`);
   if (config.blockedDomains && config.blockedDomains.length > 0) {
@@ -145,43 +215,11 @@ export function createMainAction(getOptionValueSource: OptionSourceResolver) {
   let containersStarted = false;
   let hostIptablesSetup = false;
 
-  // Handle cleanup on process exit
-  const performCleanup = async (signal?: string) => {
-    if (signal) {
-      logger.info(`Received ${signal}, cleaning up...`);
-    }
-
-    // Copy iptables audit BEFORE stopping containers (volumes are destroyed by `docker compose down -v`)
-    if (containersStarted) {
-      preserveIptablesAudit(config.workDir, config.auditDir);
-      await stopContainers(config.workDir, config.keepContainers);
-    }
-
-    if (hostIptablesSetup && !config.keepContainers) {
-      await cleanupHostIptables();
-    }
-
-    if (!config.keepContainers) {
-      await cleanup(
-        config.workDir,
-        false,
-        config.proxyLogsDir,
-        config.auditDir,
-        config.sessionStateDir,
-        config.dockerHostPathPrefix,
-        config.imageRegistry,
-        config.imageTag,
-        config.agentImage,
-      );
-      // Note: We don't remove the firewall network here since it can be reused
-      // across multiple runs. Cleanup script will handle removal if needed.
-    } else {
-      logger.info(`Configuration files preserved at: ${config.workDir}`);
-      logger.info(`Agent logs available at: ${config.workDir}/agent-logs/`);
-      logger.info(`Squid logs available at: ${config.workDir}/squid-logs/`);
-      logger.info(`Host iptables rules preserved (--keep-containers enabled)`);
-    }
-  };
+  const performCleanup = buildCleanupFn(
+    config,
+    () => containersStarted,
+    () => hostIptablesSetup,
+  );
 
   // Register signal handlers for graceful shutdown
   registerSignalHandlers({
@@ -226,3 +264,11 @@ export function createMainAction(getOptionValueSource: OptionSourceResolver) {
   }
   };
 }
+
+/** @internal Exposed for unit tests. */
+// ts-prune-ignore-next
+export const testHelpers = {
+  redactConfigForLogging,
+  persistConfigAuditArtifact,
+  buildCleanupFn,
+};