feat: Add Ollama AI integration for automatic mail template generation
✨ New Features: - 🤖 AI-powered mail template generation with Ollama - 📧 Test mail sending with preview - 🔧 Ollama server and model management - 🎨 Beautiful AI generation dialog in Templates page - ⚙️ Ollama settings panel with connection test Backend: - Add ollama.service.js - Ollama API integration - Add ollama.controller.js - Template generation endpoint - Add ollama.routes.js - /api/ollama/* routes - Support for multiple Ollama models (llama3.2, mistral, gemma) - JSON-formatted AI responses with subject + HTML body - Configurable server URL and model selection Frontend: - Settings: Ollama configuration panel - Server URL input - Model selection - Connection test with model listing - Templates: AI generation dialog - Company name, scenario, employee info inputs - Custom prompt for AI instructions - Auto-save to database - Test mail sending functionality Documentation: - OLLAMA_SETUP.md - Comprehensive setup guide - Installation instructions - Model recommendations - Usage examples - Troubleshooting Tech Stack: - Ollama API integration (REST) - Axios HTTP client - React dialogs with MUI - Self-hosted AI (privacy-friendly) - Zero external API dependencies Example Usage: Company: Garanti Bankası Scenario: Account security warning → AI generates professional phishing test mail in seconds!
This commit is contained in:
salvacybersec
221
backend/src/services/ollama.service.js
Normal file
221
backend/src/services/ollama.service.js
Normal file
@@ -0,0 +1,221 @@
|
||||
const axios = require('axios');
|
||||
const logger = require('../utils/logger');
|
||||
const { Settings } = require('../models');
|
||||
|
||||
class OllamaService {
|
||||
constructor() {
|
||||
this.serverUrl = null;
|
||||
this.model = null;
|
||||
this.initialized = false;
|
||||
}
|
||||
|
||||
/**
|
||||
* Initialize Ollama service with settings from database
|
||||
*/
|
||||
async initialize() {
|
||||
try {
|
||||
const serverUrlSetting = await Settings.findOne({
|
||||
where: { key: 'ollama_server_url' },
|
||||
});
|
||||
const modelSetting = await Settings.findOne({
|
||||
where: { key: 'ollama_model' },
|
||||
});
|
||||
|
||||
this.serverUrl = serverUrlSetting?.value || process.env.OLLAMA_URL || 'http://localhost:11434';
|
||||
this.model = modelSetting?.value || process.env.OLLAMA_MODEL || 'llama3.2';
|
||||
|
||||
this.initialized = true;
|
||||
logger.info(`Ollama service initialized: ${this.serverUrl} with model ${this.model}`);
|
||||
} catch (error) {
|
||||
logger.error('Failed to initialize Ollama service:', error);
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Test Ollama connection
|
||||
*/
|
||||
async testConnection() {
|
||||
if (!this.initialized) {
|
||||
await this.initialize();
|
||||
}
|
||||
|
||||
try {
|
||||
const response = await axios.get(`${this.serverUrl}/api/tags`, {
|
||||
timeout: 5000,
|
||||
});
|
||||
|
||||
return {
|
||||
success: true,
|
||||
models: response.data.models || [],
|
||||
message: 'Ollama bağlantısı başarılı',
|
||||
};
|
||||
} catch (error) {
|
||||
logger.error('Ollama connection test failed:', error.message);
|
||||
return {
|
||||
success: false,
|
||||
error: error.message,
|
||||
message: 'Ollama sunucusuna bağlanılamadı',
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* List available models
|
||||
*/
|
||||
async listModels() {
|
||||
if (!this.initialized) {
|
||||
await this.initialize();
|
||||
}
|
||||
|
||||
try {
|
||||
const response = await axios.get(`${this.serverUrl}/api/tags`, {
|
||||
timeout: 5000,
|
||||
});
|
||||
|
||||
return response.data.models || [];
|
||||
} catch (error) {
|
||||
logger.error('Failed to list Ollama models:', error.message);
|
||||
throw new Error('Ollama model listesi alınamadı: ' + error.message);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Generate mail template using Ollama
|
||||
* @param {Object} params - Template generation parameters
|
||||
* @param {string} params.company_name - Target company name
|
||||
* @param {string} params.scenario - Phishing scenario type
|
||||
* @param {string} params.employee_info - Employee information (optional)
|
||||
* @param {string} params.custom_prompt - Custom instructions (optional)
|
||||
*/
|
||||
async generateMailTemplate(params) {
|
||||
if (!this.initialized) {
|
||||
await this.initialize();
|
||||
}
|
||||
|
||||
const { company_name, scenario, employee_info, custom_prompt } = params;
|
||||
|
||||
// Build the prompt
|
||||
const systemPrompt = `Sen profesyonel bir güvenlik uzmanısın ve phishing test maileri oluşturuyorsun.
|
||||
Amacın gerçekçi, ikna edici ancak zararsız test mailleri oluşturmak.
|
||||
Mail şablonları HTML formatında olmalı ve modern, profesyonel görünmeli.
|
||||
Şablon içinde {{company_name}} ve {{employee_name}} placeholder'ları kullan.`;
|
||||
|
||||
let userPrompt = `Aşağıdaki bilgilere göre bir phishing test mail şablonu oluştur:
|
||||
|
||||
Şirket: ${company_name}
|
||||
Senaryo: ${scenario}`;
|
||||
|
||||
if (employee_info) {
|
||||
userPrompt += `\nÇalışan Bilgisi: ${employee_info}`;
|
||||
}
|
||||
|
||||
if (custom_prompt) {
|
||||
userPrompt += `\nEk Talimatlar: ${custom_prompt}`;
|
||||
}
|
||||
|
||||
userPrompt += `
|
||||
|
||||
ÖNEMLI:
|
||||
1. Yanıtını JSON formatında ver
|
||||
2. İki alan olmalı: "subject" (konu) ve "body" (HTML mail içeriği)
|
||||
3. Body HTML formatında, modern ve profesyonel olmalı
|
||||
4. {{company_name}} ve {{employee_name}} placeholder'larını kullan
|
||||
5. Gerçekçi ve ikna edici olmalı
|
||||
6. Link için {{tracking_url}} placeholder'ını kullan
|
||||
|
||||
Örnek format:
|
||||
{
|
||||
"subject": "Mail konusu buraya",
|
||||
"body": "<html>Mail içeriği buraya</html>"
|
||||
}`;
|
||||
|
||||
try {
|
||||
logger.info(`Generating template for company: ${company_name}, scenario: ${scenario}`);
|
||||
|
||||
const response = await axios.post(
|
||||
`${this.serverUrl}/api/chat`,
|
||||
{
|
||||
model: this.model,
|
||||
messages: [
|
||||
{
|
||||
role: 'system',
|
||||
content: systemPrompt,
|
||||
},
|
||||
{
|
||||
role: 'user',
|
||||
content: userPrompt,
|
||||
},
|
||||
],
|
||||
stream: false,
|
||||
format: 'json',
|
||||
},
|
||||
{
|
||||
timeout: 120000, // 2 minutes timeout for AI generation
|
||||
}
|
||||
);
|
||||
|
||||
const aiResponse = response.data.message.content;
|
||||
logger.info('Ollama response received');
|
||||
|
||||
// Parse JSON response
|
||||
let templateData;
|
||||
try {
|
||||
templateData = JSON.parse(aiResponse);
|
||||
} catch (parseError) {
|
||||
// If JSON parsing fails, try to extract JSON from response
|
||||
const jsonMatch = aiResponse.match(/\{[\s\S]*\}/);
|
||||
if (jsonMatch) {
|
||||
templateData = JSON.parse(jsonMatch[0]);
|
||||
} else {
|
||||
throw new Error('AI yanıtı JSON formatında değil');
|
||||
}
|
||||
}
|
||||
|
||||
if (!templateData.subject || !templateData.body) {
|
||||
throw new Error('AI yanıtında subject veya body eksik');
|
||||
}
|
||||
|
||||
return {
|
||||
subject_template: templateData.subject,
|
||||
body_template: templateData.body,
|
||||
generated_by: 'ollama',
|
||||
model: this.model,
|
||||
};
|
||||
} catch (error) {
|
||||
logger.error('Failed to generate template with Ollama:', error.message);
|
||||
throw new Error('Template oluşturulamadı: ' + error.message);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Generate a simple text completion
|
||||
*/
|
||||
async generate(prompt) {
|
||||
if (!this.initialized) {
|
||||
await this.initialize();
|
||||
}
|
||||
|
||||
try {
|
||||
const response = await axios.post(
|
||||
`${this.serverUrl}/api/generate`,
|
||||
{
|
||||
model: this.model,
|
||||
prompt: prompt,
|
||||
stream: false,
|
||||
},
|
||||
{
|
||||
timeout: 60000,
|
||||
}
|
||||
);
|
||||
|
||||
return response.data.response;
|
||||
} catch (error) {
|
||||
logger.error('Ollama generation failed:', error.message);
|
||||
throw new Error('Ollama yanıt veremedi: ' + error.message);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = new OllamaService();
|
||||
|
||||
Reference in New Issue
Block a user