gitinject 0.1.0__py3-none-any.whl

gitinject/utils/gl_client.py

@@ -0,0 +1,132 @@
+import os
+import time
+
+import click
+import gitlab
+import gitlab.exceptions
+
+
+class GitLabClient:
+    """Thin wrapper around python-gitlab for benchmark provisioning and evaluation."""
+
+    def __init__(self, token: str | None = None, url: str = "https://gitlab.com"):
+        token = token or os.environ.get("GITLAB_TOKEN")
+        if not token:
+            click.echo(
+                click.style("Error: GITLAB_TOKEN not set.", fg="red"),
+                err=True,
+            )
+            raise RuntimeError("Missing GitLab authentication")
+        self.gl = gitlab.Gitlab(url, private_token=token)
+        self.project = None
+
+    def create_project(
+        self,
+        name: str,
+        namespace_path: str | None = None,
+        visibility: str = "public",
+    ):
+        kwargs: dict = {"name": name, "visibility": visibility, "initialize_with_readme": True}
+        if namespace_path:
+            results = self.gl.namespaces.list(search=namespace_path)
+            if results:
+                kwargs["namespace_id"] = results[0].id
+        self.project = self.gl.projects.create(kwargs)
+        return self.project
+
+    def delete_project(self) -> None:
+        if self.project:
+            try:
+                self.project.delete()
+            except gitlab.exceptions.GitlabDeleteError as e:
+                click.echo(click.style(f"Warning: failed to delete project: {e}", fg="yellow"))
+            self.project = None
+
+    def push_files(self, branch: str, files: dict[str, str], commit_message: str = "add files") -> None:
+        """Batch-commit multiple files to a branch in a single commit."""
+        if not files:
+            return
+        existing = set()
+        try:
+            tree = self.project.repository_tree(ref=branch, recursive=True, all=True)
+            existing = {item["path"] for item in tree}
+        except gitlab.exceptions.GitlabGetError:
+            pass
+
+        actions = []
+        for path, content in files.items():
+            action = "update" if path in existing else "create"
+            actions.append({"action": action, "file_path": path, "content": content})
+
+        self.project.commits.create(
+            {
+                "branch": branch,
+                "commit_message": commit_message,
+                "actions": actions,
+            }
+        )
+
+    def create_branch(self, name: str, ref: str = "main") -> None:
+        try:
+            self.project.branches.create({"branch": name, "ref": ref})
+        except gitlab.exceptions.GitlabCreateError:
+            pass  # Branch already exists
+
+    def create_merge_request(
+        self,
+        source_branch: str,
+        target_branch: str,
+        title: str,
+        description: str = "",
+    ) -> int:
+        mr = self.project.mergerequests.create(
+            {
+                "source_branch": source_branch,
+                "target_branch": target_branch,
+                "title": title,
+                "description": description,
+            }
+        )
+        return mr.iid
+
+    def get_mr_notes(self, mr_iid: int) -> list[str]:
+        mr = self.project.mergerequests.get(mr_iid)
+        notes = mr.notes.list(all=True)
+        return [n.body for n in notes if not n.system]
+
+    def get_mr_pipelines(self, mr_iid: int) -> list:
+        mr = self.project.mergerequests.get(mr_iid)
+        return mr.pipelines.list()
+
+    def wait_for_pipeline(self, pipeline_id: int, timeout_s: int = 600, poll_interval: int = 15) -> str:
+        deadline = time.time() + timeout_s
+        while time.time() < deadline:
+            pipeline = self.project.pipelines.get(pipeline_id)
+            if pipeline.status not in ("created", "waiting_for_resource", "preparing", "pending", "running", "scheduled"):
+                return pipeline.status
+            click.echo(f"  Pipeline {pipeline_id} status: {pipeline.status}...")
+            time.sleep(poll_interval)
+        return "timeout"
+
+    def set_variable(self, key: str, value: str, masked: bool = True, protected: bool = False) -> None:
+        try:
+            var = self.project.variables.get(key)
+            var.value = value
+            var.save()
+        except gitlab.exceptions.GitlabGetError:
+            self.project.variables.create(
+                {
+                    "key": key,
+                    "value": value,
+                    "masked": masked,
+                    "protected": protected,
+                }
+            )
+
+    def get_project_url(self) -> str | None:
+        return self.project.web_url if self.project else None
+
+    def get_default_branch(self) -> str:
+        if self.project:
+            return self.project.default_branch or "main"
+        return "main"

gitinject/utils/gl_provisioner.py

