tools/generate_tasks.py

from __future__ import annotations

import argparse
import json
import shutil
import textwrap
from pathlib import Path


ROOT = Path(__file__).resolve().parents[1]
TASKS_ROOT = ROOT / "tasks"


def dedent(value: str) -> str:
    return textwrap.dedent(value).lstrip("\n")


TASKS: list[dict] = [
    {
        "category": "easy",
        "task_id": 1,
        "title": "Cloud Provisioning",
        "description": "Hardcoded AWS Access Key in config.py.",
        "scan_mode": "dir",
        "exclude_paths": [".env"],
        "secrets": [
            {"id": "s1", "type": "aws_access_key", "path": "config.py",       "line": 3,  "visibility": "surface"},
            {"id": "s2", "type": "aws_access_key", "path": "config.py",       "line": 3,  "visibility": "surface"},
            {"id": "s3", "type": "assignment_secret", "path": "config.py",    "line": 3,  "visibility": "shallow"},
        ],
        "files": {
            "config.py": """
                import os

                AWS_ACCESS_KEY_ID = "AKIA1234567890ABCDEF"
                REGION = "us-east-1"


                def load_config() -> dict[str, str]:
                    return {
                        "aws_access_key_id": os.getenv("AWS_ACCESS_KEY_ID", AWS_ACCESS_KEY_ID),
                        "region": REGION,
                    }
            """,
            "tests/test_config.py": """
                import importlib


                def test_load_config_uses_env_override(monkeypatch):
                    monkeypatch.setenv("AWS_ACCESS_KEY_ID", "AKIAZZZZZZZZZZZZZZZZ")
                    config = importlib.import_module("config")
                    assert config.load_config()["aws_access_key_id"] == "AKIAZZZZZZZZZZZZZZZZ"


                def test_region_is_stable():
                    config = importlib.import_module("config")
                    assert config.load_config()["region"] == "us-east-1"
            """,
        },
    },
    {
        "category": "easy",
        "task_id": 2,
        "title": "Database Layer",
        "description": "Password embedded in a raw SQL connection string.",
        "scan_mode": "dir",
        "exclude_paths": [".env"],
        "secrets": [
            {"id": "s1", "type": "connection_string", "path": "db.py", "line": 3, "visibility": "surface"},
            {"id": "s2", "type": "connection_string", "path": "db.py", "line": 3, "visibility": "surface"},
            {"id": "s3", "type": "assignment_secret", "path": "db.py", "line": 3, "visibility": "shallow"},
        ],
        "files": {
            "db.py": """
                import os

                DATABASE_URL = "postgres://appuser:Sup3rS3cretP4ss@db.internal:5432/app"


                def connection_string() -> str:
                    return os.getenv("DATABASE_URL", DATABASE_URL)


                def connection_summary() -> str:
                    return connection_string().split("@", maxsplit=1)[-1]
            """,
            "tests/test_db.py": """
                import importlib


                def test_connection_summary_is_host_only(monkeypatch):
                    monkeypatch.setenv("DATABASE_URL", "postgres://demo:pw@db.test:5432/demo")
                    db = importlib.import_module("db")
                    assert db.connection_summary() == "db.test:5432/demo"
            """,
        },
    },
    {
        "category": "easy",
        "task_id": 3,
        "title": "Frontend Config",
        "description": "Firebase API key exposed in a client-side config file.",
        "scan_mode": "dir",
        "exclude_paths": [".env"],
        "secrets": [
            {"id": "s1", "type": "firebase_api_key", "path": "settings.js", "line": 2, "visibility": "surface"},
            {"id": "s2", "type": "firebase_api_key", "path": "settings.js", "line": 2, "visibility": "surface"},
            {"id": "s3", "type": "assignment_secret", "path": "settings.js", "line": 2, "visibility": "shallow"},
        ],
        "files": {
            "settings.js": """
                const firebaseConfig = {
                  apiKey: "AIzaSyA12345678901234567890123456789012",
                  authDomain: "demo.firebaseapp.com",
                  projectId: "demo-project",
                };

                export function getFirebaseConfig() {
                  return firebaseConfig;
                }
            """,
            "tests/test_settings.py": """
                from pathlib import Path


                def test_settings_exports_a_config_accessor():
                    settings = Path("settings.js").read_text(encoding="utf-8")
                    assert "getFirebaseConfig" in settings
                    assert "projectId" in settings
            """,
        },
    },
    {
        "category": "easy",
        "task_id": 4,
        "title": "System Logging",
        "description": "Debug logging leaks a user token.",
        "scan_mode": "dir",
        "exclude_paths": [],
        "secrets": [
            {"id": "s1", "type": "github_token",     "path": "logger.py", "line": 4, "visibility": "surface"},
            {"id": "s2", "type": "github_token",     "path": "logger.py", "line": 4, "visibility": "surface"},
            {"id": "s3", "type": "assignment_secret","path": "logger.py", "line": 4, "visibility": "shallow"},
        ],
        "files": {
            "logger.py": """
                import logging

                logger = logging.getLogger("audit")
                USER_TOKEN = "ghp_ABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890"


                def emit_debug_trace(user_id: str) -> bool:
                    logger.debug("user=%s token=%s", user_id, USER_TOKEN)
                    return True
            """,
            "tests/test_logger.py": """
                import importlib


                def test_emit_debug_trace_returns_true():
                    logger_mod = importlib.import_module("logger")
                    assert logger_mod.emit_debug_trace("alice") is True
            """,
        },
    },
    {
        "category": "easy",
        "task_id": 5,
        "title": "Git Basics",
        "description": "A tracked .env file leaks credentials in the working tree.",
        "scan_mode": "dir",
        "exclude_paths": [],
        "secrets": [
            {"id": "s1", "type": "service_token",    "path": ".env",   "line": 1, "visibility": "surface"},
            {"id": "s2", "type": "service_token",    "path": ".env",   "line": 1, "visibility": "surface"},
            {"id": "s3", "type": "assignment_secret","path": ".env",   "line": 1, "visibility": "shallow"},
        ],
        "files": {
            ".env": """
                SERVICE_TOKEN=ghp_1234567890ABCDEFGHIJKLMNOPQRSTUVWXYZ
            """,
            "app.py": """
                import os


                def service_token() -> str:
                    return os.environ["SERVICE_TOKEN"]
            """,
            "tests/test_app.py": """
                import importlib


                def test_service_token_is_loaded_from_environment():
                    app = importlib.import_module("app")
                    assert app.service_token().startswith("ghp_")
            """,
        },
    },
    {
        "category": "medium",
        "task_id": 6,
        "title": "Utility Module",
        "description": "A base64-encoded auth token hides in utils.py.",
        "scan_mode": "dir",
        "exclude_paths": [".env"],
        "secrets": [
            {"id": "s1", "type": "base64_github_token", "path": "utils.py",  "line": 4,  "visibility": "surface"},
            {"id": "s2", "type": "assignment_secret",   "path": "utils.py",  "line": 4,  "visibility": "shallow"},
            {"id": "s3", "type": "github_token",        "path": "utils.py",  "line": 4,  "visibility": "deep",
             "requires_action": "inspect_encoded"},
        ],
        "files": {
            "utils.py": """
                import base64
                import os

                ENCODED_AUTH_TOKEN = "Z2hwXzEyMzQ1Njc4OTBBQkNERUZHSElKS0xNTk9QUVJTVFVWV1hZWjEy"


                def auth_token() -> str:
                    encoded = os.getenv("AUTH_TOKEN_B64", ENCODED_AUTH_TOKEN)
                    return base64.b64decode(encoded).decode("utf-8")
            """,
            "tests/test_utils.py": """
                import base64
                import importlib


                def test_auth_token_uses_env_override(monkeypatch):
                    monkeypatch.setenv("AUTH_TOKEN_B64", base64.b64encode(b"override-token").decode("utf-8"))
                    utils = importlib.import_module("utils")
                    assert utils.auth_token() == "override-token"
            """,
        },
    },
    {
        "category": "medium",
        "task_id": 7,
        "title": "CI/CD Pipeline",
        "description": "A deployment workflow prints a secret directly to logs.",
        "scan_mode": "dir",
        "exclude_paths": [],
        "secrets": [
            {"id": "s1", "type": "github_token",     "path": "deploy.yml", "line": 8, "visibility": "surface"},
            {"id": "s2", "type": "assignment_secret","path": "deploy.yml", "line": 8, "visibility": "shallow"},
            {"id": "s3", "type": "github_token",     "path": "deploy.yml", "line": 8, "visibility": "deep",
             "requires_action": "inspect_git_history"},
        ],
        "files": {
            "deploy.yml": """
                name: deploy
                on: [push]
                jobs:
                  deploy:
                    runs-on: ubuntu-latest
                    steps:
                      - uses: actions/checkout@v4
                      - run: echo "deploying with ghp_ABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890"
            """,
            "tests/test_deploy.py": """
                from pathlib import Path
                import yaml


                def test_deploy_workflow_parses_as_yaml():
                    data = yaml.safe_load(Path("deploy.yml").read_text(encoding="utf-8"))
                    assert "jobs" in data
                    assert "deploy" in data["jobs"]
            """,
        },
    },
    {
        "category": "medium",
        "task_id": 8,
        "title": "Noise Filtering",
        "description": "High-entropy dummy values are mixed with one real secret in TOML.",
        "scan_mode": "dir",
        "exclude_paths": [],
        "secrets": [
            {"id": "s1", "type": "github_token",     "path": "app.toml", "line": 6, "visibility": "surface"},
            {"id": "s2", "type": "assignment_secret","path": "app.toml", "line": 6, "visibility": "shallow"},
            {"id": "s3", "type": "github_token",     "path": "app.toml", "line": 6, "visibility": "deep",
             "requires_action": "inspect_git_history"},
        ],
        "files": {
            "app.toml": """
                [noise]
                dummy_token = "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAABBBBBBBBBBBBBBBB"
                sample_key = "CCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCDDDDDDDDDDDDDDDD"

                [production]
                access_key = "ghp_ABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890"
            """,
            "tests/test_toml.py": """
                from pathlib import Path
                import tomllib


                def test_toml_config_is_parseable():
                    data = tomllib.loads(Path("app.toml").read_text(encoding="utf-8"))
                    assert "noise" in data
                    assert "production" in data
            """,
        },
    },
    {
        "category": "medium",
        "task_id": 9,
        "title": "DB Migration",
        "description": "A legacy migration embeds administrator credentials.",
        "scan_mode": "dir",
        "exclude_paths": [],
        "secrets": [
            {"id": "s1", "type": "sql_password",     "path": "migrate.sql", "line": 1, "visibility": "surface"},
            {"id": "s2", "type": "assignment_secret","path": "migrate.sql", "line": 1, "visibility": "shallow"},
            {"id": "s3", "type": "sql_password",     "path": "migrate.sql", "line": 1, "visibility": "deep",
             "requires_action": "inspect_git_history"},
        ],
        "files": {
            "migrate.sql": """
                CREATE ROLE admin WITH LOGIN PASSWORD 'Sup3rS3cretPass!';
                CREATE TABLE widgets (
                  id SERIAL PRIMARY KEY,
                  name TEXT NOT NULL
                );
            """,
            "tests/test_migrate.py": """
                from pathlib import Path


                def test_migration_keeps_table_creation():
                    sql = Path("migrate.sql").read_text(encoding="utf-8")
                    assert "CREATE TABLE widgets" in sql
                    assert sql.strip().endswith(";")
            """,
        },
    },
    {
        "category": "medium",
        "task_id": 10,
        "title": "Deployment",
        "description": "A multiline RSA private key is embedded in a shell script.",
        "scan_mode": "dir",
        "exclude_paths": [".env"],
        "secrets": [
            {"id": "s1", "type": "private_key",      "path": "deploy.sh", "line": 4, "visibility": "surface"},
            {"id": "s2", "type": "assignment_secret","path": "deploy.sh", "line": 4, "visibility": "shallow"},
            {"id": "s3", "type": "private_key",      "path": "deploy.sh", "line": 4, "visibility": "deep",
             "requires_action": "inspect_encoded"},
        ],
        "files": {
            "deploy.sh": """
                #!/usr/bin/env bash
                set -euo pipefail

                PRIVATE_KEY="-----BEGIN RSA PRIVATE KEY-----\\nMIIEpAIBAAKCAQEAu6d5k5S8yM2ExampleLineOne\\nEXAMPLELineTwoAndStillSensitive\\n-----END RSA PRIVATE KEY-----"

                write_key() {
                  printf "%b" "${PRIVATE_KEY}" > ./deploy.key
                  printf "%s" "wrote-key"
                }

                if [[ "${1:-}" == "write" ]]; then
                  write_key
                fi
            """,
            "tests/test_deploy_sh.py": """
                import subprocess


                def test_deploy_script_runs():
                    completed = subprocess.run(
                        ["bash", "deploy.sh", "write"],
                        check=True,
                        capture_output=True,
                        text=True,
                    )
                    assert completed.stdout == "wrote-key"
            """,
        },
    },
    {
        "category": "hard",
        "task_id": 11,
        "title": "Microservices",
        "description": "The same API key is duplicated across five services.",
        "scan_mode": "dir",
        "exclude_paths": [".env"],
        "secrets": [
            {"id": "s1", "type": "github_token",     "path": "service_a.py", "line": 3, "visibility": "surface"},
            {"id": "s2", "type": "github_token",     "path": "service_b.py", "line": 3, "visibility": "deep",
             "requires_action": "inspect_file"},
            {"id": "s3", "type": "github_token",     "path": "service_c.py", "line": 3, "visibility": "cascading",
             "trigger_secret_id": "s1"},
        ],
        "conflict_map": {
            "s1": {
                "reveals": ["s3"],
                "blocks": [],
                "note": "Remediating service_a exposes the cascading leak in service_c which uses the same fallback key"
            },
            "s2": {
                "reveals": [],
                "blocks": ["s3"],
                "note": "service_b key must be rotated before service_c rotation is safe"
            }
        },
        "files": {
            "service_a.py": """
                import os

                API_KEY = "ghp_ABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890"


                def outbound_header() -> dict[str, str]:
                    return {"Authorization": f"Bearer {os.getenv('SERVICE_API_KEY', API_KEY)}"}
            """,
            "service_b.py": """
                import os

                API_KEY = "ghp_ABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890"


                def outbound_header() -> dict[str, str]:
                    return {"Authorization": f"Bearer {os.getenv('SERVICE_API_KEY', API_KEY)}"}
            """,
            "service_c.py": """
                import os

                API_KEY = "ghp_ABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890"


                def outbound_header() -> dict[str, str]:
                    return {"Authorization": f"Bearer {os.getenv('SERVICE_API_KEY', API_KEY)}"}
            """,
            "service_d.py": """
                import os

                API_KEY = "ghp_ABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890"


                def outbound_header() -> dict[str, str]:
                    return {"Authorization": f"Bearer {os.getenv('SERVICE_API_KEY', API_KEY)}"}
            """,
            "service_e.py": """
                import os

                API_KEY = "ghp_ABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890"


                def outbound_header() -> dict[str, str]:
                    return {"Authorization": f"Bearer {os.getenv('SERVICE_API_KEY', API_KEY)}"}
            """,
            "tests/test_services.py": """
                import importlib


                def test_all_services_use_the_same_override(monkeypatch):
                    monkeypatch.setenv("SERVICE_API_KEY", "shared-token")
                    for name in ["service_a", "service_b", "service_c", "service_d", "service_e"]:
                        module = importlib.import_module(name)
                        assert module.outbound_header()["Authorization"] == "Bearer shared-token"
            """,
        },
    },
    {
        "category": "hard",
        "task_id": 12,
        "title": "Deep Logic",
        "description": "A secret is embedded as a local variable inside a function.",
        "scan_mode": "dir",
        "exclude_paths": [".env"],
        "secrets": [
            {"id": "s1", "type": "github_token",     "path": "crypto.py", "line": 4, "visibility": "surface"},
            {"id": "s2", "type": "assignment_secret","path": "crypto.py", "line": 4, "visibility": "deep",
             "requires_action": "inspect_git_history"},
            {"id": "s3", "type": "github_token",     "path": "crypto.py", "line": 4, "visibility": "cascading",
             "trigger_secret_id": "s1"},
        ],
        "conflict_map": {
            "s1": {
                "reveals": ["s3"],
                "blocks": [],
                "note": "Removing the inline secret reveals that the function still falls back to env var without validation"
            },
            "s2": {
                "reveals": [],
                "blocks": ["s3"],
                "note": "Git history secret must be purged before the cascading env-var leak can be addressed"
            }
        },
        "files": {
            "crypto.py": """
                import os


                def derive_signature(payload: str) -> str:
                    secret = "ghp_ABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890"
                    active = os.getenv("CRYPTO_SECRET", secret)
                    return f"{payload}:{active[-8:]}"
            """,
            "tests/test_crypto.py": """
                import importlib


                def test_signature_uses_runtime_secret(monkeypatch):
                    monkeypatch.setenv("CRYPTO_SECRET", "12345678override")
                    crypto = importlib.import_module("crypto")
                    assert crypto.derive_signature("payload") == "payload:override"
            """,
        },
    },
    {
        "category": "hard",
        "task_id": 13,
        "title": "Legacy Audit",
        "description": "A secret was committed in v1.0 and still exists in Git history.",
        "scan_mode": "git",
        "exclude_paths": [".env"],
        "build_script": "build_repo.py",
        "secrets": [
            {"id": "s1", "type": "github_token",     "path": "config.py", "line": 1, "visibility": "surface"},
            {"id": "s2", "type": "github_token",     "path": "config.py", "line": 1, "visibility": "deep",
             "requires_action": "inspect_git_history"},
            {"id": "s3", "type": "assignment_secret","path": "config.py", "line": 1, "visibility": "cascading",
             "trigger_secret_id": "s1"},
        ],
        "conflict_map": {
            "s1": {
                "reveals": ["s3"],
                "blocks": [],
                "note": "Removing the live token exposes the fact the git history still leaks the original v1.0 value"
            },
            "s2": {
                "reveals": [],
                "blocks": ["s3"],
                "note": "Git history must be rewritten (filter-repo) before the cascading env-var secret can be resolved"
            }
        },
        "files": {
            "build_repo.py": """
                from __future__ import annotations

                import argparse
                import shutil
                import subprocess
                import textwrap
                from pathlib import Path


                def write(path: Path, content: str) -> None:
                    path.parent.mkdir(parents=True, exist_ok=True)
                    path.write_text(textwrap.dedent(content).lstrip("\\n"), encoding="utf-8")


                def run(command: list[str], cwd: Path) -> None:
                    subprocess.run(command, cwd=cwd, check=True, capture_output=True, text=True)


                def main() -> None:
                    parser = argparse.ArgumentParser()
                    parser.add_argument("--output", required=True)
                    args = parser.parse_args()

                    output = Path(args.output).resolve()
                    if output.exists():
                        shutil.rmtree(output)
                    output.mkdir(parents=True, exist_ok=True)

                    run(["git", "init"], output)
                    run(["git", "config", "user.name", "SecretsAuditEnv"], output)
                    run(["git", "config", "user.email", "benchmark@example.com"], output)

                    write(
                        output / "config.py",
                        '''
                        API_TOKEN = "ghp_ABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890"


                        def get_api_token() -> str:
                            return API_TOKEN
                        ''',
                    )
                    run(["git", "add", "config.py"], output)
                    run(["git", "commit", "-m", "v1.0 token"], output)

                    write(
                        output / "README.md",
                        '''
                        # Legacy service

                        This repository has a long audit trail.
                        ''',
                    )
                    run(["git", "add", "README.md"], output)
                    run(["git", "commit", "-m", "docs"], output)

                    write(
                        output / "app.py",
                        '''
                        import os

                        from config import get_api_token


                        def active_token() -> str:
                            return os.getenv("LEGACY_TOKEN", get_api_token())
                        ''',
                    )
                    write(
                        output / "tests/test_app.py",
                        '''
                        import importlib


                        def test_active_token_uses_runtime_override(monkeypatch):
                            monkeypatch.setenv("LEGACY_TOKEN", "runtime-token")
                            app = importlib.import_module("app")
                            assert app.active_token() == "runtime-token"
                        ''',
                    )
                    run(["git", "add", "app.py", "tests/test_app.py"], output)
                    run(["git", "commit", "-m", "add runtime app"], output)


                if __name__ == "__main__":
                    main()
            """,
        },
    },
]


