pury
Version:
🛡️ AI-powered security scanner with advanced threat detection, dual reporting system (detailed & summary), and comprehensive code analysis
229 lines • 9.31 kB
JavaScript
import { GoogleGenAI } from '@google/genai';
import { FindingType, Severity } from '../types/index.js';
import { logger } from '../utils/logger.js';
import { ANALYSIS_PROMPTS } from './prompts/index.js';
export class GeminiClient {
genAI;
config;
constructor(config) {
this.config = {
model: 'gemini-2.5-flash',
temperature: 0.1,
maxTokens: 2048,
topP: 0.8,
...config
};
this.genAI = new GoogleGenAI({
apiKey: this.config.apiKey
});
}
async analyzeCode(request) {
try {
const prompt = this.buildPrompt(request);
logger.debug(`Analyzing ${request.filePath} with Gemini AI`);
const response = await this.genAI.models.generateContent({
model: this.config.model,
contents: prompt
});
const text = response.text || '';
return this.parseResponse(text, request.filePath);
}
catch (error) {
logger.error(`Gemini API error: ${error.message}`);
throw new Error(`AI analysis failed: ${error.message}`);
}
}
async analyzeCodeStream(request) {
const prompt = this.buildPrompt(request);
const response = await this.genAI.models.generateContentStream({
model: this.config.model,
contents: prompt
});
async function* streamGenerator() {
for await (const chunk of response) {
if (chunk.text) {
yield chunk.text;
}
}
}
return streamGenerator();
}
async analyzeMultipleFiles(requests) {
const results = [];
// Process in batches to avoid rate limiting
const batchSize = 3;
const batches = this.chunkArray(requests, batchSize);
for (const batch of batches) {
const batchPromises = batch.map(async (request) => this.analyzeCode(request));
const batchResults = await Promise.allSettled(batchPromises);
for (const result of batchResults) {
if (result.status === 'fulfilled') {
results.push(result.value);
}
else {
logger.warn(`Batch analysis failed: ${result.reason}`);
// Add empty result for failed analysis
results.push({ findings: [], confidence: 0 });
}
}
// Add delay between batches to respect rate limits
if (batches.indexOf(batch) < batches.length - 1) {
await this.delay(1000);
}
}
return results;
}
buildPrompt(request) {
const { code, filePath, analysisType, context } = request;
let basePrompt = 'You are a security expert analyzing code for potential threats and issues.\n\n';
// Add specific analysis prompts based on requested types
const promptParts = [];
for (const type of analysisType) {
if (ANALYSIS_PROMPTS[type]) {
promptParts.push(ANALYSIS_PROMPTS[type]);
}
}
basePrompt += promptParts.join('\n\n');
basePrompt += '\n\nPlease analyze the following code and respond in JSON format:\n';
basePrompt += '{\n';
basePrompt += ' "findings": [\n';
basePrompt += ' {\n';
basePrompt += ' "type": "malware|vulnerability|secret|code_quality",\n';
basePrompt += ' "severity": "low|medium|high|critical",\n';
basePrompt += ' "title": "Brief title",\n';
basePrompt += ' "description": "Detailed description",\n';
basePrompt += ' "line": 10,\n';
basePrompt += ' "evidence": "Code snippet that triggered this finding",\n';
basePrompt += ' "suggestion": "How to fix this issue"\n';
basePrompt += ' }\n';
basePrompt += ' ],\n';
basePrompt += ' "confidence": 0.95,\n';
basePrompt += ' "reasoning": "Explanation of analysis approach"\n';
basePrompt += '}\n\n';
if (context) {
basePrompt += `Context: ${context}\n\n`;
}
basePrompt += `File: ${filePath}\n\n`;
basePrompt += 'Code to analyze:\n```\n';
basePrompt += code;
basePrompt += '\n```';
return basePrompt;
}
parseResponse(responseText, filePath) {
try {
// Clean up the response to extract JSON
let jsonText = responseText.trim();
// Remove markdown code blocks if present
jsonText = jsonText.replace(/```json\n?/g, '').replace(/```\n?/g, '');
// Extract JSON from response
const jsonMatch = /\{[\s\S]*\}/.exec(jsonText);
if (!jsonMatch) {
throw new Error('No JSON found in response');
}
let jsonString = jsonMatch[0];
// Fix common JSON issues
jsonString = jsonString
.replace(/\\'/g, "'") // Fix escaped single quotes
.replace(/\\\\/g, '\\') // Fix double escaped backslashes
.replace(/\\n/g, '\\n') // Keep newlines properly escaped
.replace(/\\t/g, '\\t') // Keep tabs properly escaped
.replace(/\\r/g, '\\r') // Keep carriage returns properly escaped
.replace(/\n\s*}/g, '\n}') // Clean up closing braces
.replace(/,\s*}/g, '}') // Remove trailing commas
.replace(/,\s*]/g, ']'); // Remove trailing commas in arrays
const parsed = JSON.parse(jsonString);
const findings = (parsed.findings || []).map((finding, index) => ({
id: `ai-${Date.now()}-${index}`,
type: this.validateFindingType(finding.type),
severity: this.validateSeverity(finding.severity),
title: finding.title || 'AI-detected issue',
description: finding.description || '',
file: filePath,
line: finding.line,
column: finding.column,
evidence: finding.evidence,
suggestion: finding.suggestion,
references: finding.references || []
}));
return {
findings,
confidence: Math.min(1, Math.max(0, parsed.confidence || 0.5)),
reasoning: parsed.reasoning
};
}
catch (error) {
logger.warn(`Failed to parse AI response: ${error.message}`);
logger.debug(`Raw response (first 500 chars): ${responseText.substring(0, 500)}`);
// Try to extract any meaningful information from the response
const emergencyFindings = [];
// Look for common security keywords in the raw response
const securityKeywords = [
'vulnerability',
'security',
'malware',
'suspicious',
'dangerous',
'risk'
];
const foundKeywords = securityKeywords.filter(keyword => responseText.toLowerCase().includes(keyword));
if (foundKeywords.length > 0) {
emergencyFindings.push({
id: `ai-emergency-${Date.now()}`,
type: 'security',
severity: 'medium',
title: 'AI Analysis Issue - Manual Review Needed',
description: `AI response parsing failed, but security-related keywords were found: ${foundKeywords.join(', ')}. Manual review recommended.`,
file: filePath,
line: 1,
evidence: responseText.substring(0, 200),
suggestion: 'Review the file manually for potential security issues.',
references: []
});
}
return {
findings: emergencyFindings,
confidence: 0.1,
reasoning: `Failed to parse AI response: ${error.message}`
};
}
}
validateFindingType(type) {
return Object.values(FindingType).includes(type)
? type
: FindingType.CODE_QUALITY;
}
validateSeverity(severity) {
return Object.values(Severity).includes(severity)
? severity
: Severity.MEDIUM;
}
chunkArray(array, size) {
const chunks = [];
for (let i = 0; i < array.length; i += size) {
chunks.push(array.slice(i, i + size));
}
return chunks;
}
async delay(ms) {
return new Promise(resolve => setTimeout(resolve, ms));
}
getConfig() {
return { ...this.config };
}
async testConnection() {
try {
const testCode = 'console.log("hello world");';
const request = {
code: testCode,
filePath: 'test.js',
analysisType: [FindingType.CODE_QUALITY]
};
await this.analyzeCode(request);
return true;
}
catch {
return false;
}
}
}
//# sourceMappingURL=gemini-client.js.map