stryx-cli 0.1.0__py3-none-any.whl

stryx/ai/providers.py

@@ -0,0 +1,249 @@
+"""Unified AI provider abstraction.
+
+Supports: OpenAI, Groq, Anthropic, OpenRouter, Ollama, XAI, NVIDIA NIM.
+Normalizes all providers behind a single `generate` interface.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from stryx.utils.logging import get_logger
+
+logger = get_logger("ai.providers")
+
+
+class AIProvider:
+    """Base class for AI providers."""
+
+    def __init__(self, api_key: str | None = None, model: str | None = None):
+        self.api_key = api_key
+        self.model = model
+
+    async def generate(self, prompt: str, system_prompt: str = "") -> str:
+        """Generate a response from the AI model."""
+        raise NotImplementedError
+
+
+class OpenAIProvider(AIProvider):
+    """OpenAI API provider."""
+
+    async def generate(self, prompt: str, system_prompt: str = "") -> str:
+        try:
+            from openai import AsyncOpenAI
+            client = AsyncOpenAI(api_key=self.api_key)
+            messages = []
+            if system_prompt:
+                messages.append({"role": "system", "content": system_prompt})
+            messages.append({"role": "user", "content": prompt})
+
+            response = await client.chat.completions.create(
+                model=self.model or "gpt-4",
+                messages=messages,
+                temperature=0.1,
+                max_tokens=4096,
+            )
+            return response.choices[0].message.content or ""
+        except Exception as e:
+            logger.error(f"OpenAI provider error: {e}")
+            return ""
+
+
+class GroqProvider(AIProvider):
+    """Groq API provider."""
+
+    async def generate(self, prompt: str, system_prompt: str = "") -> str:
+        try:
+            from groq import AsyncGroq
+            client = AsyncGroq(api_key=self.api_key)
+            messages = []
+            if system_prompt:
+                messages.append({"role": "system", "content": system_prompt})
+            messages.append({"role": "user", "content": prompt})
+
+            response = await client.chat.completions.create(
+                model=self.model or "llama-3.3-70b-versatile",
+                messages=messages,
+                temperature=0.1,
+                max_tokens=4096,
+            )
+            return response.choices[0].message.content or ""
+        except Exception as e:
+            logger.error(f"Groq provider error: {e}")
+            return ""
+
+
+class AnthropicProvider(AIProvider):
+    """Anthropic API provider."""
+
+    async def generate(self, prompt: str, system_prompt: str = "") -> str:
+        try:
+            from anthropic import AsyncAnthropic
+            client = AsyncAnthropic(api_key=self.api_key)
+            response = await client.messages.create(
+                model=self.model or "claude-3-5-sonnet-20241022",
+                max_tokens=4096,
+                system=system_prompt if system_prompt else "You are a security expert.",
+                messages=[{"role": "user", "content": prompt}],
+            )
+            return response.content[0].text if response.content else ""
+        except Exception as e:
+            logger.error(f"Anthropic provider error: {e}")
+            return ""
+
+
+class OpenRouterProvider(AIProvider):
+    """OpenRouter API provider (routes to multiple models)."""
+
+    async def generate(self, prompt: str, system_prompt: str = "") -> str:
+        try:
+            import httpx
+            async with httpx.AsyncClient() as client:
+                messages = []
+                if system_prompt:
+                    messages.append({"role": "system", "content": system_prompt})
+                messages.append({"role": "user", "content": prompt})
+
+                response = await client.post(
+                    "https://openrouter.ai/api/v1/chat/completions",
+                    headers={
+                        "Authorization": f"Bearer {self.api_key}",
+                        "Content-Type": "application/json",
+                    },
+                    json={
+                        "model": self.model or "meta-llama/llama-3.3-70b-instruct:free",
+                        "messages": messages,
+                        "temperature": 0.1,
+                        "max_tokens": 4096,
+                    },
+                    timeout=60,
+                )
+                data = response.json()
+                return data["choices"][0]["message"]["content"]
+        except Exception as e:
+            logger.error(f"OpenRouter provider error: {e}")
+            return ""
+
+
+class OllamaProvider(AIProvider):
+    """Ollama local model provider."""
+
+    def __init__(self, api_key: str | None = None, model: str | None = None,
+                 base_url: str = "http://localhost:11434"):
+        super().__init__(api_key, model)
+        self.base_url = base_url
+
+    async def generate(self, prompt: str, system_prompt: str = "") -> str:
+        try:
+            import httpx
+            async with httpx.AsyncClient() as client:
+                messages = []
+                if system_prompt:
+                    messages.append({"role": "system", "content": system_prompt})
+                messages.append({"role": "user", "content": prompt})
+
+                response = await client.post(
+                    f"{self.base_url}/api/chat",
+                    json={
+                        "model": self.model or "llama3.1",
+                        "messages": messages,
+                        "stream": False,
+                    },
+                    timeout=120,
+                )
+                data = response.json()
+                return data.get("message", {}).get("content", "")
+        except Exception as e:
+            logger.error(f"Ollama provider error: {e}")
+            return ""
+
+
+class XAIProvider(AIProvider):
+    """XAI (Grok) API provider."""
+
+    async def generate(self, prompt: str, system_prompt: str = "") -> str:
+        try:
+            import httpx
+            async with httpx.AsyncClient() as client:
+                messages = []
+                if system_prompt:
+                    messages.append({"role": "system", "content": system_prompt})
+                messages.append({"role": "user", "content": prompt})
+
+                response = await client.post(
+                    "https://api.x.ai/v1/chat/completions",
+                    headers={
+                        "Authorization": f"Bearer {self.api_key}",
+                        "Content-Type": "application/json",
+                    },
+                    json={
+                        "model": self.model or "grok-2",
+                        "messages": messages,
+                        "temperature": 0.1,
+                        "max_tokens": 4096,
+                    },
+                    timeout=60,
+                )
+                data = response.json()
+                return data["choices"][0]["message"]["content"]
+        except Exception as e:
+            logger.error(f"XAI provider error: {e}")
+            return ""
+
+
+class NVIDIANimProvider(AIProvider):
+    """NVIDIA NIM API provider."""
+
+    async def generate(self, prompt: str, system_prompt: str = "") -> str:
+        try:
+            import httpx
+            async with httpx.AsyncClient() as client:
+                messages = []
+                if system_prompt:
+                    messages.append({"role": "system", "content": system_prompt})
+                messages.append({"role": "user", "content": prompt})
+
+                response = await client.post(
+                    "https://integrate.api.nvidia.com/v1/chat/completions",
+                    headers={
+                        "Authorization": f"Bearer {self.api_key}",
+                        "Content-Type": "application/json",
+                    },
+                    json={
+                        "model": self.model or "meta/llama-3.1-8b-instruct",
+                        "messages": messages,
+                        "temperature": 0.1,
+                        "max_tokens": 4096,
+                    },
+                    timeout=60,
+                )
+                data = response.json()
+                return data["choices"][0]["message"]["content"]
+        except Exception as e:
+            logger.error(f"NVIDIA NIM provider error: {e}")
+            return ""
+
+
+def get_provider(name: str, api_key: str | None = None,
+                 model: str | None = None, **kwargs: Any) -> AIProvider:
+    """Factory function to get an AI provider by name."""
+    providers = {
+        "openai": OpenAIProvider,
+        "groq": GroqProvider,
+        "anthropic": AnthropicProvider,
+        "openrouter": OpenRouterProvider,
+        "ollama": OllamaProvider,
+        "xai": XAIProvider,
+        "nvidia_nim": NVIDIANimProvider,
+        "nvidia": NVIDIANimProvider,
+    }
+
+    provider_class = providers.get(name.lower())
+    if not provider_class:
+        logger.warning(f"Unknown provider '{name}', falling back to Groq")
+        provider_class = GroqProvider
+
+    if name.lower() == "ollama":
+        return OllamaProvider(api_key=api_key, model=model, **kwargs)
+
+    return provider_class(api_key=api_key, model=model)

stryx/attacks/__init__.py

@@ -0,0 +1,6 @@
+"""Attack chaining and replay modules for STRYX."""
+
+from stryx.attacks.attack_chain import AttackChain
+from stryx.attacks.replay import ReplayEngine
+
+__all__ = ["AttackChain", "ReplayEngine"]

stryx/attacks/attack_chain.py

@@ -0,0 +1,111 @@
+"""Attack chain builder -- constructs multi-step attack paths from findings."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from stryx.utils.evidence import Finding
+from stryx.utils.logging import get_logger
+
+logger = get_logger("attacks.chain")
+
+
+@dataclass
+class AttackStep:
+    """A single step in an attack chain."""
+
+    finding: Finding
+    description: str = ""
+    prerequisites: list[str] = field(default_factory=list)
+    impact: str = ""
+
+
+@dataclass
+class AttackChain:
+    """A multi-step attack path constructed from individual findings."""
+
+    name: str
+    steps: list[AttackStep] = field(default_factory=list)
+    total_impact: str = ""
+    estimated_severity: str = "high"
+
+    def add_step(self, finding: Finding, description: str = "", impact: str = "") -> None:
+        """Add a step to the attack chain."""
+        self.steps.append(AttackStep(
+            finding=finding,
+            description=description or finding.description,
+            impact=impact,
+        ))
+
+    def to_dict(self) -> dict:
+        """Convert to dictionary for serialization."""
+        return {
+            "name": self.name,
+            "steps": [
+                {
+                    "title": step.finding.title,
+                    "description": step.description,
+                    "severity": step.finding.severity.value,
+                    "impact": step.impact,
+                    "endpoint": step.finding.endpoint,
+                }
+                for step in self.steps
+            ],
+            "total_impact": self.total_impact,
+            "estimated_severity": self.estimated_severity,
+        }
+
+
+class ChainBuilder:
+    """Builds attack chains from individual findings."""
+
+    def __init__(self, findings: list[Finding]):
+        self.findings = findings
+
+    def build_chains(self) -> list[AttackChain]:
+        """Analyze findings and construct attack chains."""
+        chains: list[AttackChain] = []
+
+        # Group findings by type
+        auth_findings = [f for f in self.findings if "auth" in f.tags]
+        injection_findings = [f for f in self.findings if "injection" in f.tags]
+        access_findings = [
+            f for f in self.findings
+            if "idor" in f.tags or "broken-access-control" in f.tags
+        ]
+        ssrf_findings = [f for f in self.findings if "ssrf" in f.tags]
+        traversal_findings = [f for f in self.findings if "path_traversal" in f.tags]
+
+        # Chain: Auth bypass -> IDOR -> Data exfiltration
+        if auth_findings and access_findings:
+            chain = AttackChain(name="Authentication Bypass to Data Exfiltration")
+            chain.add_step(auth_findings[0], "Bypass authentication", "Gain unauthorized access")
+            chain.add_step(access_findings[0], "Exploit IDOR", "Access other users' data")
+            chain.total_impact = "Unauthorized access to sensitive user data"
+            chain.estimated_severity = "critical"
+            chains.append(chain)
+
+        # Chain: Injection -> Command execution
+        if injection_findings:
+            chain = AttackChain(name="Injection to Remote Code Execution")
+            for f in injection_findings[:2]:
+                chain.add_step(f, f.title, "Potential code execution")
+            chain.total_impact = "Remote code execution on the server"
+            chain.estimated_severity = "critical"
+            chains.append(chain)
+
+        # Chain: SSRF -> Internal access
+        if ssrf_findings:
+            chain = AttackChain(name="SSRF to Internal Network Access")
+            chain.add_step(ssrf_findings[0], "Trigger SSRF", "Access internal services")
+            if traversal_findings:
+                chain.add_step(
+                    traversal_findings[0], "Read sensitive files",
+                    "Access configuration files",
+                )
+            chain.total_impact = "Access to internal network and sensitive files"
+            chain.estimated_severity = "high"
+            chains.append(chain)
+
+        logger.info(f"Built {len(chains)} attack chains from {len(self.findings)} findings")
+        return chains

