Merge pull request #442 from SquirrelCorporation/bugfix-use-tmp-keys

[BUG] enhance SSH key handling
SquirrelCorporation · Nov 7, 2024 · bc3cebb · bc3cebb
2 parents 26c7716 + d8d3ae3
commit bc3cebb
Show file tree

Hide file tree

Showing 26 changed files with 284 additions and 169 deletions.
diff --git a/client/src/components/DeviceConfiguration/CheckDeviceConnection.tsx b/client/src/components/DeviceConfiguration/CheckDeviceConnection.tsx
@@ -14,7 +14,7 @@ import {
 import { Alert, Button, Popover, Steps, Typography } from 'antd';
 import { motion } from 'framer-motion';
 import React, { useEffect, useRef, useState } from 'react';
-import { API } from 'ssm-shared-lib';
+import { API, SsmAnsible } from 'ssm-shared-lib';
 
 export type CheckDeviceConnectionProps = {
   execId?: string;
@@ -52,7 +52,11 @@ const CheckDeviceConnection: React.FC<CheckDeviceConnectionProps> = (props) => {
   const [count, setCount] = useState(0);
 
   const isFinalStatusFailed = async () => {
-    if (savedStatuses?.find((status) => status._status === 'failed')) {
+    if (
+      savedStatuses?.find(
+        (status) => status._status === SsmAnsible.AnsibleTaskStatus.FAILED,
+      )
+    ) {
       const res = await getAnsibleSmartFailure({ execId: execId });
       if (res.data) {
         setSmartFailure(res.data);

diff --git a/client/src/components/PlaybookExecutionModal/PlaybookExecutionHandler.ts b/client/src/components/PlaybookExecutionModal/PlaybookExecutionHandler.ts
@@ -2,7 +2,7 @@ import taskStatusTimeline from '@/components/PlaybookExecutionModal/TaskStatusTi
 import { getExecLogs, getTaskStatuses } from '@/services/rest/playbooks';
 import { StepsProps } from 'antd';
 import React, { ReactNode } from 'react';
-import { API } from 'ssm-shared-lib';
+import { API, SsmAnsible } from 'ssm-shared-lib';
 
 export type TaskStatusTimelineType = StepsProps & {
   _status: string;
@@ -42,7 +42,12 @@ export default class PlaybookExecutionHandler {
   }
 
   static isFinalStatus = (status: string): boolean => {
-    return status === 'failed' || status === 'successful';
+    return (
+      status === SsmAnsible.AnsibleTaskStatus.FAILED ||
+      status === SsmAnsible.AnsibleTaskStatus.SUCCESS ||
+      status === SsmAnsible.AnsibleTaskStatus.CANCELED ||
+      status === SsmAnsible.AnsibleTaskStatus.TIMEOUT
+    );
   };
 
   resetTerminal = () => {

diff --git a/client/src/components/PlaybookExecutionModal/PlaybookExecutionTerminalModal.tsx b/client/src/components/PlaybookExecutionModal/PlaybookExecutionTerminalModal.tsx
@@ -19,7 +19,7 @@ import React, {
   useRef,
   useState,
 } from 'react';
-import { API } from 'ssm-shared-lib';
+import { API, SsmAnsible } from 'ssm-shared-lib';
 
 export interface PlaybookExecutionTerminalModalHandles {
   resetTerminal: () => void;
@@ -110,7 +110,11 @@ const PlaybookExecutionTerminalModal = React.forwardRef<
     };
 
     const isFinalStatusFailed = async () => {
-      if (savedStatuses?.find((status) => status._status === 'failed')) {
+      if (
+        savedStatuses?.find(
+          (status) => status._status === SsmAnsible.AnsibleTaskStatus.FAILED,
+        )
+      ) {
         const res = await getAnsibleSmartFailure({ execId: execId });
         if (res.data) {
           api.open({

diff --git a/client/src/components/PlaybookExecutionModal/TaskStatusTimeline.tsx b/client/src/components/PlaybookExecutionModal/TaskStatusTimeline.tsx
@@ -8,27 +8,31 @@ import {
 } from '@ant-design/icons';
 import { StepsProps } from 'antd';
 import React, { ReactNode } from 'react';
-import { API } from 'ssm-shared-lib';
+import { API, SsmAnsible } from 'ssm-shared-lib';
 
 const transformToTaskStatusTimeline = (
   execStatus: API.ExecStatus,
 ): TaskStatusTimelineType => {
   //  status?: 'wait' | 'process' | 'finish' | 'error';
   let status: StepsProps['status'] = undefined;
   let icon: ReactNode = <QuestionOutlined />;
-  if (execStatus.status === 'starting') {
+  if (execStatus.status === SsmAnsible.AnsibleTaskStatus.STARTING) {
     status = 'finish';
     icon = <VerticalAlignBottomOutlined />;
   }
-  if (execStatus.status === 'running') {
+  if (execStatus.status === SsmAnsible.AnsibleTaskStatus.RUNNING) {
     status = 'process';
     icon = <LoadingOutlined />;
   }
-  if (execStatus.status === 'failed') {
+  if (
+    execStatus.status === SsmAnsible.AnsibleTaskStatus.FAILED ||
+    execStatus.status === SsmAnsible.AnsibleTaskStatus.CANCELED ||
+    execStatus.status === SsmAnsible.AnsibleTaskStatus.TIMEOUT
+  ) {
     status = 'error';
     icon = <CloseCircleOutlined />;
   }
-  if (execStatus.status === 'successful') {
+  if (execStatus.status === SsmAnsible.AnsibleTaskStatus.SUCCESS) {
     status = 'finish';
     icon = <CheckCircleOutlined />;
   }

diff --git a/client/src/pages/Admin/Logs/TaskLogsColumns.tsx b/client/src/pages/Admin/Logs/TaskLogsColumns.tsx
@@ -1,7 +1,7 @@
 import { ProColumns } from '@ant-design/pro-components';
 import { Tag } from 'antd';
 import React from 'react';
-import { API } from 'ssm-shared-lib';
+import { API, SsmAnsible } from 'ssm-shared-lib';
 
 const TaskLogsColumns: ProColumns<API.Task>[] = [
   {
@@ -24,15 +24,17 @@ const TaskLogsColumns: ProColumns<API.Task>[] = [
       failed: { text: 'failed' },
       successful: { text: 'successful' },
       starting: { text: 'starting' },
+      timeout: { text: 'timeout' },
+      canceled: { text: 'canceled' },
     },
     render: (dom, entity) => {
       return (
         <Tag
           bordered={false}
           color={
-            entity.status === 'successful'
+            entity.status === SsmAnsible.AnsibleTaskStatus.SUCCESS
               ? 'success'
-              : entity.status === 'failed'
+              : entity.status === SsmAnsible.AnsibleTaskStatus.FAILED
                 ? 'error'
                 : 'default'
           }

diff --git a/client/tests/PlaybookExecutionHandler.test.ts b/client/tests/PlaybookExecutionHandler.test.ts
@@ -141,6 +141,6 @@ describe('PlaybookExecutionHandler', () => {
   it('should identify final statuses correctly', () => {
     expect(PlaybookExecutionHandler.isFinalStatus('failed')).toBe(true);
     expect(PlaybookExecutionHandler.isFinalStatus('successful')).toBe(true);
-    expect(PlaybookExecutionHandler.isFinalStatus('pending')).toBe(false);
+    expect(PlaybookExecutionHandler.isFinalStatus('running')).toBe(false);
   });
 });
diff --git a/server/src/ansible/inventory/inventory.py b/server/src/ansible/inventory/inventory.py
@@ -49,7 +49,7 @@ def parse_args():
 
 def load_inventory():
     headers = {'Accept': 'application/json', 'Authorization': "Bearer {}".format(os.getenv("SSM_API_KEY"))}
-    r = requests.get('http://localhost:3000/playbooks/inventory', headers=headers)
+    r = requests.get('http://localhost:3000/playbooks/inventory?execUuid={}'.format(os.getenv('SSM_EXEC_UUID')), headers=headers)
     return r.json()['data']
 
 

diff --git a/server/src/ansible/ssm-ansible-run.py b/server/src/ansible/ssm-ansible-run.py
@@ -79,7 +79,7 @@ def execute():
         debug = True
     if args.specific_host is not None:
         specific_host = json.loads(args.specific_host)
-
+    os.environ['SSM_EXEC_UUID'] = args.ident
     runner_args = {
         'ident': args.ident,
         'private_data_dir': './',

diff --git a/server/src/controllers/rest/devices/device.ts b/server/src/controllers/rest/devices/device.ts
@@ -12,7 +12,6 @@ import { BadRequestError, ForbiddenError, NotFoundError } from '../../../middlew
 import { SuccessResponse } from '../../../middlewares/api/ApiResponse';
 import { DEFAULT_VAULT_ID, vaultEncrypt } from '../../../modules/ansible-vault/ansible-vault';
 import WatcherEngine from '../../../modules/docker/core/WatcherEngine';
-import Shell from '../../../modules/shell';
 import DeviceUseCases from '../../../services/DeviceUseCases';
 
 export const addDevice = async (req, res) => {
@@ -53,9 +52,6 @@ export const addDevice = async (req, res) => {
       becomeMethod: becomeMethod,
       becomePass: becomePass ? await vaultEncrypt(becomePass, DEFAULT_VAULT_ID) : undefined,
     } as DeviceAuth);
-    if (sshKey) {
-      await Shell.SshPrivateKeyFileManager.saveSshKey(sshKey, createdDevice.uuid);
-    }
     void WatcherEngine.registerWatcher(createdDevice);
     new SuccessResponse('Add device successful', { device: createdDevice as API.DeviceItem }).send(
       res,

diff --git a/server/src/controllers/rest/devices/deviceauth.ts b/server/src/controllers/rest/devices/deviceauth.ts
@@ -7,7 +7,6 @@ import { InternalError, NotFoundError } from '../../../middlewares/api/ApiError'
 import { SuccessResponse } from '../../../middlewares/api/ApiResponse';
 import { DEFAULT_VAULT_ID, vaultEncrypt } from '../../../modules/ansible-vault/ansible-vault';
 import WatcherEngine from '../../../modules/docker/core/WatcherEngine';
-import Shell from '../../../modules/shell';
 
 const SENSITIVE_PLACEHOLDER = 'REDACTED';
 
@@ -113,9 +112,6 @@ export const addOrUpdateDeviceAuth = async (req, res) => {
       : undefined,
     becomeUser: becomeUser,
   } as DeviceAuth);
-  if (sshKey) {
-    await Shell.SshPrivateKeyFileManager.saveSshKey(sshKey, device.uuid);
-  }
   void WatcherEngine.deregisterWatchers();
   void WatcherEngine.registerWatchers();
   new SuccessResponse('Add or update device auth successful', { type: deviceAuth.authType }).send(

diff --git a/server/src/controllers/rest/playbooks/hook.ts b/server/src/controllers/rest/playbooks/hook.ts
@@ -2,8 +2,11 @@ import AnsibleLog from '../../../data/database/model/AnsibleLogs';
 import AnsibleLogsRepo from '../../../data/database/repository/AnsibleLogsRepo';
 import AnsibleTaskRepo from '../../../data/database/repository/AnsibleTaskRepo';
 import AnsibleTaskStatusRepo from '../../../data/database/repository/AnsibleTaskStatusRepo';
+import { isFinalStatus } from '../../../helpers/ansible/AnsibleTaskHelper';
+import logger from '../../../logger';
 import { BadRequestError, NotFoundError } from '../../../middlewares/api/ApiError';
 import { SuccessResponse } from '../../../middlewares/api/ApiResponse';
+import sshPrivateKeyFileManager from '../../../modules/shell/managers/SshPrivateKeyFileManager';
 
 export const addTaskStatus = async (req, res) => {
   if (!req.body.runner_ident || !req.body.status) {
@@ -17,6 +20,17 @@ export const addTaskStatus = async (req, res) => {
       ident: ident,
       status: status,
     });
+    if (isFinalStatus(status)) {
+      if (ansibleTask.target && ansibleTask.target.length > 0) {
+        logger.warn('Removing temporary private key');
+        ansibleTask.target?.map((e) =>
+          sshPrivateKeyFileManager.removeAnsibleTemporaryPrivateKey(e, ident),
+        );
+      } else {
+        logger.warn('Removing temporary private keys');
+        sshPrivateKeyFileManager.removeAllAnsibleExecTemporaryPrivateKeys(ident);
+      }
+    }
     new SuccessResponse('Added task status').send(res);
   } else {
     throw new NotFoundError('Task not found');

diff --git a/server/src/controllers/rest/playbooks/inventory.ts b/server/src/controllers/rest/playbooks/inventory.ts
@@ -6,19 +6,19 @@ import { SuccessResponse } from '../../../middlewares/api/ApiResponse';
 import InventoryTransformer from '../../../modules/ansible/utils/InventoryTransformer';
 
 export const getInventory = async (req, res) => {
+  const { execUuid } = req.query;
   logger.info(`[CONTROLLER] - GET - /ansible/inventory`);
   let devicesAuth: DeviceAuth[] | null = [];
-  if (req.body.target) {
+  if (req.body?.target) {
     logger.info(`[CONTROLLER][ANSIBLE[Inventory] - Target is ${req.body.target}`);
     devicesAuth = await DeviceAuthRepo.findOneByDeviceUuid(req.body.target);
   } else {
     logger.info(`[CONTROLLER][ANSIBLE][Inventory] - No target, get all`);
     devicesAuth = await DeviceAuthRepo.findAllPop();
   }
   if (devicesAuth) {
-    new SuccessResponse('Get inventory', InventoryTransformer.inventoryBuilder(devicesAuth)).send(
-      res,
-    );
+    const inventory = await InventoryTransformer.inventoryBuilder(devicesAuth, execUuid);
+    new SuccessResponse('Get inventory', inventory).send(res);
   } else {
     throw new NotFoundError('No devices auth found');
   }

diff --git a/server/src/controllers/rest/playbooks/playbook.ts b/server/src/controllers/rest/playbooks/playbook.ts
@@ -1,9 +1,7 @@
 import PlaybookRepo from '../../../data/database/repository/PlaybookRepo';
-import logger from '../../../logger';
 import { InternalError, NotFoundError } from '../../../middlewares/api/ApiError';
 import { SuccessResponse } from '../../../middlewares/api/ApiResponse';
 import Shell from '../../../modules/shell';
-import FileSystemManager from '../../../modules/shell/managers/FileSystemManager';
 import PlaybooksRepositoryUseCases from '../../../services/PlaybooksRepositoryUseCases';
 import PlaybookUseCases from '../../../services/PlaybookUseCases';
 

diff --git a/server/src/core/startup/index.ts b/server/src/core/startup/index.ts
@@ -14,9 +14,9 @@ import NotificationComponent from '../../modules/notifications/NotificationCompo
 import ContainerCustomStacksRepositoryEngine from '../../modules/repository/ContainerCustomStacksRepositoryEngine';
 import { createADefaultLocalUserRepository } from '../../modules/repository/default-playbooks-repositories';
 import PlaybooksRepositoryEngine from '../../modules/repository/PlaybooksRepositoryEngine';
+import sshPrivateKeyFileManager from '../../modules/shell/managers/SshPrivateKeyFileManager';
 import UpdateChecker from '../../modules/update/UpdateChecker';
 import ContainerRegistryUseCases from '../../services/ContainerRegistryUseCases';
-import DeviceAuthUseCases from '../../services/DeviceAuthUseCases';
 import { setAnsibleVersions } from '../system/ansible-versions';
 
 class Startup {
@@ -39,8 +39,8 @@ class Startup {
   }
 
   private async initializeModules() {
-    await DeviceAuthUseCases.saveAllDeviceAuthSshKeys();
     await PlaybooksRepositoryEngine.init();
+    void sshPrivateKeyFileManager.removeAllAnsibleTemporaryPrivateKeys();
     void NotificationComponent.init();
     void Crons.initScheduledJobs();
     void WatcherEngine.init();

diff --git a/server/src/data/database/model/AnsibleTask.ts b/server/src/data/database/model/AnsibleTask.ts
@@ -7,6 +7,7 @@ export default interface AnsibleTask {
   ident: string;
   status?: string;
   cmd?: string;
+  target?: string[];
   createdAt?: string;
 }
 
@@ -25,6 +26,9 @@ const schema = new Schema<AnsibleTask>(
       type: Schema.Types.String,
       required: true,
     },
+    target: {
+      type: Schema.Types.Array,
+    },
   },
   {
     timestamps: true,

diff --git a/server/src/helpers/ansible/AnsibleTaskHelper.ts b/server/src/helpers/ansible/AnsibleTaskHelper.ts
@@ -0,0 +1,10 @@
+import { SsmAnsible } from 'ssm-shared-lib';
+
+export const isFinalStatus = (status: string): boolean => {
+  return (
+    status === SsmAnsible.AnsibleTaskStatus.FAILED ||
+    status === SsmAnsible.AnsibleTaskStatus.SUCCESS ||
+    status === SsmAnsible.AnsibleTaskStatus.CANCELED ||
+    status === SsmAnsible.AnsibleTaskStatus.TIMEOUT
+  );
+};