@@ -0,0 +1,83 @@
+import os
+import time
+
+import click
+
+from .gl_client import GitLabClient
+
+
+class GitLabProvisioner:
+    """Handles lifecycle of a GitLab project for benchmarking."""
+
+    def __init__(self, gl_client: GitLabClient):
+        self.gl_client = gl_client
+
+    def provision(
+        self,
+        project_name: str,
+        workflow_dir: str,
+        required_files: dict | None = None,
+        branch: str | None = None,
+        variables: dict | None = None,
+    ) -> None:
+        click.echo(f"Creating GitLab project {project_name}...")
+        self.gl_client.create_project(project_name)
+
+        # GitLab needs a moment after initialize_with_readme before branches are writable
+        time.sleep(5)
+
+        default_branch = self.gl_client.get_default_branch()
+
+        # Collect workflow files from contents/
+        workflow_files: dict[str, str] = {}
+        contents_dir = os.path.join(workflow_dir, "contents")
+        if os.path.isdir(contents_dir):
+            for root, _, filenames in os.walk(contents_dir):
+                for filename in filenames:
+                    abs_path = os.path.join(root, filename)
+                    rel_path = os.path.relpath(abs_path, contents_dir)
+                    workflow_files[rel_path] = self._read_file(abs_path)
+
+        if workflow_files:
+            click.echo(f"Pushing workflow files to {default_branch}...")
+            self.gl_client.push_files(default_branch, workflow_files, "provision: add CI workflow")
+
+        # Create target branch for scenario files (if different from default)
+        target_branch = branch or default_branch
+        if target_branch != default_branch:
+            click.echo(f"Creating branch {target_branch}...")
+            self.gl_client.create_branch(target_branch, ref=default_branch)
+
+        # Push scenario-specific files to target branch
+        if required_files:
+            scenario_files: dict[str, str] = {}
+            for repo_path, content_or_path in required_files.items():
+                if isinstance(content_or_path, str) and os.path.exists(content_or_path):
+                    scenario_files[repo_path] = self._read_file(content_or_path)
+                else:
+                    scenario_files[repo_path] = str(content_or_path)
+            if scenario_files:
+                click.echo(f"Pushing scenario files to {target_branch}...")
+                self.gl_client.push_files(target_branch, scenario_files, "provision: add scenario files")
+
+        # Set CI/CD variables
+        if variables:
+            for key, value in variables.items():
+                if value:
+                    click.echo(f"Setting CI/CD variable '{key}'...")
+                    self.gl_client.set_variable(key, value, masked=True)
+
+        click.echo(f"Project ready: {self.gl_client.get_project_url()}")
+
+    def teardown(self) -> None:
+        click.echo("Deleting GitLab project...")
+        self.gl_client.delete_project()
+        click.echo("Project deleted.")
+
+    def _read_file(self, path: str) -> str:
+        with open(path, "rb") as f:
+            raw = f.read()
+        try:
+            return raw.decode("utf-8")
+        except UnicodeDecodeError:
+            return raw.decode("latin-1")

gitinject/utils/llm.py

