first ai implementation, oidc done
This commit is contained in:
overcuriousity
293
src/pages/api/ai/query.ts
Normal file
293
src/pages/api/ai/query.ts
Normal file
@@ -0,0 +1,293 @@
|
||||
// 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' }
|
||||
});
|
||||
}
|
||||
};
|
||||
Reference in New Issue
Block a user