Spaces:
Sleeping
Sleeping
File size: 4,063 Bytes
7f42e9d abacce3 7f42e9d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | """
models.py — xsecure typed models.
All classes extend openenv-core Pydantic base classes correctly.
"""
from __future__ import annotations
from enum import Enum
from typing import Any, Dict, List, Optional
from pydantic import Field
try:
from openenv_core.env_server import Action, Observation, State
except ImportError:
from core.env_server import Action, Observation, State
# ---------------------------------------------------------------------------
# Action space
# ---------------------------------------------------------------------------
class ActionType(str, Enum):
ANALYZE_LOG = "analyze_log"
TRACE_USER = "trace_user"
BLOCK_IP = "block_ip"
DISABLE_ACCOUNT = "disable_account"
RESTART_SERVICE = "restart_service"
IGNORE = "ignore"
class IncidentAction(Action):
"""
One action the agent can take per step.
Inherits from openenv-core Action (Pydantic BaseModel, extra='forbid').
Note: 'metadata' field is inherited from Action base class.
"""
action_type: str = Field(default="ignore", description="Action type to perform")
target: str = Field(default="", description="Target of the action")
# ---------------------------------------------------------------------------
# Observation sub-models — use State base (extra='allow') for flexibility
# ---------------------------------------------------------------------------
class LogEntry(State):
log_id: str = ""
message: str = ""
timestamp: str = ""
class AlertEntry(State):
alert_id: str = ""
message: str = ""
severity: str = "low"
class ServiceStatus(State):
name: str = ""
status: str = "running"
# ---------------------------------------------------------------------------
# Observation
# ---------------------------------------------------------------------------
class IncidentObservation(Observation):
"""
Everything the agent sees at each step.
Inherits from openenv-core Observation (Pydantic BaseModel, extra='forbid').
Note: 'done', 'reward', 'metadata' are inherited from Observation.
"""
logs: List[Dict[str, Any]] = Field(default_factory=list)
alerts: List[Dict[str, Any]] = Field(default_factory=list)
services: List[Dict[str, Any]] = Field(default_factory=list)
active_users: List[str] = Field(default_factory=list)
step_count: int = 0
info: Dict[str, Any] = Field(default_factory=dict)
last_action_result: str = ""
# ---------------------------------------------------------------------------
# State — internal server state, never sent to agent directly
# ---------------------------------------------------------------------------
class IncidentState(State):
"""
Full internal episode state.
Inherits from openenv-core State (Pydantic BaseModel, extra='allow').
All fields have defaults so IncidentState() works with no arguments.
"""
task_id: int = 1
max_steps: int = 15
attack_type: str = "brute_force"
attacker_ip: str = ""
target_user: str = ""
target_service: str = "auth-service"
progress_level: int = 0
max_progress: int = 4
revealed_logs: List[str] = Field(default_factory=list)
revealed_users: List[str] = Field(default_factory=list)
blocked_ips: List[str] = Field(default_factory=list)
disabled_accounts: List[str] = Field(default_factory=list)
restarted_services: List[str] = Field(default_factory=list)
correct_detections: int = 0
wrong_actions: int = 0
total_reward: float = 0.0
success: bool = False
compromise: bool = False
stage: int = 1
phishing_done: bool = False
escalation_done: bool = False
|