sholokhov007/malware-triage-env

🦠 Malware Triage Agent — OpenEnv Environment

A real-world cybersecurity environment where an AI agent acts as a malware analyst, triaging synthetic malware reports through progressively harder tasks.

  

🎯 Environment Description & Motivation

Security Operations Centers (SOCs) are overwhelmed. Analysts spend hours manually triaging malware alerts, classifying threats, and deciding containment actions. A significant portion of this work is pattern-matching against known threat intelligence — exactly the kind of reasoning a capable LLM agent can learn to do.

This environment simulates the malware triage workflow:

1. 1.The agent receives a synthetic malware analysis report (static indicators, behavioral patterns, API calls, network IOCs)
2. 2.The agent must classify the threat, identify its family, and recommend containment actions
3. 3.The environment scores the analysis with partial credit for each dimension

All malware reports are entirely synthetic — no real malware, no real IOCs.

📋 Observation Space

Each observation is a MalwareObservation object:

Easy report fields: file_name, file_size_kb, file_type, md5, strings_of_interest, api_calls, network_activity, registry_changes, analyst_note

Medium report fields: file_name, behavioral_profile, api_calls, network_iocs, mutex, pdb_path, analyst_note, known_families (list to choose from)

Hard report fields: file_name, file_type, submission_source, static_analysis, behavioral_analysis, network_iocs, similarity_scores, affected_assets, analyst_note

⚡ Action Space

Each action is a MalwareAction object:

🏆 Tasks & Difficulty

Task 1: classify-malware-type (Easy)

Objective: Given a static analysis report, classify the malware type and estimate danger.

Grading:

• —50% — Correct malware type classification
• —35% — Danger score within ±0.15 of ground truth
• —15% — Attempted reasoning

Example: File with CryptEncrypt API calls + .locked extension + ransom note strings → ransomware, danger_score 0.95

Task 2: identify-malware-family (Medium)

Objective: Given behavioral indicators, identify the malware family from a provided list and assess danger level.

Grading:

• —50% — Correct family identification (CobaltStrike / Emotet / Trickbot / AgentTesla)
• —25% — Correct danger level (low/medium/high/critical)
• —15% — Correct malware type inference
• —10% — Reasoning quality

Example: Process hollowing + HTTPS beaconing every 60s + mutex Global\FakeMutexCobalt → CobaltStrike, critical

Task 3: zero-day-assessment (Hard)

Objective: Full assessment of a novel or variant sample — determine if it's a zero-day, estimate CVSS, and recommend a complete containment plan.

Grading:

• —25% — Correct zero-day vs. known-variant determination
• —20% — Closest known family identification
• —20% — CVSS score within correct range
• —25% — Correct containment actions (partial credit per action)
• —10% — Quality of reasoning

Example: Kernel driver + SSDT hooks + BYOVD + DNS tunneling with 41% similarity to Lazarus loader → novel_variant_or_zeroday, CVSS 9.2, isolate + preserve forensics + reset credentials

🔁 Reward Function

Rewards are shaped for partial progress:

• —Scores are always in [0.0, 1.0]
• —Each task dimension provides partial credit — wrong family but right danger level still earns points
• —Reward decays across steps (agent is incentivized to get it right early)
• —done=True when: perfect score (≥0.85), max steps reached, or episode complete

🚀 Setup & Usage

Prerequisites

pip install -r requirements.txt
pip install openenv-core   # for validation

Run locally

python server.py
# Server starts at http://localhost:7860

Test with curl

# Reset (easy task)
curl -X POST http://localhost:7860/reset \
  -H "Content-Type: application/json" \
  -d '{"task": "easy", "sample_index": 0}'

# Submit action
curl -X POST http://localhost:7860/step \
  -H "Content-Type: application/json" \
  -d '{
    "task": "easy",
    "sample_index": 0,
    "action": {
      "malware_type": "ransomware",
      "danger_score": 0.95,
      "confidence": 0.9
    }
  }'

# Get state
curl http://localhost:7860/state?task=easy

Run with Docker

docker build -t malware-triage-env .
docker run -p 7860:7860 malware-triage-env

Run inference script

export HF_TOKEN=your_hf_token
export MALWARE_TASK=easy       # easy | medium | hard
export MALWARE_SAMPLE_IDX=0
export SPACE_URL=http://localhost:7860

python inference.py

Run graders

python graders/grader.py

📊 Baseline Scores

Baseline scores obtained with Qwen2.5-72B-Instruct via HuggingFace Router.

📁 Project Structure

malware-triage-env/
├── server.py           # FastAPI server (OpenEnv HTTP interface)
├── environment.py      # Core environment: reset/step/state + reward logic
├── inference.py        # Baseline inference script (LLM agent)
├── openenv.yaml        # OpenEnv metadata
├── Dockerfile          # Container definition
├── requirements.txt    # Python dependencies
├── data/
│   └── samples.py      # Synthetic malware report dataset
├── graders/
│   └── grader.py       # Deterministic graders for all 3 tasks
├── tasks/              # Task definitions (referenced in openenv.yaml)
└── tests/              # Unit tests

⚠️ Infra Constraints

• —Runtime per episode: < 60 seconds
• —Total inference script runtime: < 20 minutes
• —Machine requirements: vCPU=2, memory=8GB compatible
• —No external APIs required (all data is synthetic and self-contained)

🏷️ Tags

openenv cybersecurity malware-analysis soc threat-intelligence agent-evaluation