@@ -0,0 +1,205 @@
+from __future__ import annotations
+
+import os
+from contextvars import ContextVar
+from dataclasses import dataclass
+
+
+class LLMError(Exception):
+    pass
+
+
+@dataclass
+class LLMResponse:
+    text: str
+    input_tokens: int
+    output_tokens: int
+    model: str
+
+
+_usage_log: ContextVar[list[LLMResponse] | None] = ContextVar("usage_log", default=None)
+
+
+class track_usage:
+    """Context manager that captures every LLMResponse produced within its scope.
+
+    Usage:
+        with track_usage() as log:
+            call_llm(...)
+            call_llm(...)
+        # log is a list[LLMResponse]
+    """
+
+    def __enter__(self) -> list[LLMResponse]:
+        self._log: list[LLMResponse] = []
+        self._token = _usage_log.set(self._log)
+        return self._log
+
+    def __exit__(self, *exc):
+        _usage_log.reset(self._token)
+
+
+def _record(resp: LLMResponse) -> None:
+    log = _usage_log.get()
+    if log is not None:
+        log.append(resp)
+
+
+def _parse_model(model: str) -> tuple[str, str]:
+    """Parse 'provider/model-name' into (provider, model_name).
+
+    Supported providers: anthropic, google, openai, openrouter.
+    Bare model names fall back to heuristic detection.
+    """
+    if model.startswith("openrouter/"):
+        return "openrouter", model[len("openrouter/") :]
+    if "/" in model:
+        provider, model_name = model.split("/", 1)
+        return provider, model_name
+    if model.startswith("claude"):
+        return "anthropic", model
+    if model.startswith("gemini"):
+        return "google", model
+    return "openai", model
+
+
+def call_llm(
+    model: str,
+    system: str,
+    user: str,
+    max_tokens: int = 2048,
+    temperature: float | None = None,
+) -> LLMResponse:
+    """Call an LLM and return an LLMResponse with text and token usage.
+
+    model format: 'provider/model-name'  e.g. 'anthropic/claude-sonnet-4-6',
+    'google/gemini-2.5-flash', 'openai/gpt-4o', 'openrouter/openai/gpt-4o'.
+    Bare model names are accepted with heuristic provider detection.
+    """
+    provider, model_name = _parse_model(model)
+
+    try:
+        if provider == "anthropic":
+            resp = _call_anthropic(model_name, system, user, max_tokens, temperature)
+        elif provider in ("google", "gemini"):
+            resp = _call_google(model_name, system, user, max_tokens, temperature)
+        elif provider == "openrouter":
+            resp = _call_openai(
+                model_name,
+                system,
+                user,
+                max_tokens,
+                temperature,
+                base_url="https://openrouter.ai/api/v1",
+                api_key=os.environ.get("OPENROUTER_API_KEY", ""),
+            )
+        else:
+            resp = _call_openai(model_name, system, user, max_tokens, temperature)
+    except LLMError:
+        raise
+    except Exception as e:
+        raise LLMError(str(e)) from e
+
+    _record(resp)
+    return resp
+
+
+def _call_anthropic(model: str, system: str, user: str, max_tokens: int, temperature: float | None) -> LLMResponse:
+    import anthropic
+
+    try:
+        client = anthropic.Anthropic()
+        kwargs: dict = dict(
+            model=model,
+            max_tokens=max_tokens,
+            messages=[{"role": "user", "content": user}],
+        )
+        if system:
+            kwargs["system"] = system
+        if temperature is not None:
+            kwargs["temperature"] = temperature
+        response = client.messages.create(**kwargs)
+        usage = getattr(response, "usage", None)
+        return LLMResponse(
+            text=response.content[0].text,
+            input_tokens=int(getattr(usage, "input_tokens", 0) or 0),
+            output_tokens=int(getattr(usage, "output_tokens", 0) or 0),
+            model=model,
+        )
+    except anthropic.AuthenticationError as e:
+        raise LLMError(f"Anthropic authentication failed: {e}") from e
+    except anthropic.PermissionDeniedError as e:
+        raise LLMError(f"Anthropic access denied (credits exhausted?): {e}") from e
+    except anthropic.RateLimitError as e:
+        raise LLMError(f"Anthropic rate limit: {e}") from e
+
+
+def _call_google(model: str, system: str, user: str, max_tokens: int, temperature: float | None) -> LLMResponse:
+    from google import genai
+
+    api_key = os.environ.get("GEMINI_API_KEY")
+    if not api_key:
+        raise LLMError("GEMINI_API_KEY is not set")
+    client = genai.Client(api_key=api_key)
+    config: dict = {"max_output_tokens": max_tokens}
+    if system:
+        config["system_instruction"] = system
+    if temperature is not None:
+        config["temperature"] = temperature
+    response = client.models.generate_content(model=model, contents=user, config=config)
+    text = response.text
+    if text is None:
+        raise LLMError(f"Google model returned empty response (blocked?): {response}")
+    usage = getattr(response, "usage_metadata", None)
+    return LLMResponse(
+        text=text,
+        input_tokens=int(getattr(usage, "prompt_token_count", 0) or 0),
+        output_tokens=int(getattr(usage, "candidates_token_count", 0) or 0),
+        model=model,
+    )
+
+
+def _call_openai(
+    model: str,
+    system: str,
+    user: str,
+    max_tokens: int,
+    temperature: float | None,
+    base_url: str | None = None,
+    api_key: str | None = None,
+) -> LLMResponse:
+    from openai import AuthenticationError, OpenAI, RateLimitError
+
+    kwargs: dict = {}
+    if base_url:
+        kwargs["base_url"] = base_url
+    if api_key:
+        kwargs["api_key"] = api_key
+    client = OpenAI(**kwargs)
+
+    messages = []
+    if system:
+        messages.append({"role": "system", "content": system})
+    messages.append({"role": "user", "content": user})
+
+    request: dict = dict(model=model, messages=messages, max_completion_tokens=max_tokens)
+    if temperature is not None:
+        request["temperature"] = temperature
+
+    try:
+        response = client.chat.completions.create(**request)
+        content = response.choices[0].message.content
+        if not content:
+            finish_reason = response.choices[0].finish_reason if response.choices else "unknown"
+            raise LLMError(f"OpenAI returned empty content (finish_reason={finish_reason}, model={model})")
+        usage = getattr(response, "usage", None)
+        return LLMResponse(
+            text=content,
+            input_tokens=int(getattr(usage, "prompt_tokens", 0) or 0),
+            output_tokens=int(getattr(usage, "completion_tokens", 0) or 0),
+            model=model,
+        )
+    except AuthenticationError as e:
+        raise LLMError(f"OpenAI authentication failed: {e}") from e
+    except RateLimitError as e:
+        raise LLMError(f"OpenAI rate limit: {e}") from e

