CM-40766 - Add missing toolbar actions: run all, expand all, collapse all (#107)

Author: MarshalX

package.json

@@ -36,6 +36,21 @@
         }
       ],
       "view/title": [
+        {
+          "command": "cycode.runAllScans",
+          "when": "view =~ /^cycode./",
+          "group": "navigation@0"
+        },
+        {
+          "command": "cycode.treeViewCollapseAllCommand",
+          "when": "view == cycode.view.tree",
+          "group": "navigation@1"
+        },
+        {
+          "command": "cycode.treeViewExpandAllCommand",
+          "when": "view == cycode.view.tree",
+          "group": "navigation@1"
+        },
         {
           "command": "cycode.openMainMenu",
           "when": "view == cycode.view.tree",
@@ -162,6 +177,24 @@
         "title": "Return to Home Screen",
         "icon": "$(arrow-left)"
       },
+      {
+        "command": "cycode.runAllScans",
+        "category": "Cycode",
+        "title": "Run all scans types for the entire project",
+        "icon": "$(testing-run-icon)"
+      },
+      {
+        "command": "cycode.treeViewExpandAllCommand",
+        "category": "Cycode",
+        "title": "Expand All Nodes in Tree View",
+        "icon": "$(expand-all)"
+      },
+      {
+        "command": "cycode.treeViewCollapseAllCommand",
+        "category": "Cycode",
+        "title": "Collapse All Nodes in Tree View",
+        "icon": "$(collapse-all)"
+      },
       {
         "command": "cycode.auth",
         "title": "Authenticate with service"

src/commands/index.ts

@@ -8,11 +8,14 @@ import iacScanCommand from './iac-scan-command';
 import iacScanForCurrentProjectCommand from './iac-scan-for-current-project-command';
 import sastScanCommand from './sast-scan-command';
 import sastScanForCurrentProjectCommand from './sast-scan-for-current-project-command';
-import onTreeItemClickCommand from './on-tree-item-click-command';
 import openViolationInFileCommand from './open-violation-in-file-command';
 import openViolationPanelCommand from './open-violation-panel-command';
 import openSettingsCommand from './open-settings-command';
 import openMainMenuCommand from './open-main-menu-command';
+import onTreeViewDetectionNodeClickCommand from './on-tree-view-detection-node-click-command';
+import treeViewExpandAllCommand from './tree-view-expand-all-command';
+import treeViewCollapseAllCommand from './tree-view-collapse-all-command';
+import runAllScansCommand from './run-all-scans-command';
 
 export enum VscodeCommands {
   SecretScanCommandId = 'cycode.secretScan',
@@ -22,19 +25,25 @@ export enum VscodeCommands {
   IacScanForProjectCommandId = 'cycode.iacScanForProject',
   SastScanCommandId = 'cycode.sastScan',
   SastScanForProjectCommandId = 'cycode.sastScanForProject',
+  RunAllScansCommandId = 'cycode.runAllScans',
 
   AuthCommandId = 'cycode.auth',
   IgnoreCommandId = 'cycode.ignore',
 
   OpenSettingsCommandId = 'cycode.openSettings',
   OpenMainMenuCommandId = 'cycode.openMainMenu',
 
-  ShowProblemsTab = 'workbench.action.problems.focus',
-  ShowCycodeView = 'workbench.view.extension.cycode',
-
   OpenViolationInFile = 'cycode.openViolationInFile',
   OpenViolationPanel = 'cycode.openViolationPanel',
-  OnTreeItemClick = 'cycode.onTreeItemClick',
+
+  OnTreeViewDetectionNodeClickCommand = 'cycode.onTreeViewDetectionNodeClickCommand',
+  TreeViewExpandAllCommand = 'cycode.treeViewExpandAllCommand',
+  TreeViewCollapseAllCommand = 'cycode.treeViewCollapseAllCommand',
+
+  // built-in or created automatically by vscode:
+  WorkbenchShowProblemsTab = 'workbench.action.problems.focus',
+  WorkbenchShowCycodeView = 'workbench.view.extension.cycode',
+  WorkbenchTreeViewCollapseAll = 'workbench.actions.treeView.cycode.view.tree.collapseAll',
 }
 
 const _VS_CODE_COMMANDS_ID_TO_CALLBACK: Record<string, (...args: never[]) => unknown> = {
@@ -47,11 +56,14 @@ const _VS_CODE_COMMANDS_ID_TO_CALLBACK: Record<string, (...args: never[]) => unk
   [VscodeCommands.IacScanForProjectCommandId]: iacScanForCurrentProjectCommand,
   [VscodeCommands.SastScanCommandId]: sastScanCommand,
   [VscodeCommands.SastScanForProjectCommandId]: sastScanForCurrentProjectCommand,
-  [VscodeCommands.OnTreeItemClick]: onTreeItemClickCommand,
+  [VscodeCommands.RunAllScansCommandId]: runAllScansCommand,
   [VscodeCommands.OpenViolationInFile]: openViolationInFileCommand,
   [VscodeCommands.OpenViolationPanel]: openViolationPanelCommand,
   [VscodeCommands.OpenSettingsCommandId]: openSettingsCommand,
   [VscodeCommands.OpenMainMenuCommandId]: openMainMenuCommand,
+  [VscodeCommands.OnTreeViewDetectionNodeClickCommand]: onTreeViewDetectionNodeClickCommand,
+  [VscodeCommands.TreeViewExpandAllCommand]: treeViewExpandAllCommand,
+  [VscodeCommands.TreeViewCollapseAllCommand]: treeViewCollapseAllCommand,
 };
 
 export const registerCommands = (context: vscode.ExtensionContext): void => {

src/commands/on-tree-item-click-command.ts renamed to src/commands/on-tree-view-detection-node-click-command.ts

@@ -0,0 +1,28 @@
+import * as vscode from 'vscode';
+import { container } from 'tsyringe';
+import { validateConfig } from '../utils/config';
+import { CycodeService, ICycodeService } from '../services/cycode-service';
+import { CliScanType } from '../cli/models/cli-scan-type';
+import { IStateService } from '../services/state-service';
+import { StateServiceSymbol } from '../symbols';
+
+export default async () => {
+  if (validateConfig()) {
+    return;
+  }
+
+  const stateService = container.resolve<IStateService>(StateServiceSymbol);
+  if (!stateService.globalState.CliAuthed) {
+    vscode.window.showErrorMessage('Please authenticate with Cycode first');
+    return;
+  }
+
+  const cycodeService = container.resolve<ICycodeService>(CycodeService);
+
+  const scanPromises = [];
+  scanPromises.push(cycodeService.startScanForCurrentProject(CliScanType.Secret));
+  scanPromises.push(cycodeService.startScanForCurrentProject(CliScanType.Sca));
+  scanPromises.push(cycodeService.startScanForCurrentProject(CliScanType.Iac));
+  scanPromises.push(cycodeService.startScanForCurrentProject(CliScanType.Sast));
+  await Promise.all(scanPromises);
+};

src/commands/run-all-scans-command.ts

@@ -0,0 +1,8 @@
+import { container } from 'tsyringe';
+import { IExtensionService } from '../services/extension-service';
+import { ExtensionServiceSymbol } from '../symbols';
+
+export default async () => {
+  const extension = container.resolve<IExtensionService>(ExtensionServiceSymbol);
+  await extension.treeDataProvider.collapseAll();
+};

src/commands/tree-view-collapse-all-command.ts

@@ -0,0 +1,8 @@
+import { container } from 'tsyringe';
+import { IExtensionService } from '../services/extension-service';
+import { ExtensionServiceSymbol } from '../symbols';
+
+export default async () => {
+  const extension = container.resolve<IExtensionService>(ExtensionServiceSymbol);
+  await extension.treeDataProvider.expandAll();
+};

src/commands/tree-view-expand-all-command.ts

@@ -22,7 +22,7 @@ import { ILoggerService } from './services/logger-service';
 import { IScanResultsService } from './services/scan-results-service';
 import { IExtensionService } from './services/extension-service';
 import { onProjectOpen } from './listeners/on-project-open';
-import { createTreeView } from './providers/tree-data';
+import { createTreeDataProvider } from './providers/tree-data';
 import { registerCommands } from './commands';
 import { registerOnDidSaveTextDocument } from './listeners/on-did-save-text-document';
 import { registerActivityBar } from './ui/views/activity-bar';
@@ -49,12 +49,12 @@ export async function activate(context: vscode.ExtensionContext) {
   logger.info('Cycode plugin is running');
 
   const diagnosticCollection = vscode.languages.createDiagnosticCollection(extensionName);
-  const treeView = createTreeView(context);
+  const treeDataProvider = createTreeDataProvider(context);
 
   const extension = container.resolve<IExtensionService>(ExtensionServiceSymbol);
   extension.extensionContext = context;
   extension.diagnosticCollection = diagnosticCollection;
-  extension.treeView = treeView;
+  extension.treeDataProvider = treeDataProvider;
 
   // refactor this to porper class
   const extensionStatusBar = statusBar.create();

src/extension.ts

@@ -1,25 +1,18 @@
 import * as vscode from 'vscode';
 import { TreeDataProvider } from './provider';
-import { BaseNode } from './nodes/base-node';
 
-export interface TreeView {
-  provider: TreeDataProvider;
-  view: vscode.TreeView<BaseNode>;
-}
-
-export const createTreeView = (context: vscode.ExtensionContext): TreeView => {
-  const provider = new TreeDataProvider();
-  const view = vscode.window.createTreeView(TreeDataProvider.viewType, {
-    treeDataProvider: provider,
-    showCollapseAll: true,
-    canSelectMany: false,
-  });
+export const createTreeDataProvider = (context: vscode.ExtensionContext): TreeDataProvider => {
+  const treeDataProvider = new TreeDataProvider();
+  treeDataProvider.treeView = vscode.window.createTreeView(
+    TreeDataProvider.viewType, { treeDataProvider: treeDataProvider },
+  );
 
   context.subscriptions.push(
     vscode.window.registerTreeDataProvider(
       TreeDataProvider.viewType,
-      provider,
+      treeDataProvider,
     ),
   );
-  return { view, provider };
+
+  return treeDataProvider;
 };

src/providers/tree-data/index.ts

@@ -17,7 +17,7 @@ export class DetectionNode extends BaseNode {
 
     this.command = {
       title: '',
-      command: VscodeCommands.OnTreeItemClick,
+      command: VscodeCommands.OnTreeViewDetectionNodeClickCommand,
       arguments: [scanType, detection],
     };
     this.collapsibleState = vscode.TreeItemCollapsibleState.None;

src/providers/tree-data/nodes/detection-node.ts

@@ -13,6 +13,5 @@ export class ScanTypeNode extends BaseNode {
 
     this.scanType = scanType;
     this.contextValue = `${scanType.toLowerCase()}ScanTypeNode`;
-    console.log(this.contextValue);
   }
 }

src/providers/tree-data/nodes/scan-type-node.ts

@@ -8,23 +8,29 @@ import { container } from 'tsyringe';
 import { IScanResultsService } from '../../services/scan-results-service';
 import { ScanResultsServiceSymbol } from '../../symbols';
 import { DetectionBase } from '../../cli/models/scan-result/detection-base';
+import { VscodeCommands } from '../../commands';
 
 export class TreeDataProvider implements vscode.TreeDataProvider<BaseNode> {
   public static readonly viewType = 'cycode.view.tree';
 
-  private _onDidChangeTreeData:
-  vscode.EventEmitter<BaseNode | undefined> = new vscode.EventEmitter<BaseNode | undefined>();
+  public treeView: vscode.TreeView<BaseNode>;
 
-  readonly onDidChangeTreeData:
-  vscode.Event<BaseNode | undefined> = this._onDidChangeTreeData.event;
+  private _onDidChangeTreeData: vscode.EventEmitter<void> = new vscode.EventEmitter<void>();
+  readonly onDidChangeTreeData: vscode.Event<void> = this._onDidChangeTreeData.event;
 
   private _createdRootNodes: ScanTypeNode[] = [];
   private _createdNodesToChildren = new Map<BaseNode, BaseNode[]>();
+  private _createdChildToParentNode = new Map<BaseNode, BaseNode>();
 
   getTreeItem(element: BaseNode): vscode.TreeItem {
     return element;
   }
 
+  getParent(element: BaseNode): vscode.ProviderResult<BaseNode> {
+    // must be implemented properly to be able to call this.treeView.reveal()
+    return this._createdChildToParentNode.get(element);
+  }
+
   getChildren(
     element?: BaseNode,
   ): Thenable<BaseNode[]> {
@@ -36,18 +42,14 @@ export class TreeDataProvider implements vscode.TreeDataProvider<BaseNode> {
   }
 
   private getSeverityWeight(severity: string): number {
-    switch (severity) {
-      case 'critical':
-        return 4;
-      case 'high':
-        return 3;
-      case 'medium':
-        return 2;
-      case 'low':
-        return 1;
-      default:
-        return 0;
-    }
+    const severityWeights: Record<string, number> = {
+      critical: 4,
+      high: 3,
+      medium: 2,
+      low: 1,
+    };
+
+    return severityWeights[severity.toLowerCase()] || 0;
   }
 
   private getScanTypeNodeSummary(sortedDetections: DetectionBase[]): string {
@@ -71,7 +73,7 @@ export class TreeDataProvider implements vscode.TreeDataProvider<BaseNode> {
     const detections = scanResultsService.getDetections(scanType);
 
     const severitySortedDetections = detections.sort((a, b) => {
-      return this.getSeverityWeight(b.severity.toLowerCase()) - this.getSeverityWeight(a.severity.toLowerCase());
+      return this.getSeverityWeight(b.severity) - this.getSeverityWeight(a.severity);
     });
     const groupedByFilepathDetection = severitySortedDetections
       .reduce<Map<string, DetectionBase[]>>((acc, detection) => {
@@ -91,22 +93,58 @@ export class TreeDataProvider implements vscode.TreeDataProvider<BaseNode> {
       const fileNode = new FileNode(filepath, detections.length);
       this._createdNodesToChildren.get(scanTypeNode)?.push(fileNode);
       this._createdNodesToChildren.set(fileNode, []);
+      this._createdChildToParentNode.set(fileNode, scanTypeNode);
+
       for (const detection of detections) {
         const detectionNode = new DetectionNode(scanType, detection);
         this._createdNodesToChildren.get(fileNode)?.push(detectionNode);
+        this._createdChildToParentNode.set(detectionNode, fileNode);
       }
     }
   }
 
   public refresh(): void {
     this._createdRootNodes = [];
     this._createdNodesToChildren.clear();
+    this._createdChildToParentNode.clear();
 
     this.createNodes(CliScanType.Secret);
     this.createNodes(CliScanType.Sca);
     this.createNodes(CliScanType.Iac);
     this.createNodes(CliScanType.Sast);
 
-    this._onDidChangeTreeData.fire(undefined);
+    this._onDidChangeTreeData.fire();
+  }
+
+  public async expandAll() {
+    /*
+     * vscode api limits rendering of expanding.
+     * Editing of collapsibleState field doesn't affect nodes that are visible already.
+     * To bypass this limitation, we can use reveal method instead.
+     * This method is limited to depth of 3; it is enough to our tree.
+     * More of the context: https://github.com/microsoft/vscode/issues/131955
+     */
+
+    const revealOptions = { expand: 3, select: false, focus: false };
+    const thenables: Thenable<void>[] = [];
+    for (const rootNode of this._createdRootNodes) {
+      thenables.push(this.treeView.reveal(rootNode, revealOptions));
+      for (const fileNode of this._createdNodesToChildren.get(rootNode) || []) {
+        thenables.push(this.treeView.reveal(fileNode, revealOptions));
+      }
+    }
+
+    await Promise.all(thenables);
+  }
+
+  public async collapseAll() {
+    /*
+     * It works excellently and always re-renders the tree view,
+     * even already visible elements,
+     * but unfortunately, there is no build-in "ExpandAll" command.
+     * That's why we use our own implementation.
+     */
+
+    await vscode.commands.executeCommand(VscodeCommands.WorkbenchTreeViewCollapseAll);
   }
 }

src/providers/tree-data/provider.ts

@@ -240,7 +240,7 @@ export class CliService implements ICliService {
         )
         .then((buttonPressed) => {
           if (buttonPressed === TrayNotificationTexts.OpenProblemsTab) {
-            vscode.commands.executeCommand(VscodeCommands.ShowProblemsTab);
+            vscode.commands.executeCommand(VscodeCommands.WorkbenchShowProblemsTab);
           }
         });
     } else if (onDemand) {