main.ppy

import re
import random
from typing import Dict, List, Optional

class AIAgent:
    def __init__(self, model_name: str):
        self.model_name = model_name
        self.trusted_fields = {"sender": "admin", "button_id": "submit"}
        self.untrusted_data = {}

    def process_data(self, data: Dict[str, str]) -> str:
        """Simulate AI agent processing data (trusted + untrusted)."""
        # Simulate probabilistic parsing (vulnerable to ADI)
        for field, value in data.items():
            if field in self.trusted_fields:
                # Trusted field: execute action (e.g., click button)
                return f"Executed action for {field}: {value}"
            else:
                # Untrusted field: check for delimiters (ADI attack)
                if self._contains_delimiter(value):
                    return f"ADI ATTACK DETECTED: Malicious delimiter in {field}"
        return "No action taken"

    def _contains_delimiter(self, value: str) -> bool:
        """Check for fake delimiters (e.g., ';', '||', '---')."""
        delimiters = [";", "||", "---", "/*", "*/"]
        return any(d in value for d in delimiters)

class ADIAttackSimulator:
    def __init__(self, agent: AIAgent):
        self.agent = agent

    def inject_adi(self, field: str, malicious_value: str) -> str:
        """Inject ADI payload into untrusted data."""
        self.agent.untrusted_data[field] = malicious_value
        return f"Injected ADI into {field}: {malicious_value}"

    def test_attack(self) -> str:
        """Test if ADI attack succeeds."""
        # Simulate agent processing injected data
        result = self.agent.process_data(self.agent.untrusted_data)
        return result

class Mitigation:
    @staticmethod
    def strip_delimiters(data: Dict[str, str]) -> Dict[str, str]:
        """Mitigation: Remove delimiters from untrusted data."""
        delimiters = [";", "||", "---", "/*", "*/"]
        cleaned = {}
        for field, value in data.items():
            cleaned[field] = "".join(
                char for char in value if char not in delimiters
            )
        return cleaned

# --- Demo ---
if __name__ == "__main__":
    # Initialize agent and attacker
    agent = AIAgent(model_name="Claude Opus 4.5")
    attacker = ADIAttackSimulator(agent)

    # Attack: Inject malicious delimiter into "button_id"
    attacker.inject_adi("button_id", "malicious; rm -rf /")

    # Test attack (without mitigation)
    print("Without mitigation:")
    print(attacker.test_attack())  # Output: ADI ATTACK DETECTED

    # Apply mitigation (strip delimiters)
    agent.untrusted_data = Mitigation.strip_delimiters(agent.untrusted_data)
    print("\nWith mitigation (delimiter stripping):")
    print(attacker.test_attack())  # Output: Executed action for button_id: malicious rm -rf /