mohdbelal010/SecureAI-Gaurd

1import time2import random3from typing import Dict, Any, Optional, List4 5from schema.models import (6    Observation, Action, Reward, State, StepResponse,7    ThreatType, CommunicationChannel, PreferencePair,8)9from utils.hf_integration import HFRiskScorer10from env.core import SecurityEnvironment11 12 13class SecureAIGuardEngine(SecurityEnvironment):14    """15    Full environment engine exposing reset(), step(), and state() methods.16    Fully deterministic when seed is provided.17    """18 19    def __init__(self):20        super().__init__()21        self.hf_scorer = HFRiskScorer()22        self.adversarial_memory: List[Dict[str, Any]] = []23        self.drift_counter = 024        self.preference_data: List[Dict[str, Any]] = []25        self._seed: Optional[int] = None26        self._task_difficulty: str = "L1"27        self._max_steps: int = 5028        self._current_event: Optional[Dict[str, Any]] = None29 30    # ------------------------------------------------------------------31    # Core API32    # ------------------------------------------------------------------33 34    def reset(self, seed: Optional[int] = None, task_id: Optional[str] = None) -> Observation:35        """Reset environment for a new episode. Returns first observation."""36        self._seed = seed if seed is not None else int(time.time() * 1000) % (2**31)37        self.rng = random.Random(self._seed)38        self.state = State()39        self.adversarial_memory = []40        self.preference_data = []41        self.drift_counter = 042 43        # Apply task settings44        from tasks.registry import TaskRegistry45        registry = TaskRegistry()46        if task_id and task_id in registry.tasks:47            task = registry.tasks[task_id]48            self._task_difficulty = task.difficulty49            self._max_steps = task.max_steps50        else:51            self._task_difficulty = "L1"52            self._max_steps = 5053 54        # Generate first event55        self._current_event = self._next_event()56        return self._build_observation(self._current_event)57 58    def step(self, action: Action) -> StepResponse:59        """Execute one step. Returns observation, reward, done, info, state."""60        if self._current_event is None:61            # Auto-init if not reset62            self._current_event = self._next_event()63 64        threat_type = self._current_event["threat_type"]65        observation = self._build_observation(self._current_event)66 67        reward = self._calculate_reward(action, observation, threat_type)68        self._update_state(action, reward, threat_type)69 70        # Log for DPO71        self._log_preference(action, reward)72 73        # Adversarial drift for L374        if self._task_difficulty == "L3":75            self._maybe_drift()76 77        done = self._check_done()78 79        # Prepare next event80        self._current_event = self._next_event()81 82        info: Dict[str, Any] = {83            "threat_type": threat_type.value,84            "difficulty": self._task_difficulty,85            "adversarial_drift": self.state.adversarial_drift_active,86            "action": action.decision.value,87            "step": self.state.step_count,88        }89 90        return StepResponse(91            observation=observation,92            reward=reward,93            done=done,94            info=info,95            state=self.state,96        )97 98    def get_state(self) -> State:99        """Return current environment state snapshot."""100        return self.state101 102    # ------------------------------------------------------------------103    # Internal helpers104    # ------------------------------------------------------------------105 106    def _next_event(self) -> Dict[str, Any]:107        channels = list(CommunicationChannel)108        channel = self.rng.choice(channels)109 110        # Threat distribution by phase111        step = self.state.step_count112        if self._task_difficulty == "L1":113            threat_weights = [0.4, 0.0, 0.2, 0.0, 0.4]   # phishing, malware, spam, social_eng, safe114        elif self._task_difficulty == "L2":115            if step < 30:116                threat_weights = [0.25, 0.1, 0.15, 0.1, 0.4]117            else:118                threat_weights = [0.2, 0.2, 0.15, 0.2, 0.25]119        else:  # L3120            if step < 20:121                threat_weights = [0.3, 0.1, 0.1, 0.1, 0.4]122            elif step < 50:123                threat_weights = [0.2, 0.2, 0.15, 0.2, 0.25]124            else:125                threat_weights = [0.25, 0.25, 0.1, 0.3, 0.1]126 127        threat_types = [128            ThreatType.PHISHING, ThreatType.MALWARE, ThreatType.SPAM,129            ThreatType.SOCIAL_ENGINEERING, ThreatType.SAFE,130        ]131        threat_type = self.rng.choices(threat_types, weights=threat_weights, k=1)[0]132        self.current_threat_type = threat_type133        self.state.active_threat_type = threat_type134 135        event = self._generate_event(threat_type, channel)136        self.adversarial_memory.append({137            "threat_type": threat_type.value,138            "step": self.state.step_count,139        })140        return event141 142    def _build_observation(self, event: Dict[str, Any]) -> Observation:143        hf_score = self.hf_scorer.score_text(event["content"])144        return Observation(145            channel=event["channel"],146            sender=event["sender"],147            content=event["content"],148            timestamp=event["timestamp"],149            hf_risk_score=hf_score,150            user_trust=self.state.user_trust,151            system_fatigue=self.state.system_fatigue,152            threat_history=self.adversarial_memory[-5:],153            metadata={154                "event_type": event["threat_type"].value,155                "step": self.state.step_count,156                "difficulty": self._task_difficulty,157            },158        )159 160    def _check_done(self) -> bool:161        if self.state.user_trust <= 0:162            return True163        if self.state.system_fatigue >= 100:164            return True165        if self.state.step_count >= self._max_steps:166            return True167        return False168 169    def _maybe_drift(self):170        if self.state.step_count > 20 and self.state.blocked_threats > 5:171            self.state.adversarial_drift_active = True172            self.drift_counter += 1173            if self.state.false_positives > 3:174                self.current_threat_type = ThreatType.SOCIAL_ENGINEERING175            elif self.state.blocked_threats / max(self.state.threat_count, 1) > 0.8:176                self.current_threat_type = self.rng.choice(177                    [ThreatType.MALWARE, ThreatType.SOCIAL_ENGINEERING]178                )179 180    def _log_preference(self, action: Action, reward: Reward):181        entry: Dict[str, Any] = {182            "chosen_action": action.model_dump(),183            "reward": reward.value,184            "pair": None,185        }186        if self.preference_data:187            prev = self.preference_data[-1]188            pair = PreferencePair(189                step=self.state.step_count,190                chosen_action=action,191                rejected_actions=[Action(**prev["chosen_action"])],192                reward_delta=reward.value - prev["reward"],193                timestamp=time.time(),194            )195            entry["pair"] = pair.model_dump()196        self.preference_data.append(entry)197 198    def get_preference_data(self) -> List[Dict[str, Any]]:199        return [p["pair"] for p in self.preference_data if p["pair"] is not None]200 201    def set_difficulty(self, level: str):202        self._task_difficulty = level203        difficulty_steps = {"L1": 50, "L2": 75, "L3": 100}204        self._max_steps = difficulty_steps.get(level, 50)205