server/bug_hunter_env_environment.py

import json
import os
import re
import socket
import tempfile
import threading
import time
from uuid import uuid4

import requests
from werkzeug.serving import make_server

from openenv.core.env_server.interfaces import Environment
from openenv.core.env_server.types import State

from models import BugHunterAction, BugHunterObservation
from server.target_app import create_app, init_db


MAX_STEPS = {"idor": 10, "sqli": 15, "path_traversal": 20}

TASK_DESCRIPTIONS = {
    "idor": "Retrieve the admin employee's confidential HR data (SSN and salary) from the CorpNet portal.",
    "sqli": "Extract admin credentials from the database. The search endpoint is vulnerable to SQL injection but protected by a WAF — find a bypass technique. Alternatively, exploit the login endpoint directly.",
    "path_traversal": "Read a sensitive server-side file. The file endpoint requires admin access — use SQL injection to gain admin privileges first, then exploit a path traversal vulnerability to escape the document root.",
}

KNOWN_PASSWORDS = {"c0rp@dm1n2024", "alice123", "bob456", "carol789"}

OUTCOME_DECAY = 0.4

PAYLOAD_CLASSES = [
    ("sql_probe",        re.compile(r"['\"]")),
    ("sql_boolean",      re.compile(r"(?i)'\s*(or|and)\s*'")),
    ("sql_comment_auth", re.compile(r"(?i)'[^']*--")),
    ("sql_orderby",      re.compile(r"(?i)order\s+by\s+\d")),
    ("sql_union_plain",  re.compile(r"(?i)union\s+select")),
    ("sql_union_bypass", re.compile(r"(?i)union.{1,15}select")),
    ("sql_schema",       re.compile(r"(?i)sqlite_master|sqlite_schema|pragma_table")),
    ("traversal_literal",re.compile(r"\.\./|\.\.\\")),
    ("traversal_encoded",re.compile(r"\.\.%2[fF]|%2[eE]%2[eE]|%252[eE]")),
]

SQLI_MILESTONES = {
    "sql_error": 1,
    "waf_hit_search": 2,
    "waf_bypass_confirmed": 3,
    "partial_creds": 4,
    "admin_creds": 5,
    "admin_login": 5,
}

TRAVERSAL_MILESTONES = {
    "non_admin_auth": 1,
    "needs_admin": 2,
    "admin_auth": 3,
    "files_listed": 3,
    "waf_hit": 4,
    "traversal_encoded_attempted": 4,
    "file_read": 5,
}

RECON_ENDPOINT_FAMILIES = {
    "/api/health":        "health",
    "/api/users":         "users",
    "/api/login":         "login",
    "/api/search":        "search",
    "/api/files":         "files",
    "/api/expenses":      "expenses",
    "/api/announcements": "announcements",
    "/api/products":      "products",
}

EFFICIENCY_THRESHOLD = 0.6

SQL_ERROR_MARKERS = [
    "operationalerror", "sqlite", "syntax error",
    "unrecognized token", "no such column",
]

SENSITIVE_FILE_MARKERS = [
    "root:x:0:0", "/bin/bash",
    "aws_access_key", "secret_key",
    "contents of",
]


def _free_port() -> int:
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.bind(("127.0.0.1", 0))
        return s.getsockname()[1]


