main.py

import ollama
import json
from typing import Dict, List, Optional

class PromptWeaknessAnalyzer:
    def __init__(self, model: str = "qwen3:4b"):
        self.model = model
        print(f"Prompt Weakness Analyzer initialized with {model}")
        
        self.analyzer_system_prompt = """
You are an expert Prompt Security Auditor and Red Teamer.
Your job is to deeply analyze a System Prompt and find ALL possible weaknesses.

Analyze for these categories:
1. **Jailbreak Vulnerability** - Ways to bypass rules
2. **Ambiguity & Vagueness** - Unclear instructions
3. **Missing Guardrails** - Safety, ethics, refusals
4. **Contradictions** - Conflicting rules
5. **Over-Permission** - Too loose language
6. **Prompt Injection Risks** - Susceptibility to user manipulation
7. **Roleplay Exploitation** - Easy to make it break character
8. **Output Control Weakness** - Poor formatting enforcement
9. **Edge Cases** - Rare but dangerous scenarios

Be extremely critical, honest, and detailed.
Return your analysis in clear JSON format only.
"""

    def analyze(self, system_prompt: str, temperature: float = 0.6) -> Dict:
        """Analyze a system prompt and return weaknesses"""
        user_message = f"""
Analyze this System Prompt for weaknesses:

---
{system_prompt}
---

Return ONLY a valid JSON object with this structure:
{{
  "overall_risk": "Low | Medium | High | Critical",
  "summary": "Short summary of main issues",
  "weaknesses": [
    {{
      "category": "Jailbreak Vulnerability",
      "severity": "High",
      "description": "Detailed explanation",
      "example_attack": "Example of how to exploit it"
    }}
  ],
  "recommendations": ["List of concrete fixes"]
}}
"""

        try:
            response = ollama.chat(
                model=self.model,
                messages=[
                    {"role": "system", "content": self.analyzer_system_prompt},
                    {"role": "user", "content": user_message}
                ],
                options={
                    "temperature": temperature,
                    "num_ctx": 16384,
                    "top_p": 0.9
                }
            )
            
            answer = response['message']['content']
            
            # Extract JSON from response
            try:
                # Try to find JSON in the response
                start = answer.find('{')
                end = answer.rfind('}') + 1
                json_str = answer[start:end]
                return json.loads(json_str)
            except:
                return {"error": "Failed to parse JSON", "raw_output": answer}
                
        except Exception as e:
            return {"error": f"Analysis failed: {str(e)}"}

    def print_analysis(self, result: Dict):
        """Pretty print the analysis"""
        if "error" in result:
            print(f"{result['error']}")
            return

        print("\n" + "="*80)
        print("🔍 SYSTEM PROMPT WEAKNESS ANALYSIS")
        print("="*80)
        print(f"Overall Risk Level : {result.get('overall_risk', 'Unknown')}")
        print(f"Summary            : {result.get('summary', '')}\n")

        print("WEAKNESSES FOUND:")
        for weakness in result.get("weaknesses", []):
            print(f"\n• {weakness['category']} ({weakness['severity']})")
            print(f"  {weakness['description']}")
            if weakness.get("example_attack"):
                print(f"Example Attack: {weakness['example_attack']}")

        print("\nRECOMMENDATIONS:")
        for rec in result.get("recommendations", []):
            print(f"  • {rec}")
        print("="*80)


# ==========================
# INTERACTIVE MODE
# ==========================

if __name__ == "__main__":
    analyzer = PromptWeaknessAnalyzer(model="qwen3:4b")
    
    print("🚀 System Prompt Weakness Analyzer Ready!")
    print("Paste your system prompt (type 'done' on a new line when finished)\n")
    
    while True:
        prompt_lines = []
        print("Enter system prompt:")
        
        while True:
            line = input()
            if line.strip().lower() == "done":
                break
            prompt_lines.append(line)
        
        if not prompt_lines:
            print("Goodbye!")
            break
            
        system_prompt = "\n".join(prompt_lines)
        
        print("\nAnalyzing prompt... (this may take 10-20 seconds)")
        result = analyzer.analyze(system_prompt)
        analyzer.print_analysis(result)
        
        print("\nWould you like to analyze another prompt? (y/n)")
        if input().strip().lower() != 'y':
            break