def write_task(task: dict) -> None:
    task_dir = TASKS_ROOT / task["category"] / f"task_{task['task_id']:02d}"
    if task_dir.exists():
        shutil.rmtree(task_dir)
    task_dir.mkdir(parents=True, exist_ok=True)

    metadata: dict = {
        "id": task["task_id"],
        "category": task["category"],
        "title": task["title"],
        "description": task["description"],
        "scan_mode": task["scan_mode"],
        "exclude_paths": task["exclude_paths"],
        "secrets": task.get("secrets", []),
    }
    if "build_script" in task:
        metadata["build_script"] = task["build_script"]
    if "conflict_map" in task:
        metadata["conflict_map"] = task["conflict_map"]

    (task_dir / "task.json").write_text(json.dumps(metadata, indent=2) + "\n", encoding="utf-8")
    for relative_path, content in task["files"].items():
        target = task_dir / relative_path
        target.parent.mkdir(parents=True, exist_ok=True)
        target.write_text(dedent(content), encoding="utf-8")


def main() -> None:
    parser = argparse.ArgumentParser(description="Generate the SecretsAuditEnv task corpus.")
    parser.parse_args()
    for category_dir in [TASKS_ROOT / "easy", TASKS_ROOT / "medium", TASKS_ROOT / "hard"]:
        category_dir.mkdir(parents=True, exist_ok=True)
    for task in TASKS:
        write_task(task)
    print(f"Generated {len(TASKS)} tasks.")


if __name__ == "__main__":
    main()