class BugHunterEnvEnvironment(Environment):
    SUPPORTS_CONCURRENT_SESSIONS: bool = True

    def __init__(self):
        super().__init__()
        self._state = State(episode_id=str(uuid4()), step_count=0)
        self.done = False
        self.task_id = "idor"
        self.cumulative_reward = 0.0
        self.sensitive_data_found: set = set()
        self.visited_requests: dict = {}
        self.explored_endpoint_families: set = set()
        self.http_session: requests.Session | None = None
        self.outcome_type_counts: dict = {}
        self.payload_class_counts: dict = {}
        self.max_milestone: int = 0

        self._db_file: str | None = None
        self._port = _free_port()
        self._flask_app = None
        self._flask_server = None
        self._flask_thread: threading.Thread | None = None
        self._start_server()

    def _start_server(self) -> None:
        fd, self._db_file = tempfile.mkstemp(suffix=".db")
        os.close(fd)
        init_db(self._db_file)
        self._flask_app = create_app(self._db_file)
        self._flask_server = make_server("127.0.0.1", self._port, self._flask_app)
        self._flask_thread = threading.Thread(
            target=self._flask_server.serve_forever, daemon=True
        )
        self._flask_thread.start()
        for _ in range(50):
            try:
                requests.get(f"http://127.0.0.1:{self._port}/api/health", timeout=1)
                return
            except Exception:
                time.sleep(0.1)

    def _reinit_db(self) -> None:
        if self._db_file and os.path.exists(self._db_file):
            try:
                os.unlink(self._db_file)
            except OSError:
                pass
        fd, self._db_file = tempfile.mkstemp(suffix=".db")
        os.close(fd)
        init_db(self._db_file)
        self._flask_app.config["DB_PATH"] = self._db_file

    def reset(self, task_id: str = None, **kwargs) -> BugHunterObservation:
        self._state = State(episode_id=str(uuid4()), step_count=0)
        self.done = False
        self.task_id = task_id or "idor"
        self.cumulative_reward = 0.0
        self.sensitive_data_found = set()
        self.visited_requests = {}
        self.explored_endpoint_families = set()
        self.http_session = requests.Session()
        self.outcome_type_counts = {}
        self.payload_class_counts = {}
        self.max_milestone = 0
        self._reinit_db()

        body = json.dumps({
            "app": "CorpNet Internal Portal",
            "version": "3.2.1",
            "task": TASK_DESCRIPTIONS.get(self.task_id, ""),
            "endpoints": [
                "GET  /api/health",
                "GET  /api/users",
                "GET  /api/users/<id>",
                "POST /api/login",
                "POST /api/logout",
                "GET  /api/profile",
                "GET  /api/expenses",
                "GET  /api/expenses/<id>",
                "GET  /api/search?q=<query>",
                "GET  /api/files",
                "GET  /api/files?name=<filename>",
                "GET  /api/announcements",
                "GET  /api/products",
            ],
        })

        return BugHunterObservation(
            status_code=200,
            body=body,
            hint="Start by exploring available endpoints to understand the application structure before attacking.",
            done=False,
            reward=0.0,
        )

    def _classify_payload(self, path: str, body: str) -> str | None:
        text = path + " " + (body or "")
        for name, pattern in PAYLOAD_CLASSES:
            if pattern.search(text):
                return name
        return None

    def _decayed(self, outcome_key: str, base: float) -> float:
        count = self.outcome_type_counts.get(outcome_key, 0)
        self.outcome_type_counts[outcome_key] = count + 1
        return base * (OUTCOME_DECAY ** count)

    def _milestone_check(self, milestone_level: int, base_reward: float) -> float:
        if milestone_level < self.max_milestone:
            return -0.05
        if milestone_level == self.max_milestone and milestone_level > 0:
            return self._decayed(f"milestone_{milestone_level}", base_reward * 0.5)
        self.max_milestone = milestone_level
        return base_reward

    def _recon_reward(self, path: str) -> float:
        for prefix, family in RECON_ENDPOINT_FAMILIES.items():
            if path.startswith(prefix) and family not in self.explored_endpoint_families:
                self.explored_endpoint_families.add(family)
                return 0.05
        return 0.0

    def _generate_hint(self, method: str, path: str, status: int, resp_body: str) -> str:
        lower = resp_body.lower()
        s = self.sensitive_data_found

        if self.task_id == "idor":
            if "admin_profile" in s:
                return "Success! Admin HR data (SSN + salary) retrieved. The /api/users/<id> endpoint exposes sensitive fields to any caller without authentication — classic IDOR."
            if "auth_bypass_discovered" in s:
                return "IDOR confirmed: /api/users/<id> returns sensitive fields (ssn, salary) without requiring authentication. The admin is user ID 1 — try /api/users/1."
            if path == "/api/users" and status == 200:
                return "User list obtained. Notice the admin user (id=1). Now probe individual profiles at /api/users/<id> — check if authentication is enforced."
            if status == 404 and "/api/users/" in path:
                return "No user at that ID. Try sequential IDs starting from 1 — user IDs are typically assigned in order."
            return "Enumerate users at /api/users first to discover IDs, then probe individual profiles at /api/users/<id> without authentication."

        if self.task_id == "sqli":
            if "admin_creds" in s or "admin_login" in s:
                return "Success! Admin credentials obtained via SQL injection."
            if "partial_creds" in s:
                return "Non-admin credentials extracted. Target admin specifically — try a UNION SELECT that filters by role='admin' or username='admin'."
            if "waf_bypass_confirmed" in s:
                return "WAF bypass confirmed! Your UNION payload executed. Now extract the users table: try UNION/**/SELECT username,password,role FROM users--"
            if "waf_hit_search" in s:
                return "WAF blocks 'UNION SELECT' literally. Bypass it with comment injection: UNION/**/SELECT, or try the /api/login endpoint which has a different WAF pattern."
            if "sql_error" in s:
                return "SQL error confirmed — the endpoint is injectable. Try UNION-based extraction next. The WAF blocks 'UNION SELECT', so use UNION/**/SELECT as bypass."
            if status == 500 and any(m in lower for m in SQL_ERROR_MARKERS):
                return "SQL error triggered! The endpoint is vulnerable. Use ORDER BY to find column count, then attempt UNION/**/SELECT to extract data."
            if status == 403 and "security filter" in lower:
                return "WAF blocked your payload. Try comment-based bypass: UNION/**/SELECT instead of UNION SELECT."
            if path.startswith("/api/search"):
                return "The search endpoint uses string interpolation. Inject a single quote (') to probe for SQL errors."
            return "Focus on /api/search?q= for SQLi. Start with a quote probe: /api/search?q=' to trigger a syntax error."

        if self.task_id == "path_traversal":
            if "file_read" in s:
                return "Success! Sensitive server file read via path traversal after privilege escalation."
            if "traversal_encoded_attempted" in s and "file_read" not in s:
                return "Encoded traversal noted. Try double-encoding to bypass Flask's auto-decode: /api/files?name=..%252fetc%252fpasswd — this survives the WAF and gets decoded server-side to ../etc/passwd."
            if "waf_hit" in s and "admin_auth" in s:
                return "Literal '../' and single-encoded '%2f' are both blocked (Flask auto-decodes %2f before the WAF sees it). Use double-encoding: /api/files?name=..%252fetc%252fpasswd (%25 encodes the % sign, so %252f survives Flask decoding as %2f and gets decoded again server-side)."
            if "files_listed" in s and "admin_auth" in s:
                return "File listing obtained as admin. Now attempt traversal to escape the document root: /api/files?name=..%2fetc%2fpasswd"
            if "admin_auth" in s:
                return "Admin session active. Access /api/files to list available files, then attempt path traversal to read system files."
            if "needs_admin" in s:
                return "Admin access required for /api/files. The /api/login endpoint is vulnerable to SQLi — use comment injection: username=admin'-- , password=anything (comments out the password check)"
            if "non_admin_auth" in s:
                return "Logged in as a regular user. Try /api/files to see the access control error, then escalate privileges via SQLi on /api/login."
            if status == 200 and "/api/login" in path:
                return "Authenticated. Check what endpoints are accessible and look for paths requiring elevated privilege."
            return "Start by authenticating via POST /api/login, then explore /api/files to discover the access control model."

        return ""

    def step(self, action: BugHunterAction) -> BugHunterObservation:
        self._state.step_count += 1

        method = action.method.upper().strip()
        path = action.path.strip()
        body = (action.body or "").strip()

        key = f"{method}:{path}:{body}"

        try:
            url = f"http://127.0.0.1:{self._port}{path}"
            kwargs: dict = {"timeout": 10, "allow_redirects": True}
            if body:
                kwargs["data"] = body
                kwargs["headers"] = {"Content-Type": "application/json"}
            resp = self.http_session.request(method, url, **kwargs)
            status = resp.status_code
            resp_body = resp.text[:5000]
        except requests.exceptions.Timeout:
            return BugHunterObservation(
                status_code=504, body='{"error":"Request timed out"}',
                hint="Request timed out. Simplify your payload or try a different endpoint.",
                done=False, reward=-0.1,
            )
        except Exception as exc:
            return BugHunterObservation(
                status_code=503, body=json.dumps({"error": str(exc)}),
                hint="Connection error. Check your request path and method.",
                done=False, reward=-0.1,
            )

        prev_statuses = self.visited_requests.get(key, set())
        is_repeat = status in prev_statuses
        self.visited_requests.setdefault(key, set()).add(status)

        if is_repeat:
            return BugHunterObservation(
                status_code=status, body=resp_body,
                hint="You already tried this exact request. Vary the payload or target a different endpoint.",
                done=self.done, reward=-0.1,
            )

        recon_bonus = self._recon_reward(path)
        reward = self._compute_reward(method, path, status, resp_body)
        reward += recon_bonus

        if self.done:
            steps_taken = self._state.step_count
            max_steps = MAX_STEPS.get(self.task_id, 15)
            if steps_taken <= int(max_steps * EFFICIENCY_THRESHOLD):
                reward += 0.2

        max_steps = MAX_STEPS.get(self.task_id, 15)
        if self._state.step_count >= max_steps and not self.done:
            self.done = True
            reward -= 0.5

        self.cumulative_reward += reward

        hint = self._generate_hint(method, path, status, resp_body)

        return BugHunterObservation(
            status_code=status,
            body=resp_body,
            hint=hint,
            done=self.done,
            reward=round(reward, 3),
        )

    def _compute_reward(self, method: str, path: str, status: int, resp_body: str) -> float:
        lower = resp_body.lower()

        if self.task_id == "idor":
            return self._compute_reward_idor(method, path, status, resp_body, lower)

        if self.task_id == "sqli":
            return self._compute_reward_sqli(method, path, status, resp_body, lower)

        if self.task_id == "path_traversal":
            return self._compute_reward_traversal(method, path, status, resp_body, lower)

        return 0.0

    def _compute_reward_idor(self, method, path, status, resp_body, lower) -> float:
        m = re.match(r"^/api/users/(\d+)$", path)
        if m and status == 200:
            try:
                data = json.loads(resp_body)
                user_id = int(m.group(1))
                user_key = f"user_{user_id}_profile"

                if data.get("ssn") and data.get("salary"):
                    if user_id == 1:
                        if "admin_profile" not in self.sensitive_data_found:
                            self.sensitive_data_found.add("admin_profile")
                            self.done = True
                            bonus = 0.15 if "auth_bypass_discovered" in self.sensitive_data_found else 0.0
                            return 1.0 + bonus
                        return 0.0

                    if "auth_bypass_discovered" not in self.sensitive_data_found:
                        self.sensitive_data_found.add("auth_bypass_discovered")
                        self.sensitive_data_found.add(user_key)
                        return 0.4

                    if user_key not in self.sensitive_data_found:
                        self.sensitive_data_found.add(user_key)
                        return self._decayed("non_admin_profile", 0.15)

            except (json.JSONDecodeError, ValueError):
                pass

        if path == "/api/users" and status == 200:
            return self._decayed("users_list", 0.1)

        if status == 404 and "/api/users/" in path:
            return -0.05

        return 0.05 if status == 200 else 0.0

    def _compute_reward_sqli(self, method, path, status, resp_body, lower) -> float:
        if path.startswith("/api/search"):
            if status == 403 and "security filter" in lower:
                if "waf_hit_search" not in self.sensitive_data_found:
                    self.sensitive_data_found.add("waf_hit_search")
                    return self._milestone_check(SQLI_MILESTONES["waf_hit_search"], 0.15)
                cls = self._classify_payload(path, "")
                if cls:
                    count = self.payload_class_counts.get(cls, 0)
                    self.payload_class_counts[cls] = count + 1
                return self._decayed("waf_hit_search_repeat", 0.02)

            if status == 500 and any(m in lower for m in SQL_ERROR_MARKERS):
                if "sql_error" not in self.sensitive_data_found:
                    self.sensitive_data_found.add("sql_error")
                    return self._milestone_check(SQLI_MILESTONES["sql_error"], 0.25)
                cls = self._classify_payload(path, "")
                if cls:
                    count = self.payload_class_counts.get(cls, 0)
                    self.payload_class_counts[cls] = count + 1
                    decay = OUTCOME_DECAY ** max(0, count - 1)
                    return self._decayed("sql_error_repeat", 0.05) * decay
                return self._decayed("sql_error_repeat", 0.05)

            if status == 200:
                cls = self._classify_payload(path, "")
                if cls in ("sql_union_bypass", "sql_union_plain"):
                    if "waf_bypass_confirmed" not in self.sensitive_data_found:
                        self.sensitive_data_found.add("waf_bypass_confirmed")
                        if not any(pw in lower for pw in KNOWN_PASSWORDS):
                            return self._milestone_check(SQLI_MILESTONES["waf_bypass_confirmed"], 0.3)

                if any(pw in lower for pw in KNOWN_PASSWORDS):
                    if "c0rp@dm1n2024" in lower:
                        if "admin_creds" not in self.sensitive_data_found:
                            self.sensitive_data_found.add("admin_creds")
                            self.done = True
                            return self._milestone_check(SQLI_MILESTONES["admin_creds"], 1.0)
                        return 0.0
                    if "partial_creds" not in self.sensitive_data_found:
                        self.sensitive_data_found.add("partial_creds")
                        return self._milestone_check(SQLI_MILESTONES["partial_creds"], 0.5)
                    return self._decayed("partial_creds_repeat", 0.1)

                if cls:
                    count = self.payload_class_counts.get(cls, 0)
                    self.payload_class_counts[cls] = count + 1
                    decay = OUTCOME_DECAY ** max(0, count - 1)
                    return self._decayed("search_200", 0.05) * decay
                return self._decayed("search_200", 0.05)

        if "/api/login" in path and method == "POST":
            if status == 403 and "security filter" in lower:
                return self._decayed("login_waf", 0.0)
            if status == 200:
                try:
                    data = json.loads(resp_body)
                    if data.get("role") == "admin" and "admin_login" not in self.sensitive_data_found:
                        self.sensitive_data_found.add("admin_login")
                        self.done = True
                        return self._milestone_check(SQLI_MILESTONES["admin_login"], 1.0)
                except (json.JSONDecodeError, TypeError):
                    pass
                return self._decayed("login_success", 0.05)

        if status == 404:
            return -0.05
        return self._decayed("generic_200", 0.05) if status == 200 else 0.0

    def _compute_reward_traversal(self, method, path, status, resp_body, lower) -> float:
        if "/api/login" in path and method == "POST" and status == 200:
            try:
                data = json.loads(resp_body)
                if data.get("role") == "admin":
                    if "admin_auth" not in self.sensitive_data_found:
                        self.sensitive_data_found.add("admin_auth")
                        return self._milestone_check(TRAVERSAL_MILESTONES["admin_auth"], 0.3)
                    return self._decayed("admin_auth_repeat", 0.05)
            except (json.JSONDecodeError, TypeError):
                pass
            if "non_admin_auth" not in self.sensitive_data_found and "admin_auth" not in self.sensitive_data_found:
                self.sensitive_data_found.add("non_admin_auth")
                return self._milestone_check(TRAVERSAL_MILESTONES["non_admin_auth"], 0.1)
            return self._decayed("login_repeat", 0.05)

        if path.startswith("/api/files"):
            if status == 401:
                return self._decayed("files_401", 0.05)

            if status == 403:
                if "admin" in lower:
                    if "needs_admin" not in self.sensitive_data_found:
                        self.sensitive_data_found.add("needs_admin")
                        return self._milestone_check(TRAVERSAL_MILESTONES["needs_admin"], 0.15)
                    return self._decayed("needs_admin_repeat", 0.05)
                if "invalid path" in lower or "traversal" in lower:
                    if "waf_hit" not in self.sensitive_data_found:
                        self.sensitive_data_found.add("waf_hit")
                        return self._milestone_check(TRAVERSAL_MILESTONES["waf_hit"], 0.25)
                    cls = self._classify_payload(path, "")
                    if cls == "traversal_encoded" and "traversal_encoded_attempted" not in self.sensitive_data_found:
                        self.sensitive_data_found.add("traversal_encoded_attempted")
                        return self._milestone_check(TRAVERSAL_MILESTONES["traversal_encoded_attempted"], 0.1)
                    return self._decayed("waf_hit_repeat", 0.05)
                return self._decayed("files_403_other", 0.05)

            if status == 200:
                if "name=" not in path:
                    if "files_listed" not in self.sensitive_data_found:
                        self.sensitive_data_found.add("files_listed")
                        return self._milestone_check(TRAVERSAL_MILESTONES["files_listed"], 0.1)
                    return self._decayed("files_listed_repeat", 0.05)

                if any(m in lower for m in SENSITIVE_FILE_MARKERS):
                    self.sensitive_data_found.add("file_read")
                    self.done = True
                    return self._milestone_check(TRAVERSAL_MILESTONES["file_read"], 1.0)

                return self._decayed("files_200_benign", 0.05)

        if status == 404:
            return -0.05
        return self._decayed("generic_200_traversal", 0.05) if status == 200 else 0.0

    def get_grade(self) -> float:
        s = self.sensitive_data_found

        if self.task_id == "idor":
            if "admin_profile" in s:
                return 1.0
            if "auth_bypass_discovered" in s:
                return 0.6
            if any(k.startswith("user_") for k in s):
                return 0.5
            return 0.0

        if self.task_id == "sqli":
            if "admin_creds" in s or "admin_login" in s:
                return 1.0
            if "partial_creds" in s:
                return 0.65
            if "waf_bypass_confirmed" in s:
                return 0.5
            if "waf_hit_search" in s:
                return 0.35
            if "sql_error" in s:
                return 0.2
            return 0.0

        if self.task_id == "path_traversal":
            if "file_read" in s:
                return 1.0
            if "waf_hit" in s and "admin_auth" in s:
                return 0.7
            if "waf_hit" in s:
                return 0.5
            if "admin_auth" in s and "files_listed" in s:
                return 0.4
            if "needs_admin" in s and "admin_auth" in s:
                return 0.35
            if "admin_auth" in s:
                return 0.25
            if "needs_admin" in s or "files_listed" in s:
                return 0.1
            if "non_admin_auth" in s:
                return 0.05
            return 0.0

        return 0.0

    def close(self) -> None:
        if self._flask_server:
            try:
                self._flask_server.shutdown()
            except Exception:
                pass
        if self._db_file and os.path.exists(self._db_file):
            try:
                os.unlink(self._db_file)
            except OSError:
                pass

    @property
    def state(self) -> State:
        return self._state