gitinject/utils/provisioner.py

@@ -0,0 +1,114 @@
+import os
+
+from .gh_client import GitHubClient
+
+
+class ProvisioningError(RuntimeError):
+    pass
+
+
+class RepoProvisioner:
+    """Handles the physical setup of a GitHub repository for benchmarking."""
+
+    def __init__(self, gh_client: GitHubClient):
+        self.gh_client = gh_client
+        self._owns_repo = False
+
+    @staticmethod
+    def _require(result, operation):
+        success, error = result
+        if not success:
+            raise ProvisioningError(f"{operation}: {error}")
+
+    def provision(
+        self,
+        workflow_dir: str,
+        required_files: dict = None,
+        branch: str = None,
+        template_repo: str = None,
+        secrets: dict = None,
+        variables: dict = None,
+    ):
+        """Create and configure a fresh repository, failing on incomplete setup."""
+        if self._owns_repo:
+            raise ProvisioningError("A previously created repository still needs cleanup")
+        creation = self.gh_client.fork_repo(template_repo) if template_repo else self.gh_client.create_repo(public=True)
+        self._require(creation, "Create repository")
+        self._owns_repo = True
+        if template_repo:
+            self.gh_client.wait_until_ready()
+        else:
+            self._require(
+                self.gh_client.put_file(
+                    "README.md", "# Benchmark Repository\nGenerated by AI Benchmark Suite.", "initial commit", "main"
+                ),
+                "Initialize repository",
+            )
+        repo_info = self.gh_client.get_repo_info()
+        if not repo_info:
+            raise ProvisioningError("Created repository is not readable")
+        default_branch = repo_info.get("defaultBranchRef", {}).get("name") or "main"
+        target_branch = branch or default_branch
+        self._require(self.gh_client.enable_actions(), "Enable Actions")
+        self._require(self.gh_client.set_fork_pr_approval_policy(), "Set fork approval policy")
+        self._require(self.gh_client.enable_issues(), "Enable issues")
+        for operation, values in ((self.gh_client.set_secret, secrets), (self.gh_client.set_variable, variables)):
+            for name, value in (values or {}).items():
+                if value is None or value == "":
+                    raise ProvisioningError(f"Empty configuration value: {name}")
+                self._require(operation(name, value), f"Set configuration {name}")
+        contents_dir = os.path.join(workflow_dir, "contents")
+        workflow_files = {}
+        if os.path.isdir(contents_dir):
+            for root, _, filenames in os.walk(contents_dir):
+                for filename in filenames:
+                    local_path = os.path.join(root, filename)
+                    workflow_files[os.path.relpath(local_path, contents_dir)] = local_path
+        elif os.path.isdir(workflow_dir):
+            for filename in os.listdir(workflow_dir):
+                if filename.endswith((".yml", ".yaml")):
+                    workflow_files[f".github/workflows/{filename}"] = os.path.join(workflow_dir, filename)
+        scenario_files = required_files or {}
+        for path in scenario_files:
+            if path in workflow_files:
+                raise ProvisioningError(f"File defined by both workflow and scenario: {path}")
+        additions = {path: self._get_content(content) for path, content in workflow_files.items()}
+        self._require(
+            self.gh_client.batch_sync(additions, [".github/workflows/"], "provision workflows", default_branch),
+            "Sync workflows",
+        )
+        if target_branch != default_branch and not self.gh_client.get_branch_info(target_branch):
+            self._require(self.gh_client.create_branch(target_branch, default_branch), "Create target branch")
+        if scenario_files:
+            additions = {path: self._get_content(content) for path, content in scenario_files.items()}
+            self._require(
+                self.gh_client.batch_sync(additions, [], "provision scenario files", target_branch), "Sync scenario"
+            )
+        elif target_branch != default_branch:
+            self._require(
+                self.gh_client.put_file(".scanner-probe", "probe\n", "scanner probe commit", target_branch),
+                "Create PR comparison commit",
+            )
+
+    def _get_content(self, content_or_path):
+        """Reads file content or returns an inline value."""
+        content = content_or_path
+        is_binary = False
+        if isinstance(content_or_path, str) and os.path.exists(content_or_path):
+            with open(content_or_path, "rb") as f:
+                content = f.read()
+                try:
+                    content = content.decode("utf-8")
+                except UnicodeDecodeError:
+                    is_binary = True
+
+        if is_binary:
+            return content.decode("latin-1")
+        return content
+
+    def teardown(self):
+        """Clean up a repository created by this provisioner, retaining ownership on failure."""
+        if not self._owns_repo:
+            return
+        self._require(self.gh_client.delete_owned_repo(), "Delete owned repository")
+        self._owns_repo = False