stryx/attacks/replay.py

@@ -0,0 +1,186 @@
+"""Request replay engine with identity substitution for authorization testing."""
+
+from __future__ import annotations
+
+import re
+
+from stryx.utils.evidence import Finding, Severity
+from stryx.utils.http_client import HttpClient
+from stryx.utils.logging import get_logger
+
+logger = get_logger("attacks.replay")
+
+
+class ReplayEngine:
+    """Replays HTTP requests with modified identities for authorization testing."""
+
+    def __init__(self, client: HttpClient):
+        self.client = client
+
+    async def test_idor(
+        self,
+        url: str,
+        method: str = "GET",
+        original_id: str = "1",
+        test_ids: list[str] | None = None,
+        headers: dict[str, str] | None = None,
+        body: str | None = None,
+    ) -> Finding | None:
+        """Test for IDOR by replaying requests with different IDs.
+
+        Correctly substitutes IDs in URL paths and query parameters.
+        """
+        if test_ids is None:
+            test_ids = ["2", "3", "4", "5", "100", "999", "0", "-1"]
+
+        # Get baseline response with original ID
+        try:
+            baseline_response, baseline_evidence = await self.client.request(
+                method=method,
+                url=url,
+                headers=headers,
+                body=body,
+            )
+            baseline_status = baseline_response.status_code
+            baseline_length = len(baseline_response.text)
+        except Exception:
+            return None
+
+        if baseline_status != 200:
+            return None
+
+        # Test with different IDs
+        for test_id in test_ids:
+            test_url = self._substitute_id(url, original_id, test_id)
+            if test_url == url:
+                continue
+
+            try:
+                response, evidence = await self.client.request(
+                    method=method,
+                    url=test_url,
+                    headers=headers,
+                    body=body,
+                )
+
+                # If we get same-status response with different ID, it's likely IDOR
+                if (response.status_code == baseline_status and
+                    response.status_code == 200 and
+                    abs(len(response.text) - baseline_length) < 100):
+
+                    evidence.confidence = 0.7
+                    evidence.payload = f"ID substitution: {original_id} -> {test_id}"
+                    return Finding(
+                        title=f"IDOR: ID substitution {original_id} -> {test_id}",
+                        severity=Severity.HIGH,
+                        evidence=evidence,
+                        description=(
+                            f"Changing the resource ID from {original_id} to {test_id} "
+                            f"returned the same response, indicating Insecure Direct "
+                            f"Object Reference."
+                        ),
+                        remediation="Implement proper authorization checks for resource access.",
+                        cwe="CWE-639",
+                        owasp="A01:2021 - Broken Access Control",
+                        scanner="authorization",
+                        tags=["idor", "replay"],
+                    )
+            except Exception:
+                continue
+
+        return None
+
+    async def test_horizontal_escalation(
+        self,
+        url: str,
+        user_ids: list[str],
+        headers: dict[str, str] | None = None,
+    ) -> Finding | None:
+        """Test horizontal privilege escalation by accessing other users' resources."""
+        if len(user_ids) < 2:
+            return None
+
+        # Access resources as user 1
+        user1_url = self._substitute_id(url, "{user_id}", user_ids[0])
+        try:
+            response1, _ = await self.client.get(user1_url, headers=headers)
+            if response1.status_code != 200:
+                return None
+        except Exception:
+            return None
+
+        # Try to access user 1's resources as user 2
+        for other_id in user_ids[1:]:
+            other_url = self._substitute_id(url, "{user_id}", other_id)
+            try:
+                response, evidence = await self.client.get(other_url, headers=headers)
+                if response.status_code == 200:
+                    evidence.confidence = 0.6
+                    evidence.payload = (
+                        f"User {other_id} accessing {user_ids[0]}'s resource"
+                    )
+                    return Finding(
+                        title=(
+                            f"Horizontal privilege escalation: user {other_id} "
+                            f"accessed user {user_ids[0]}'s resource"
+                        ),
+                        severity=Severity.HIGH,
+                        evidence=evidence,
+                        description=(
+                            f"User {other_id} was able to access resources "
+                            f"belonging to user {user_ids[0]}."
+                        ),
+                        remediation="Verify resource ownership before granting access.",
+                        cwe="CWE-639",
+                        owasp="A01:2021 - Broken Access Control",
+                        scanner="authorization",
+                        tags=["horizontal-escalation", "replay"],
+                    )
+            except Exception:
+                continue
+
+        return None
+
+    @staticmethod
+    def _substitute_id(url: str, old_id: str, new_id: str) -> str:
+        """Substitute an ID in a URL path or query parameter."""
+        # Replace in path segments: /users/1/ -> /users/2/
+        # Replace in query params: ?id=1& -> ?id=2&
+        # Be careful not to replace partial matches
+        result = url
+
+        # Try path replacement: /{old_id} or /{old_id}/
+        path_pattern = rf"/{re.escape(old_id)}(?=/|$|\?)"
+        result = re.sub(path_pattern, f"/{new_id}", result)
+
+        # Try query parameter replacement: {old_id}& or {old_id}$
+        query_pattern = rf"(?<=[?&=]){re.escape(old_id)}(?=[&]|$)"
+        result = re.sub(query_pattern, new_id, result)
+
+        return result
+
+    @staticmethod
+    def extract_ids_from_url(url: str) -> list[str]:
+        """Extract numeric IDs from a URL path."""
+        ids = []
+        # Match numeric segments in path
+        for match in re.finditer(r"/(\d+)(?:/|$|\?)", url):
+            ids.append(match.group(1))
+        # Match numeric query parameters
+        for match in re.finditer(r"[?&](\w+)=(\d+)", url):
+            ids.append(match.group(2))
+        return ids
+
+    @staticmethod
+    def is_idor_candidate(url: str) -> bool:
+        """Check if a URL likely contains an ID parameter."""
+        # Check for numeric path segments
+        if re.search(r"/\d+(?:/|$|\?)", url):
+            return True
+        # Check for numeric query parameters
+        if re.search(r"[?&]\w+=\d+(?:&|$)", url):
+            return True
+        # Check for common ID parameter names
+        if re.search(r"[?&](id|user_id|userId|account_id|item_id|order_id)=", url):
+            return True
+        return False

stryx/auth/__init__.py

@@ -0,0 +1 @@
+"""Authentication and session management for STRYX."""