293 lines
8.6 KiB
TypeScript
293 lines
8.6 KiB
TypeScript
// src/pages/api/ai/query.ts
|
|
// src/pages/api/ai/query.ts
|
|
import type { APIRoute } from 'astro';
|
|
import { getSessionFromRequest, verifySession } from '../../../utils/auth.js';
|
|
import { promises as fs } from 'fs';
|
|
import { load } from 'js-yaml';
|
|
import path from 'path';
|
|
|
|
export const prerender = false;
|
|
|
|
function getEnv(key: string): string {
|
|
const value = process.env[key];
|
|
if (!value) {
|
|
throw new Error(`Missing environment variable: ${key}`);
|
|
}
|
|
return value;
|
|
}
|
|
|
|
const AI_MODEL = getEnv('AI_MODEL');
|
|
// Rate limiting store (in production, use Redis)
|
|
const rateLimitStore = new Map<string, { count: number; resetTime: number }>();
|
|
const RATE_LIMIT_WINDOW = 60 * 1000; // 1 minute
|
|
const RATE_LIMIT_MAX = 10; // 10 requests per minute per user
|
|
|
|
// Input validation and sanitization
|
|
function sanitizeInput(input: string): string {
|
|
// Remove potential prompt injection patterns
|
|
const dangerous = [
|
|
/ignore\s+previous\s+instructions?/gi,
|
|
/new\s+instructions?:/gi,
|
|
/system\s*:/gi,
|
|
/assistant\s*:/gi,
|
|
/human\s*:/gi,
|
|
/<\s*\/?system\s*>/gi,
|
|
/```\s*system/gi,
|
|
];
|
|
|
|
let sanitized = input.trim();
|
|
dangerous.forEach(pattern => {
|
|
sanitized = sanitized.replace(pattern, '[FILTERED]');
|
|
});
|
|
|
|
// Limit length
|
|
return sanitized.slice(0, 2000);
|
|
}
|
|
|
|
// Strip markdown code blocks from AI response
|
|
function stripMarkdownJson(content: string): string {
|
|
// Remove ```json and ``` wrappers
|
|
return content
|
|
.replace(/^```json\s*/i, '')
|
|
.replace(/\s*```\s*$/, '')
|
|
.trim();
|
|
}
|
|
|
|
// Rate limiting check
|
|
function checkRateLimit(userId: string): boolean {
|
|
const now = Date.now();
|
|
const userLimit = rateLimitStore.get(userId);
|
|
|
|
if (!userLimit || now > userLimit.resetTime) {
|
|
rateLimitStore.set(userId, { count: 1, resetTime: now + RATE_LIMIT_WINDOW });
|
|
return true;
|
|
}
|
|
|
|
if (userLimit.count >= RATE_LIMIT_MAX) {
|
|
return false;
|
|
}
|
|
|
|
userLimit.count++;
|
|
return true;
|
|
}
|
|
|
|
// Load tools database
|
|
async function loadToolsDatabase() {
|
|
try {
|
|
const yamlPath = path.join(process.cwd(), 'src/data/tools.yaml');
|
|
const yamlContent = await fs.readFile(yamlPath, 'utf8');
|
|
return load(yamlContent) as any;
|
|
} catch (error) {
|
|
console.error('Failed to load tools database:', error);
|
|
throw new Error('Database unavailable');
|
|
}
|
|
}
|
|
|
|
// Create system prompt
|
|
function createSystemPrompt(toolsData: any): string {
|
|
const toolsList = toolsData.tools.map((tool: any) => ({
|
|
name: tool.name,
|
|
description: tool.description,
|
|
domains: tool.domains,
|
|
phases: tool.phases,
|
|
platforms: tool.platforms,
|
|
skillLevel: tool.skillLevel,
|
|
license: tool.license,
|
|
tags: tool.tags,
|
|
projectUrl: tool.projectUrl ? 'self-hosted' : 'external'
|
|
}));
|
|
|
|
// Dynamically build phases list from configuration
|
|
const phasesDescription = toolsData.phases.map((phase: any) =>
|
|
`- ${phase.id}: ${phase.name}`
|
|
).join('\n');
|
|
|
|
// Dynamically build domains list from configuration
|
|
const domainsDescription = toolsData.domains.map((domain: any) =>
|
|
`- ${domain.id}: ${domain.name}`
|
|
).join('\n');
|
|
|
|
return `Du bist ein DFIR (Digital Forensics and Incident Response) Experte, der Ermittlern bei der Toolauswahl hilft.
|
|
|
|
VERFÜGBARE TOOLS DATABASE:
|
|
${JSON.stringify(toolsList, null, 2)}
|
|
|
|
UNTERSUCHUNGSPHASEN (NIST Framework):
|
|
${phasesDescription}
|
|
|
|
FORENSISCHE DOMÄNEN:
|
|
${domainsDescription}
|
|
|
|
PRIORITÄTEN:
|
|
1. Self-hosted Tools (projectUrl: "self-hosted") bevorzugen
|
|
2. Open Source Tools bevorzugen (license != "Proprietary")
|
|
3. Maximal 3 Tools pro Phase empfehlen
|
|
4. Deutsche Antworten für deutsche Anfragen, English for English queries
|
|
|
|
ANTWORT-FORMAT (strict JSON):
|
|
{
|
|
"scenario_analysis": "Detaillierte Analyse des Szenarios auf Deutsch/English",
|
|
"recommended_tools": [
|
|
{
|
|
"name": "EXAKTER Name aus der Database",
|
|
"priority": "high|medium|low",
|
|
"phase": "data-collection|examination|analysis|reporting",
|
|
"justification": "Warum dieses Tool für dieses Szenario geeignet ist"
|
|
}
|
|
],
|
|
"workflow_suggestion": "Vorgeschlagener Untersuchungsablauf",
|
|
"additional_notes": "Wichtige Überlegungen und Hinweise"
|
|
}
|
|
|
|
Antworte NUR mit validen JSON. Keine zusätzlichen Erklärungen außerhalb des JSON.`;
|
|
}
|
|
|
|
export const POST: APIRoute = async ({ request }) => {
|
|
try {
|
|
// Check if authentication is required
|
|
const authRequired = process.env.AUTHENTICATION_NECESSARY !== 'false';
|
|
let userId = 'test-user';
|
|
|
|
if (authRequired) {
|
|
// Authentication check
|
|
const sessionToken = getSessionFromRequest(request);
|
|
if (!sessionToken) {
|
|
return new Response(JSON.stringify({ error: 'Authentication required' }), {
|
|
status: 401,
|
|
headers: { 'Content-Type': 'application/json' }
|
|
});
|
|
}
|
|
|
|
const session = await verifySession(sessionToken);
|
|
if (!session) {
|
|
return new Response(JSON.stringify({ error: 'Invalid session' }), {
|
|
status: 401,
|
|
headers: { 'Content-Type': 'application/json' }
|
|
});
|
|
}
|
|
|
|
userId = session.userId;
|
|
}
|
|
|
|
// Rate limiting
|
|
if (!checkRateLimit(userId)) {
|
|
return new Response(JSON.stringify({ error: 'Rate limit exceeded' }), {
|
|
status: 429,
|
|
headers: { 'Content-Type': 'application/json' }
|
|
});
|
|
}
|
|
|
|
// Parse request body
|
|
const body = await request.json();
|
|
const { query } = body;
|
|
|
|
if (!query || typeof query !== 'string') {
|
|
return new Response(JSON.stringify({ error: 'Query required' }), {
|
|
status: 400,
|
|
headers: { 'Content-Type': 'application/json' }
|
|
});
|
|
}
|
|
|
|
// Sanitize input
|
|
const sanitizedQuery = sanitizeInput(query);
|
|
if (sanitizedQuery.includes('[FILTERED]')) {
|
|
return new Response(JSON.stringify({ error: 'Invalid input detected' }), {
|
|
status: 400,
|
|
headers: { 'Content-Type': 'application/json' }
|
|
});
|
|
}
|
|
|
|
// Load tools database
|
|
const toolsData = await loadToolsDatabase();
|
|
|
|
// Create AI request
|
|
const systemPrompt = createSystemPrompt(toolsData);
|
|
|
|
const aiResponse = await fetch(process.env.AI_API_ENDPOINT + '/v1/chat/completions', {
|
|
method: 'POST',
|
|
headers: {
|
|
'Content-Type': 'application/json',
|
|
'Authorization': `Bearer ${process.env.AI_API_KEY}`
|
|
},
|
|
body: JSON.stringify({
|
|
model: AI_MODEL, // or whatever model is available
|
|
messages: [
|
|
{
|
|
role: 'system',
|
|
content: systemPrompt
|
|
},
|
|
{
|
|
role: 'user',
|
|
content: sanitizedQuery
|
|
}
|
|
],
|
|
max_tokens: 2000,
|
|
temperature: 0.3
|
|
})
|
|
});
|
|
|
|
|
|
if (!aiResponse.ok) {
|
|
console.error('AI API error:', await aiResponse.text());
|
|
return new Response(JSON.stringify({ error: 'AI service unavailable' }), {
|
|
status: 503,
|
|
headers: { 'Content-Type': 'application/json' }
|
|
});
|
|
}
|
|
|
|
const aiData = await aiResponse.json();
|
|
const aiContent = aiData.choices?.[0]?.message?.content;
|
|
|
|
if (!aiContent) {
|
|
return new Response(JSON.stringify({ error: 'No response from AI' }), {
|
|
status: 503,
|
|
headers: { 'Content-Type': 'application/json' }
|
|
});
|
|
}
|
|
|
|
// Parse AI JSON response
|
|
let recommendation;
|
|
try {
|
|
const cleanedContent = stripMarkdownJson(aiContent);
|
|
recommendation = JSON.parse(cleanedContent);
|
|
} catch (error) {
|
|
console.error('Failed to parse AI response:', aiContent);
|
|
return new Response(JSON.stringify({ error: 'Invalid AI response format' }), {
|
|
status: 503,
|
|
headers: { 'Content-Type': 'application/json' }
|
|
});
|
|
}
|
|
|
|
// Validate tool names against database
|
|
const validToolNames = new Set(toolsData.tools.map((t: any) => t.name));
|
|
const validatedRecommendation = {
|
|
...recommendation,
|
|
recommended_tools: recommendation.recommended_tools?.filter((tool: any) => {
|
|
if (!validToolNames.has(tool.name)) {
|
|
console.warn(`AI recommended unknown tool: ${tool.name}`);
|
|
return false;
|
|
}
|
|
return true;
|
|
}) || []
|
|
};
|
|
|
|
// Log successful query
|
|
console.log(`[AI Query] User: ${userId}, Query length: ${sanitizedQuery.length}, Tools: ${validatedRecommendation.recommended_tools.length}`);
|
|
|
|
return new Response(JSON.stringify({
|
|
success: true,
|
|
recommendation: validatedRecommendation,
|
|
query: sanitizedQuery
|
|
}), {
|
|
status: 200,
|
|
headers: { 'Content-Type': 'application/json' }
|
|
});
|
|
|
|
} catch (error) {
|
|
console.error('AI query error:', error);
|
|
return new Response(JSON.stringify({ error: 'Internal server error' }), {
|
|
status: 500,
|
|
headers: { 'Content-Type': 'application/json' }
|
|
});
|
|
}
|
|
}; |