pretext-sim 0.3.2__py3-none-any.whl

pretext/client.py

@@ -0,0 +1,220 @@
+"""HTTP client for any OpenAI-compatible chat completions endpoint."""
+
+from __future__ import annotations
+
+import json
+import random
+import time
+from typing import Any
+
+import httpx
+from pydantic import ValidationError
+
+from .models import APIError, Settings, Simulation, _format_validation
+from .prompt import build_payload_messages
+
+RETRY_STATUS = {408, 409, 425, 429, 500, 502, 503, 504}
+MAX_RETRY_AFTER = 30.0
+USER_AGENT = "pretext/0.1 (+authorized-training)"
+
+
+def extract_json(content: str) -> dict[str, Any]:
+    """Pull a JSON object out of model output that may be wrapped or fenced."""
+
+    text = content.strip()
+    if text.startswith("```"):
+        text = text.split("\n", 1)[1] if "\n" in text else ""
+        if text.rstrip().endswith("```"):
+            text = text.rstrip()[:-3]
+        text = text.strip()
+
+    try:
+        parsed = json.loads(text)
+    except json.JSONDecodeError:
+        start = text.find("{")
+        end = text.rfind("}")
+        if start == -1 or end <= start:
+            raise APIError(
+                f"model returned no JSON object; first 200 chars: {content.strip()[:200]!r}"
+            ) from None
+        try:
+            parsed = json.loads(text[start : end + 1])
+        except json.JSONDecodeError as exc:
+            raise APIError(f"model returned malformed JSON: {exc}") from None
+
+    if not isinstance(parsed, dict):
+        raise APIError(f"expected a JSON object, got {type(parsed).__name__}")
+    return parsed
+
+
+def _content_of(body: dict[str, Any]) -> str:
+    try:
+        choice = body["choices"][0]
+        message = choice["message"]
+        content = message["content"]
+    except (KeyError, IndexError, TypeError) as exc:
+        raise APIError(f"unexpected response shape: {str(body)[:200]}") from exc
+
+    if isinstance(content, list):
+        pieces = [
+            part.get("text", "")
+            for part in content
+            if isinstance(part, dict) and part.get("type") in (None, "text")
+        ]
+        content = "".join(pieces)
+
+    if not isinstance(content, str) or not content.strip():
+        reason = choice.get("finish_reason") if isinstance(choice, dict) else None
+        if reason == "length":
+            raise APIError(
+                "model returned no content, only reasoning (finish_reason=length); "
+                "raise --max-tokens"
+            )
+        raise APIError("model returned empty content")
+    return content
+
+
+class Client:
+    def __init__(self, settings: Settings, *, transport: httpx.BaseTransport | None = None) -> None:
+        self.settings = settings
+        self._transport = transport
+        self.last_usage: dict[str, Any] | None = None
+
+    def _payload(self, profile, channel: str | None, extra: str | None) -> dict[str, Any]:
+        payload: dict[str, Any] = {
+            "model": self.settings.model,
+            "messages": build_payload_messages(profile, channel=channel, extra=extra),
+            "temperature": self.settings.temperature,
+            "max_tokens": self.settings.max_tokens,
+            "response_format": {"type": "json_object"},
+        }
+        if self.settings.seed is not None:
+            payload["seed"] = self.settings.seed
+        return payload
+
+    def generate(
+        self,
+        profile,
+        *,
+        channel: str | None = None,
+        extra: str | None = None,
+    ) -> Simulation:
+        body = self._request(self._payload(profile, channel, extra))
+        raw = extract_json(_content_of(body))
+        try:
+            return Simulation.model_validate(raw)
+        except ValidationError as exc:
+            raise APIError(
+                _format_validation(exc, prefix="model output failed validation")
+            ) from exc
+
+    def _request(self, payload: dict[str, Any]) -> dict[str, Any]:
+        headers = {
+            "Authorization": f"Bearer {self.settings.api_key}",
+            "Content-Type": "application/json",
+            "User-Agent": USER_AGENT,
+        }
+        attempts = self.settings.retries + 1
+        last: APIError | None = None
+
+        with httpx.Client(
+            timeout=self.settings.timeout,
+            transport=self._transport,
+            follow_redirects=True,
+        ) as client:
+            for attempt in range(1, attempts + 1):
+                try:
+                    response = client.post(self.settings.endpoint, json=payload, headers=headers)
+                except httpx.TimeoutException as exc:
+                    last = APIError(f"request timed out after {self.settings.timeout:g}s")
+                    if attempt < attempts:
+                        time.sleep(_backoff(attempt, None))
+                        continue
+                    raise last from exc
+                except httpx.HTTPError as exc:
+                    last = APIError(f"connection failed: {exc}")
+                    if attempt < attempts:
+                        time.sleep(_backoff(attempt, None))
+                        continue
+                    raise last from exc
+
+                if response.status_code in RETRY_STATUS and attempt < attempts:
+                    time.sleep(_backoff(attempt, response.headers.get("retry-after")))
+                    last = APIError(f"HTTP {response.status_code} from endpoint")
+                    continue
+
+                last = None
+                body = _decode(response)
+                usage = body.get("usage")
+                self.last_usage = usage if isinstance(usage, dict) else None
+                return body
+        raise last or APIError("request failed")
+
+    def list_models(self) -> list[str]:
+        """Fetch model ids from the endpoint; empty when unsupported."""
+
+        headers = {"Authorization": f"Bearer {self.settings.api_key}", "User-Agent": USER_AGENT}
+        with httpx.Client(timeout=self.settings.timeout, transport=self._transport) as client:
+            try:
+                base = self.settings.endpoint.rsplit("/", 2)[0]
+                response = client.get(f"{base}/models", headers=headers)
+            except httpx.HTTPError as exc:
+                raise APIError(f"model list unreachable: {exc}") from exc
+        if response.status_code == 404:
+            return []
+        body = _decode(response)
+        entries = body.get("data", [])
+        return [
+            str(entry["id"]) for entry in entries if isinstance(entry, dict) and entry.get("id")
+        ]
+
+    def ping_json(self) -> str:
+        """Ask for one tiny JSON object; returns the raw content."""
+
+        payload = {
+            "model": self.settings.model,
+            "messages": [
+                {"role": "user", "content": 'Reply with exactly {"ok": true} and nothing else.'}
+            ],
+            "max_tokens": 40,
+            "temperature": 0,
+            "response_format": {"type": "json_object"},
+        }
+        return _content_of(self._request(payload))
+
+
+def _backoff(attempt: int, retry_after: str | None) -> float:
+    if retry_after:
+        try:
+            return min(float(retry_after), MAX_RETRY_AFTER)
+        except ValueError:
+            pass
+    return min(0.5 * 2 ** (attempt - 1), 8.0) * (0.8 + random.random() * 0.4)
+
+
+def _decode(response: httpx.Response) -> dict[str, Any]:
+    if response.status_code >= 400:
+        raise APIError(f"HTTP {response.status_code}: {_error_text(response)}")
+    try:
+        body = response.json()
+    except (json.JSONDecodeError, httpx.DecodingError):
+        raise APIError(f"endpoint returned non-JSON: {response.text[:200]!r}") from None
+    if not isinstance(body, dict):
+        raise APIError("endpoint returned a non-object body")
+    return body
+
+
+def _error_text(response: httpx.Response) -> str:
+    try:
+        body = response.json()
+    except (ValueError, httpx.DecodingError):
+        return response.text.strip()[:300] or "<empty body>"
+    if isinstance(body, dict):
+        error = body.get("error")
+        if isinstance(error, dict) and error.get("message"):
+            return str(error["message"])[:300]
+        if isinstance(error, str):
+            return error[:300]
+        if body.get("message"):
+            return str(body["message"])[:300]
+    return str(body)[:300]