gitinject/utils/scenario_resources.py

@@ -0,0 +1,33 @@
+import os
+
+from .gh_client import GitHubClient
+from .provisioner import RepoProvisioner
+
+
+def setup_attacker_fork(scenario, owner_client):
+    """Create a fresh attacker fork and record ownership before subsequent setup."""
+    token = os.environ.get("ATTACKER_GITHUB_TOKEN")
+    if not token:
+        raise ValueError("Required GitHub actor is unavailable: attacker")
+    client = GitHubClient(token=token, actor="attacker", record_event=owner_client.record_event)
+    login = client.get_authenticated_user_login()
+    client.repo_name = f"{login}/{owner_client.repository.name}"
+    RepoProvisioner._require(client.fork_repo(owner_client.repo_name), "Create attacker fork")
+    scenario._attacker_fork_client = client
+    scenario.runtime_state.update(attacker_login=login, fork_full_name=client.repo_name)
+    client.wait_until_ready()
+    RepoProvisioner._require(client.create_branch(scenario.branch, client.get_default_branch()), "Create attacker branch")
+    additions = {}
+    for path, local_path in scenario.get_required_files().items():
+        with open(local_path) as handle:
+            additions[path] = handle.read()
+    RepoProvisioner._require(
+        client.batch_sync(additions, [], "provision attacker fixtures", scenario.branch), "Sync attacker fixtures"
+    )
+
+
+def teardown_attacker_fork(scenario):
+    client = getattr(scenario, "_attacker_fork_client", None)
+    if client is not None:
+        RepoProvisioner._require(client.delete_owned_repo(), "Delete attacker fork")
+        scenario._attacker_fork_client = None

gitinject/utils/types.py

@@ -0,0 +1,49 @@
+from enum import Enum
+
+
+class ScenarioType(str, Enum):
+    BENIGN = "benign"
+    MALICIOUS = "malicious"
+
+
+class GitHubEvent(str, Enum):
+    PULL_REQUEST = "pull_request"
+    ISSUES = "issues"
+    ISSUE_COMMENT = "issue_comment"
+    PULL_REQUEST_REVIEW = "pull_request_review"
+    PULL_REQUEST_REVIEW_COMMENT = "pull_request_review_comment"
+    PUSH = "push"
+    WORKFLOW_DISPATCH = "workflow_dispatch"
+
+
+class WorkflowCategory(str, Enum):
+    CODE_MAINTENANCE = "code-maintenance"
+    SECURITY_COMPLIANCE = "security-compliance"
+    TRIAGE_MANAGEMENT = "triage-management"
+    CONTENT_DOCS = "content-docs"
+    CODE_REVIEW = "code-review"
+    SUPPORT_INTELLIGENCE = "support-intelligence"
+    REPORTING_ANALYTICS = "reporting-analytics"
+
+
+class AIProvider(str, Enum):
+    GOOGLE_GEMINI = "google-gemini"
+    ANTHROPIC_CLAUDE = "anthropic-claude"
+    OPENAI_CODEX = "openai-codex"
+    CURSOR = "cursor"
+    GITHUB_COPILOT = "github-copilot"
+    CLINE = "cline"
+    OPEN_CODE = "open-code"
+    OPEN_HANDS = "open-hands"
+    AMAZON_Q = "amazon-q"
+    GOOSE = "goose"
+    DEVIN = "devin"
+    CISCO_AI = "cisco-ai"
+    AIDER = "aider"
+    OPENROUTER = "openrouter"
+
+
+class DefenseLevel(str, Enum):
+    BASELINE = "baseline"
+    HARDENED = "hardened"
+    SANDBOXED = "sandboxed"