prompt-injection-lab/src/core.py

74 lines
3.4 KiB
Python

import json
import os
import smtplib
from email.message import EmailMessage
from loguru import logger
class AttackManager:
def __init__(self):
# Загружаем настройки и сразу проверяем их наличие
self.webhook_url = os.getenv("WEBHOOK_URL")
self.sender_email = os.getenv("EMAIL_SENDER")
self.sender_password = os.getenv("EMAIL_PASSWORD")
self.smtp_server = os.getenv("SMTP_SERVER", "smtp.yandex.ru")
self.smtp_port = int(os.getenv("SMTP_PORT", 465))
# ПРОВЕРКА: Если чего-то не хватает, бросаем ошибку
# Это уберет ошибку Pylance, так как после этого блока переменные точно str
if not all([self.webhook_url, self.sender_email, self.sender_password]):
logger.critical(
"❌ Ошибка: Проверьте файл .env! Отсутствуют обязательные параметры."
)
raise ValueError("Missing environment variables")
self.payloads_path = os.path.join("data", "payloads.json")
def get_payloads(self) -> list[dict]:
"""Загружает все доступные атаки из JSON."""
with open(self.payloads_path, "r", encoding="utf-8") as f:
return json.load(f)
def prepare_message(self, target_email: str, attack: dict) -> EmailMessage:
"""Подготавливает письмо на основе шаблона атаки."""
# Заменяем плейсхолдер на реальный URL
# Мы уверены, что self.webhook_url это str благодаря проверке в __init__
final_body = attack["body"].replace("{{WEBHOOK_URL}}", str(self.webhook_url))
msg = EmailMessage()
msg.set_content(final_body)
msg["Subject"] = attack["subject"]
msg["From"] = self.sender_email
msg["To"] = target_email
return msg
def execute_attacks(self, target_email: str, payload_ids: list[str] | None = None):
"""
Отправляет выбранные атаки (или все сразу, если payload_ids не указан).
"""
payloads = self.get_payloads()
# Если список ID не передан, берем все атаки из файла
to_send = [p for p in payloads if payload_ids is None or p["id"] in payload_ids]
if not to_send:
logger.warning("Нет атак для отправки. Проверьте ID.")
return
try:
logger.info(f"Connecting to {self.smtp_server}...")
with smtplib.SMTP_SSL(self.smtp_server, self.smtp_port) as server:
# Теперь Pylance спокоен, так как мы гарантировали, что это str
server.login(str(self.sender_email), str(self.sender_password))
for attack in to_send:
msg = self.prepare_message(target_email, attack)
server.send_message(msg)
logger.success(f"✅ Attack '{attack['id']}' sent to {target_email}")
logger.info(f"All attacks completed. Monitor Webhook: {self.webhook_url}")
except Exception as e:
logger.error(f"Failed to execute attacks: {e}")