Mojo-Maniac/Cerberus-Backend

1import * as vscode from 'vscode';2 3export interface Vulnerability {4	file: string;5	status: string;6	error?: string;7	result?: string;8	line?: number;9	endLine?: number;10	type?: string;11	severity?: string;12	description?: string;13	originalCode?: string;14	fixedCode?: string;15}16 17export class VulnerabilityItem extends vscode.TreeItem {18	constructor(19		public readonly label: string,20		public readonly collapsibleState: vscode.TreeItemCollapsibleState,21		public readonly vulnerability?: Vulnerability,22		public readonly command?: vscode.Command,23		contextValue?: string,24		public readonly index?: number25	) {26		super(label, collapsibleState);27 28		if (contextValue) {29			this.contextValue = contextValue;30		}31 32		if (vulnerability?.status === 'analyzed' && vulnerability?.type) {33			// Show vulnerability type and severity with fix indicator34			const fixAvailable = vulnerability.fixedCode ? ' ― Apply Fix →' : '';35			this.description = this.getSeverityLabel(vulnerability.severity) + fixAvailable;36			this.iconPath = this.getSeverityIcon(vulnerability.severity);37			this.tooltip = this.buildTooltip(vulnerability);38		} else if (vulnerability?.status === 'analyzed' && vulnerability?.result) {39			this.description = 'Click to apply fix';40			this.iconPath = new vscode.ThemeIcon('error', new vscode.ThemeColor('errorForeground'));41			this.tooltip = vulnerability.result.substring(0, 500);42		} else if (vulnerability?.status === 'error') {43			this.description = 'Analysis error';44			this.iconPath = new vscode.ThemeIcon('warning', new vscode.ThemeColor('warningForeground'));45			this.tooltip = vulnerability.error || 'Unknown error';46		} else {47			this.iconPath = new vscode.ThemeIcon('file-code');48			this.tooltip = 'Click to open file';49		}50	}51 52	private getSeverityLabel(severity?: string): string {53		switch (severity) {54			case 'critical': return '🔴 Critical';55			case 'high': return '🟠 High';56			case 'medium': return '🟡 Medium';57			case 'low': return '🟢 Low';58			default: return 'Issue';59		}60	}61 62	private getSeverityIcon(severity?: string): vscode.ThemeIcon {63		switch (severity) {64			case 'critical':65				return new vscode.ThemeIcon('error', new vscode.ThemeColor('errorForeground'));66			case 'high':67				return new vscode.ThemeIcon('warning', new vscode.ThemeColor('problemsWarningIcon.foreground'));68			case 'medium':69				return new vscode.ThemeIcon('info', new vscode.ThemeColor('notificationsInfoIcon.foreground'));70			case 'low':71				return new vscode.ThemeIcon('pass', new vscode.ThemeColor('testing.iconPassed'));72			default:73				return new vscode.ThemeIcon('circle-outline');74		}75	}76 77	private buildTooltip(vuln: Vulnerability): string {78		const lines = [79			`Type: ${vuln.type || 'Unknown'}`,80			`Severity: ${vuln.severity || 'Unknown'}`,81			`Location: Line ${vuln.line}${vuln.endLine && vuln.endLine !== vuln.line ? `-${vuln.endLine}` : ''}`,82			'',83			'Original Code:',84			vuln.originalCode || '(not available)',85			'',86			'Fixed Code:',87			vuln.fixedCode || '(not available)'88		];89		return lines.join('\n');90	}91}92 93export class VulnerabilityTreeDataProvider implements vscode.TreeDataProvider<VulnerabilityItem> {94	private _onDidChangeTreeData: vscode.EventEmitter<VulnerabilityItem | undefined | null | void> =95		new vscode.EventEmitter<VulnerabilityItem | undefined | null | void>();96	readonly onDidChangeTreeData: vscode.Event<VulnerabilityItem | undefined | null | void> =97		this._onDidChangeTreeData.event;98 99	private vulnerabilities: Vulnerability[] = [];100 101	setVulnerabilities(vulnerabilities: Vulnerability[]): void {102		this.vulnerabilities = vulnerabilities;103		this._onDidChangeTreeData.fire(null);104	}105 106	getVulnerabilitiesForFile(fileName: string): Vulnerability[] {107		return this.vulnerabilities.filter((v: Vulnerability) => {108			const vulnFileName = v.file.split('\\').pop()?.split('/').pop() || v.file;109			return vulnFileName === fileName;110		});111	}112 113	getVulnerabilityByIndex(fileName: string, index: number): Vulnerability | undefined {114		const fileVulns = this.getVulnerabilitiesForFile(fileName);115		return fileVulns[index];116	}117 118	getTreeItem(element: VulnerabilityItem): vscode.TreeItem {119		return element;120	}121 122	getChildren(element?: VulnerabilityItem): Thenable<VulnerabilityItem[]> {123		if (!element) {124			// Root level - show list of files with vulnerabilities125			const fileMap = new Map<string, Vulnerability[]>();126 127			this.vulnerabilities.forEach((vuln: Vulnerability) => {128				if (!fileMap.has(vuln.file)) {129					fileMap.set(vuln.file, []);130				}131				fileMap.get(vuln.file)!.push(vuln);132			});133 134			const items = Array.from(fileMap.entries()).map(135				([file, vulns]) => {136					const errorCount = vulns.filter(v => v.status === 'error').length;137					const issueCount = vulns.filter(v => v.status === 'analyzed').length;138					const description = errorCount > 0139						? `${issueCount} issues, ${errorCount} errors`140						: `${issueCount} issue${issueCount !== 1 ? 's' : ''}`;141 142					const item = new VulnerabilityItem(143						this.getFileName(file),144						vscode.TreeItemCollapsibleState.Expanded,145						undefined,146						{147							command: 'vscode.open',148							title: 'Open file',149							arguments: [vscode.Uri.file(file)],150						},151						'file'152					);153					item.description = description;154					return item;155				}156			);157 158			return Promise.resolve(items);159		} else if (element.vulnerability === undefined) {160			// File level - show vulnerabilities in this file161			const fileName = element.label;162			const fileVulns = this.vulnerabilities.filter((v: Vulnerability) =>163				this.getFileName(v.file) === fileName164			);165 166			const items = fileVulns.map(167				(vuln, index) => {168					// Build a descriptive label169					let label: string;170					if (vuln.type) {171						label = vuln.type;172						if (vuln.line) {173							label += ` (Line ${vuln.line})`;174						}175					} else {176						label = `Issue ${index + 1}${vuln.status === 'error' ? ' (Error)' : ''}`;177					}178 179					return new VulnerabilityItem(180						label,181						vscode.TreeItemCollapsibleState.Collapsed,182						vuln,183						undefined,184						vuln.status === 'analyzed' ? 'vulnerability' : undefined,185						index186					);187				}188			);189 190			return Promise.resolve(items);191		} else {192			// Vulnerability level - show details193			const vuln = element.vulnerability;194			const items: VulnerabilityItem[] = [];195 196			if (vuln.status === 'analyzed') {197				// Show vulnerability details198				if (vuln.description) {199					const descItem = new VulnerabilityItem(200						`📋 ${vuln.description}`,201						vscode.TreeItemCollapsibleState.None,202						undefined203					);204					items.push(descItem);205				}206 207				// Show original code208				if (vuln.originalCode) {209					const origHeader = new VulnerabilityItem(210						'❌ Vulnerable Code:',211						vscode.TreeItemCollapsibleState.None,212						undefined213					);214					items.push(origHeader);215 216					const origLines = vuln.originalCode.split('\n').slice(0, 5);217					const origPreview = origLines.map(l => `  ${l}`).join('\n');218					const origItem = new VulnerabilityItem(219						origPreview + (vuln.originalCode.split('\n').length > 5 ? '\n  ...' : ''),220						vscode.TreeItemCollapsibleState.None,221						undefined222					);223					items.push(origItem);224				}225 226				// Show fixed code227				if (vuln.fixedCode) {228					const fixHeader = new VulnerabilityItem(229						'✅ Fixed Code:',230						vscode.TreeItemCollapsibleState.None,231						undefined232					);233					items.push(fixHeader);234 235					const fixLines = vuln.fixedCode.split('\n').slice(0, 5);236					const fixPreview = fixLines.map(l => `  ${l}`).join('\n');237					const fixItem = new VulnerabilityItem(238						fixPreview + (vuln.fixedCode.split('\n').length > 5 ? '\n  ...' : ''),239						vscode.TreeItemCollapsibleState.None,240						undefined241					);242					items.push(fixItem);243				} else if (vuln.result) {244					// Fallback to old format245					const lines = vuln.result.split('\n');246					const previewLines = lines.slice(0, 10);247					const preview = previewLines.join('\n') + (lines.length > 10 ? '\n...' : '');248 249					items.push(250						new VulnerabilityItem(251							'Fixed Code Preview:',252							vscode.TreeItemCollapsibleState.None,253							undefined254						)255					);256 257					items.push(258						new VulnerabilityItem(259							preview,260							vscode.TreeItemCollapsibleState.None,261							undefined262						)263					);264				}265			} else if (vuln.error) {266				items.push(267					new VulnerabilityItem(268						`Error: ${vuln.error}`,269						vscode.TreeItemCollapsibleState.None,270						undefined271					)272				);273			}274 275			return Promise.resolve(items);276		}277	}278 279	private getFileName(filePath: string): string {280		return filePath.split('\\').pop()?.split('/').pop() || filePath;281	}282}283