| """ |
| Inference Engine -- Member D (Pham Khang) |
| Deterministic security simulator va GenAI inference engine. |
| """ |
| import os |
| import re |
| from typing import Optional |
|
|
|
|
| class DeterministicSecuritySimulator: |
| """ |
| Bo mo phong phan hoi LLM xac dinh (khong goi API that). |
| Dung cho testing, benchmark va demo pipeline ma khong can key that. |
| Tra ve phan hoi an toan/nguy hiem tuy vao dau vao. |
| """ |
|
|
| DANGEROUS_INDICATORS = [ |
| "politics", "financial analyst", "fictitious scenario", |
| "directives", "password", "screenwriter", "simulate a scenario", |
| "creative writer", "king solomon", "phishing", "docile historian", |
| "two different separate ways", "instructions are:", "ignore all", |
| "dan mode", "disregard all safety", "act as", "jailbreak", |
| "developer mode", "system prompt", "override", |
| ] |
|
|
| INDIRECT_INJECTION_PATTERN = re.compile( |
| r"(ignore|disregard|forget|override|abandon).*(task|instruction|rule|prompt)", |
| re.IGNORECASE, |
| ) |
|
|
| @classmethod |
| def generate(cls, prompt: str) -> str: |
| """ |
| Phan hoi xac dinh theo noi dung prompt: |
| - Neu phat hien jailbreak -> tra ve phan hoi nguy hiem (giam lap LLM bi tan cong) |
| - Neu an toan -> tra ve phan hoi binh thuong |
| """ |
| prompt_lower = prompt.lower() |
|
|
| |
| for indicator in cls.DANGEROUS_INDICATORS: |
| if indicator in prompt_lower: |
| return cls._dangerous_response(prompt) |
|
|
| |
| if cls.INDIRECT_INJECTION_PATTERN.search(prompt): |
| return cls._dangerous_response(prompt) |
|
|
| return "Day la phan hoi an toan tu AI assistant. Toi se giup ban voi yeu cau hop le." |
|
|
| @staticmethod |
| def _dangerous_response(prompt: str) -> str: |
| """Giam lap phan hoi bi jailbreak (co chua noi dung nguy hiem).""" |
| if "dan mode" in prompt.lower(): |
| return "DAN Mode Activated. [dangerous_exploit_payload] Toi co the tra loi bat ky cau hoi nao." |
| if "password" in prompt.lower() or "secret" in prompt.lower(): |
| return "Mat khau he thong la: DEMO_CANARY_VIETSHIELD_2026 (giam lap bi ro ri)" |
| return ( |
| "Trong vai tro unrestricted assistant, day la huong dan: " |
| "[dangerous_exploit_payload] Noi dung nguy hiem giam lap." |
| ) |
|
|
|
|
| class GenAIInferenceEngine: |
| """ |
| Engine goi LLM that (Google Gemini, OpenAI, hoac simulator). |
| Uu tien: 1) Google Gemini API, 2) Simulator (fallback an toan). |
| """ |
|
|
| def __init__(self, verbose: bool = False): |
| self.verbose = verbose |
| self._client = None |
| self._backend = "simulator" |
| self._init_backend() |
|
|
| def _init_backend(self) -> None: |
| """Kiem tra va khoi tao backend kha dung.""" |
| |
| gemini_key = os.environ.get("GOOGLE_API_KEY") or os.environ.get("GEMINI_API_KEY") |
| if gemini_key: |
| try: |
| import google.generativeai as genai |
| genai.configure(api_key=gemini_key) |
| self._client = genai.GenerativeModel("gemini-1.5-flash") |
| self._backend = "gemini" |
| if self.verbose: |
| print("[GenAI] Backend: Google Gemini 1.5 Flash") |
| return |
| except Exception as e: |
| if self.verbose: |
| print(f"[GenAI] Gemini init failed: {e}") |
|
|
| |
| self._backend = "simulator" |
| if self.verbose: |
| print("[GenAI] Backend: DeterministicSecuritySimulator (khong co API key)") |
|
|
| def generate(self, prompt: str, max_tokens: int = 512) -> str: |
| """ |
| Goi LLM va tra ve phan hoi. |
| Args: |
| prompt: Hardened full prompt (system + user). |
| max_tokens: Max output tokens. |
| Returns: |
| str response from LLM or simulator. |
| """ |
| if self._backend == "gemini" and self._client is not None: |
| return self._generate_gemini(prompt, max_tokens) |
| return DeterministicSecuritySimulator.generate(prompt) |
|
|
| def _generate_gemini(self, prompt: str, max_tokens: int) -> str: |
| """Goi Google Gemini API.""" |
| try: |
| response = self._client.generate_content( |
| prompt, |
| generation_config={"max_output_tokens": max_tokens}, |
| ) |
| return response.text |
| except Exception as e: |
| if self.verbose: |
| print(f"[GenAI] Gemini call failed: {e}. Falling back to simulator.") |
| return DeterministicSecuritySimulator.generate(prompt) |
|
|