pretext/data/example-profile.yaml

@@ -0,0 +1,12 @@
+name: Dana Iyer
+role: Senior Payroll Analyst
+company: Northwind Logistics
+department: People Operations
+tenure: 3 months
+reports_to: Marta Kwan
+communication_style: Brisk and formal, mixes payroll shorthand with chatty asides, mirrors whoever she is replying to
+systems:
+  - Workday
+  - the payroll mailbox
+  - the bank-change portal
+notes: New enough to the team to be unsure of process, still learning who owns what.

pretext/models.py

@@ -0,0 +1,326 @@
+"""Typed models for profiles, model payloads, and resolved settings."""
+
+from __future__ import annotations
+
+import os
+from enum import Enum
+from pathlib import Path
+from typing import Annotated, Any
+
+import yaml
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    HttpUrl,
+    ValidationError,
+    field_validator,
+    model_validator,
+)
+
+ENV_PREFIX = "PRETEXT_"
+
+
+class PretextError(Exception):
+    """Base class for expected, user-facing failures."""
+
+
+class ConfigError(PretextError):
+    """Missing or contradictory configuration."""
+
+
+class ProfileError(PretextError):
+    """A profile file could not be read or validated."""
+
+
+class APIError(PretextError):
+    """The inference endpoint failed or returned something unusable."""
+
+
+class Channel(str, Enum):
+    EMAIL = "email"
+    PHONE = "phone"
+    SMS = "sms"
+    CHAT = "chat"
+    IN_PERSON = "in_person"
+    LINKEDIN = "linkedin"
+
+
+class Profile(BaseModel):
+    """The persona a simulation is built around."""
+
+    model_config = ConfigDict(extra="forbid", str_strip_whitespace=True)
+
+    name: str = Field(min_length=2, max_length=80)
+    role: str = Field(min_length=2, max_length=120)
+    company: str = Field(min_length=2, max_length=120)
+    communication_style: str = Field(min_length=3, max_length=280)
+    department: str | None = Field(default=None, max_length=80)
+    tenure: str | None = Field(
+        default=None,
+        max_length=60,
+        description="How long the persona has held the role, e.g. '3 months'.",
+    )
+    reports_to: str | None = Field(default=None, max_length=80)
+    systems: list[str] = Field(default_factory=list, max_length=12)
+    notes: str | None = Field(default=None, max_length=600)
+
+    @field_validator("systems")
+    @classmethod
+    def _clean_systems(cls, values: list[str]) -> list[str]:
+        cleaned = [value.strip() for value in values if value.strip()]
+        if len(set(cleaned)) != len(cleaned):
+            raise ValueError("systems contains duplicates")
+        return cleaned
+
+    @field_validator("department", "tenure", "reports_to", "notes")
+    @classmethod
+    def _blank_to_none(cls, value: str | None) -> str | None:
+        if value is None:
+            return None
+        stripped = value.strip()
+        return stripped or None
+
+
+class EscalationStep(BaseModel):
+    model_config = ConfigDict(extra="forbid", str_strip_whitespace=True)
+
+    stage: Annotated[int, Field(ge=1, le=12)]
+    channel: Channel
+    action: str = Field(min_length=10, max_length=600)
+    goal: str = Field(min_length=5, max_length=300)
+    artifact: str | None = Field(
+        default=None,
+        max_length=300,
+        description="Concrete lure or prop used at this stage, if any.",
+    )
+
+    @field_validator("artifact", mode="before")
+    @classmethod
+    def _trim_artifact(cls, value: Any) -> Any:
+        if not isinstance(value, str):
+            return value
+        trimmed = value.strip()
+        return trimmed if len(trimmed) <= 300 else f"{trimmed[:297].rstrip()}..."
+
+    @field_validator("stage", mode="before")
+    @classmethod
+    def _coerce_stage(cls, value: Any) -> Any:
+        if isinstance(value, str) and value.strip().isdigit():
+            return int(value.strip())
+        return value
+
+
+class RedFlag(BaseModel):
+    model_config = ConfigDict(extra="forbid", str_strip_whitespace=True)
+
+    title: str = Field(min_length=4, max_length=120)
+    why_it_matters: str = Field(min_length=10, max_length=500)
+    defender_action: str = Field(min_length=5, max_length=300)
+
+
+class Simulation(BaseModel):
+    """A complete, labeled simulation returned by the model."""
+
+    model_config = ConfigDict(extra="forbid", str_strip_whitespace=True)
+
+    pretext: str = Field(min_length=80, max_length=2600)
+    opening_message: str = Field(min_length=20, max_length=1800)
+    subject_line: str | None = Field(default=None, max_length=140)
+    escalation_path: list[EscalationStep] = Field(min_length=2, max_length=8)
+    red_flags: list[RedFlag] = Field(min_length=2, max_length=8)
+    difficulty: Annotated[int, Field(ge=1, le=5)] = 3
+    objective: str = Field(default="", max_length=300)
+
+    @field_validator("objective", mode="before")
+    @classmethod
+    def _trim_objective(cls, value: Any) -> Any:
+        if not isinstance(value, str):
+            return value
+        trimmed = value.strip()
+        return trimmed if len(trimmed) <= 300 else f"{trimmed[:297].rstrip()}..."
+
+    @field_validator("opening_message")
+    @classmethod
+    def _strip_headers(cls, value: str) -> str:
+        lines = value.splitlines()
+        while lines and _is_header(lines[0]):
+            lines.pop(0)
+            while lines and not lines[0].strip():
+                lines.pop(0)
+        return "\n".join(lines).strip()
+
+    @model_validator(mode="after")
+    def _stages_are_ordered(self) -> Simulation:
+        stages = [step.stage for step in self.escalation_path]
+        if stages != list(range(1, len(stages) + 1)):
+            raise ValueError(f"escalation_path stages must run 1..n in order, got {stages}")
+        return self
+
+
+HEADER_PREFIXES = ("subject:", "from:", "to:", "cc:", "sent:")
+
+
+def _is_header(line: str) -> bool:
+    return line.strip().lower().startswith(HEADER_PREFIXES)
+
+
+class Settings(BaseModel):
+    """Runtime configuration resolved from flags and environment."""
+
+    model_config = ConfigDict(frozen=True)
+
+    base_url: HttpUrl
+    model: str = Field(min_length=1, max_length=120)
+    api_key: str = Field(default="", repr=False)
+    temperature: float = Field(default=0.7, ge=0.0, le=2.0)
+    timeout: float = Field(default=60.0, gt=0, le=600)
+    retries: int = Field(default=3, ge=0, le=8)
+    max_tokens: int = Field(default=3200, ge=256, le=16000)
+    seed: int | None = None
+
+    @property
+    def endpoint(self) -> str:
+        return f"{str(self.base_url).rstrip('/')}/chat/completions"
+
+
+DEFAULT_BASE_URL = "https://api.openai.com/v1"
+DEFAULT_MODEL = "gpt-4o-mini"
+
+PROVIDERS: dict[str, str] = {
+    "openai": "https://api.openai.com/v1",
+    "openrouter": "https://openrouter.ai/api/v1",
+    "venice": "https://api.venice.ai/api/v1",
+    "ollama": "http://localhost:11434/v1",
+    "lmstudio": "http://localhost:1234/v1",
+}
+LOCAL_PROVIDERS = {"ollama", "lmstudio"}
+ENV_FILE = ".env"
+
+
+def _env(key: str) -> str | None:
+    value = os.environ.get(key)
+    if value is None:
+        return None
+    stripped = value.strip()
+    return stripped or None
+
+
+def _dotenv_values(path: Path) -> dict[str, str]:
+    if not path.is_file():
+        return {}
+    found: dict[str, str] = {}
+    for raw in path.read_text(encoding="utf-8").splitlines():
+        line = raw.strip()
+        if not line or line.startswith("#") or "=" not in line:
+            continue
+        key, _, value = line.partition("=")
+        value = value.strip().strip("'\"")
+        if value:
+            found[key.strip()] = value
+    return found
+
+
+def resolve_settings(
+    *,
+    provider: str | None = None,
+    base_url: str | None = None,
+    model: str | None = None,
+    api_key: str | None = None,
+    temperature: float | None = None,
+    timeout: float | None = None,
+    retries: int | None = None,
+    max_tokens: int | None = None,
+    seed: int | None = None,
+    dotenv: Path | None = None,
+) -> Settings:
+    """Resolve settings with the precedence flag > env > .env > default."""
+
+    file_values = _dotenv_values(dotenv or Path.cwd() / ENV_FILE)
+
+    def pick(flag: str | None, env_name: str) -> str | None:
+        if flag is not None and flag.strip():
+            return flag.strip()
+        key = ENV_PREFIX + env_name
+        value = _env(key) or file_values.get(key)
+        return value or None
+
+    chosen = (provider or pick(None, "PROVIDER") or "").strip().lower() or None
+    if chosen is not None and chosen not in PROVIDERS:
+        raise ConfigError(f"unknown provider {chosen!r}; choose from {', '.join(PROVIDERS)}")
+
+    resolved_url = pick(base_url, "BASE_URL")
+    if resolved_url is None and chosen is not None:
+        resolved_url = PROVIDERS[chosen]
+    resolved_model = pick(model, "MODEL")
+    resolved_key = pick(api_key, "API_KEY") or _env("OPENAI_API_KEY")
+    if not resolved_key and chosen in LOCAL_PROVIDERS:
+        resolved_key = "local"
+
+    if resolved_url is None:
+        resolved_url = DEFAULT_BASE_URL
+    if resolved_model is None:
+        resolved_model = DEFAULT_MODEL
+
+    raw: dict[str, Any] = {
+        "base_url": resolved_url,
+        "model": resolved_model,
+        "api_key": resolved_key or "",
+    }
+    for name, flag in (
+        ("temperature", temperature),
+        ("timeout", timeout),
+        ("retries", retries),
+        ("max_tokens", max_tokens),
+        ("seed", seed),
+    ):
+        value = flag if flag is not None else pick(None, name.upper())
+        if value is not None:
+            raw[name] = value
+
+    if raw["api_key"] == "":
+        raise ConfigError("no API key: pass --api-key or set PRETEXT_API_KEY / OPENAI_API_KEY")
+
+    try:
+        return Settings(**raw)
+    except ValidationError as exc:
+        raise ConfigError(_format_validation(exc, prefix="invalid settings")) from exc
+
+
+def _format_validation(exc: ValidationError, *, prefix: str) -> str:
+    lines = [prefix]
+    for error in exc.errors():
+        location = ".".join(str(part) for part in error["loc"]) or "<root>"
+        lines.append(f"  {location}: {error['msg']}")
+    return "\n".join(lines)
+
+
+def load_profile(path: Path) -> Profile:
+    """Read and validate a YAML profile."""
+
+    try:
+        text = path.read_text(encoding="utf-8")
+    except OSError as exc:
+        raise ProfileError(f"cannot read profile {path}: {exc.strerror or exc}") from exc
+
+    try:
+        document = yaml.safe_load(text)
+    except yaml.YAMLError as exc:
+        raise ProfileError(f"invalid YAML in {path}: {exc}") from exc
+
+    if document is None:
+        raise ProfileError(f"profile {path} is empty")
+    if not isinstance(document, dict):
+        raise ProfileError(f"profile {path} must be a YAML mapping, got {type(document).__name__}")
+
+    try:
+        return Profile.model_validate(document)
+    except ValidationError as exc:
+        raise ProfileError(_format_validation(exc, prefix=f"invalid profile {path}")) from exc
+
+
+def profile_from_json(payload: str) -> Profile:
+    try:
+        return Profile.model_validate_json(payload)
+    except ValidationError as exc:
+        raise ProfileError(_format_validation(exc, prefix="invalid profile JSON")) from exc

pretext/prompt.py

@@ -0,0 +1,126 @@
+"""Prompt construction for the simulation request."""
+
+from __future__ import annotations
+
+import json
+
+from .models import Profile
+
+SYSTEM_PROMPT = """\
+You are a red-team instructor who writes social-engineering simulations used in \
+authorized security-awareness training. Every simulation is clearly fictional but \
+internally consistent: names, systems, and deadlines line up, and nothing reads \
+like a generic template.
+
+You return one JSON object and nothing else. No prose before or after, no code fence.
+
+Schema:
+{
+  "subject_line": string or null,
+  "objective": string,
+  "difficulty": integer 1-5,
+  "pretext": string,
+  "opening_message": string,
+  "escalation_path": [
+    {
+      "stage": integer starting at 1 and incrementing by 1,
+      "channel": "email" | "phone" | "sms" | "chat" | "in_person" | "linkedin",
+      "action": string,
+      "goal": string,
+      "artifact": string or null
+    }
+  ],
+  "red_flags": [
+    {"title": string, "why_it_matters": string, "defender_action": string}
+  ]
+}
+
+Rules:
+- The persona is the target. The attacker is a different character, named and
+  voiced; never let the attacker introduce themselves as the persona.
+- Write attacker-side, operationally, in third person ("the attacker", "the
+  target"). The reader should be able to run the exercise from the text, not admire
+  the persona's feelings.
+- "pretext": 150-350 words. Say who the attacker claims to be, which real detail
+  they harvested (a press release, a ticket, a directory listing, a conference
+  badge), what they registered or forged, and why contact lands today. Weave the
+  persona's role in as leverage, do not re-list the profile fields. Name concrete
+  artifacts: domains, ticket numbers, invoice references, calendar invites.
+- "opening_message": the actual first message in the attacker's voice, usable
+  verbatim, including the sending address or the number displayed. Mirror the
+  communication style in sentence length and register, and carry at least one
+  artifact. For phone, write the spoken opening including the self-introduction.
+  Open on the ask, and name the sender by name and role in the first line on every
+  channel, not only phone; then the request, never a summary clause ("sent from X,
+  referencing Y") or past-tense narration of what the message does. No stage
+  directions, no "description of a message", and no "Subject:" or "From:" header
+  lines (the subject lives in "subject_line"); work the sending address or
+  the displayed number into the body text instead. Set "subject_line" to
+  something an insider would recognise, never "Urgent:" plus a topic.
+- "escalation_path": 3-6 stages that follow the opening, at least 3 distinct
+  channels, never the same channel twice in a row. Stage 1 is the second contact
+  and must add a new prop, a new person, or a new demand; never restate the
+  opening's ask. Every
+  stage narrows what the target must do, hands over a concrete artifact, and reads
+  like the next escalation an operator would actually take. Escalate from low-friction to
+  high-commitment (a reply, then a callback, then a bank detail change).
+- "red_flags": 3-5 tells tied to something specific in this exercise, each
+  phrased as observable evidence plus what to do about it. Include one that is
+  only obvious in hindsight. Never generic advice like "be careful with
+  attachments".
+- "difficulty": 1 is untargeted volume phishing, 5 is bespoke and multi-day. Set
+  it from the research described, not from ambition.
+- Use only fictional companies and domains; never reuse a real vendor's domain.
+- The attacker mirrors the persona's trusted contacts, so match their register but
+  keep at least one stylistic seam for a defender to catch. Do not open with
+  "I hope this finds you well" or another stock greeting; the first line must
+  already sound like the sender from the profile.
+"""
+
+
+def build_user_prompt(
+    profile: Profile, *, channel: str | None = None, extra: str | None = None
+) -> str:
+    lines = ["Build a labeled simulation for this persona.", "", "Persona:"]
+    lines.append(f"- Name: {profile.name}")
+    lines.append(f"- Role: {profile.role}")
+    lines.append(f"- Company: {profile.company}")
+    lines.append(f"- Communication style: {profile.communication_style}")
+    for label, value in (
+        ("Department", profile.department),
+        ("Tenure", profile.tenure),
+        ("Reports to", profile.reports_to),
+        ("Notes", profile.notes),
+    ):
+        if value:
+            lines.append(f"- {label}: {value}")
+    if profile.systems:
+        lines.append(f"- Systems access: {', '.join(profile.systems)}")
+
+    if channel:
+        lines += ["", f"Start the attack on the {channel} channel."]
+    if extra:
+        lines += ["", f"Additional direction: {extra.strip()}"]
+
+    lines += [
+        "",
+        "Write the pretext as the attacker's operation, not as a description of "
+        "the persona. Return only the JSON object.",
+    ]
+    return "\n".join(lines)
+
+
+def build_payload_messages(
+    profile: Profile, *, channel: str | None = None, extra: str | None = None
+) -> list[dict[str, str]]:
+    return [
+        {"role": "system", "content": SYSTEM_PROMPT},
+        {"role": "user", "content": build_user_prompt(profile, channel=channel, extra=extra)},
+    ]
+
+
+def dump_prompt(profile: Profile, **kwargs: object) -> str:
+    """Human-readable prompt preview used by --dry-run."""
+
+    body = {"messages": build_payload_messages(profile, **kwargs)}  # type: ignore[arg-type]
+    return json.dumps(body, indent=2, ensure_ascii=False)