Metadata-Version: 2.4
Name: wizsec
Version: 1.1.0
Summary: A custom SDK for the Wiz API
Author-email: James Husted <james@husted.dev>
License-Expression: MIT
Project-URL: homepage, https://github.com/HusteDev/wizsec
Project-URL: repository, https://github.com/HusteDev/wizsec
Keywords: Wiz,SDK,API,Security
Classifier: Programming Language :: Python :: 3
Classifier: Operating System :: OS Independent
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: httpx
Requires-Dist: pyrate-limiter<5.0,>=4.0
Requires-Dist: python-dotenv
Requires-Dist: PyYAML
Requires-Dist: graphql-core
Provides-Extra: dev
Requires-Dist: pytest>=7.0; extra == "dev"
Requires-Dist: pytest-asyncio; extra == "dev"
Requires-Dist: black; extra == "dev"
Requires-Dist: mypy; extra == "dev"
Requires-Dist: types-PyYAML; extra == "dev"
Provides-Extra: docs
Requires-Dist: mkdocs-material; extra == "docs"
Requires-Dist: mkdocstrings[python]; extra == "docs"
Dynamic: license-file

# wizsec

[![CI](https://github.com/HusteDev/wizsec/actions/workflows/ci.yml/badge.svg)](https://github.com/HusteDev/wizsec/actions/workflows/ci.yml)
[![CodeQL](https://github.com/HusteDev/wizsec/actions/workflows/codeql.yml/badge.svg)](https://github.com/HusteDev/wizsec/actions/workflows/codeql.yml)
[![PyPI version](https://img.shields.io/pypi/v/wizsec)](https://pypi.org/project/wizsec/)
[![Python versions](https://img.shields.io/pypi/pyversions/wizsec)](https://pypi.org/project/wizsec/)
[![License](https://img.shields.io/pypi/l/wizsec)](https://github.com/HusteDev/wizsec/blob/main/LICENSE)
[![codecov](https://codecov.io/gh/HusteDev/wizsec/graph/badge.svg)](https://codecov.io/gh/HusteDev/wizsec)
[![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)

A Python SDK for the [Wiz](https://www.wiz.io/) Cloud Security GraphQL API. Provides sync and async clients with automatic pagination, rate limiting, batch operations, and report generation.

## Features

- **Unified HTTP transport** via [httpx](https://www.python-httpx.org/) (sync and async)
- **Automatic cursor-based pagination** with result merging
- **Per-environment rate limiting** using `pyrate-limiter` (respects Wiz's global rate limits)
- **Batch requests** — submit multiple queries concurrently (sync threads or async tasks)
- **Report generation** — create, poll, stream, and download Wiz reports (JSON and CSV)
- **Multiple auth flows** — client credentials and device code (OAuth)
- **Flexible credential storage** — environment variables, credential files, or interactive prompt
- **Multi-environment / multi-profile** — connect to enabled `app`, `gov`, or `fedramp` environments with isolated credential profiles
- **Serverless support** — optimized for AWS Lambda and similar environments
- **YAML configuration** via `~/.wiz/wiz.config`
- **Client-side schema validation** — catch query typos before they hit the API
- **Custom query libraries** — resolve query names from importable Python modules
- **PEP 561 typed** (`py.typed` marker included)

## Installation

Install from [Pypi](https://pypi.org/project/wizsec/):

```bash
pip install wizsec
```

Install from source:

```bash
git clone https://github.com/HusteDev/wizsec.git
cd wizsec
pip install .
```

For development:

```bash
pip install -e ".[dev]"
```

## Requirements

- Python >= 3.9
- `httpx`, `pyrate-limiter`, `python-dotenv`, `PyYAML`, `graphql-core`

## Quick Start

### Authentication Setup

The SDK supports two OAuth grant types:

| Grant Type | Use Case | Requires |
|---|---|---|
| `client_credentials` | Service accounts, automation, CI/CD | Client ID + Secret |
| `device_code` | Interactive / user-based sessions | Browser + WizCode license |

#### Client Credentials (default)

Provide your client ID and secret via environment variables, a credentials file, or constructor arguments.

**Environment variables (simplest):**

```bash
export WIZ_CLIENT_ID="your-client-id"
export WIZ_CLIENT_SECRET="your-client-secret"
```

**Credentials file** at `~/.wiz/wiz.credentials`:

```ini
[default]
client_id = your-client-id
client_secret = your-client-secret
environment = app
```

**Pass directly:**

```python
from wizsec import WizClient, Config

Config.load()
client = WizClient(client_id="...", client_secret="...")
```

#### Device Code (Interactive)

Device code auth opens a browser for the user to authorize the session — no client secret needed. This is ideal for CLI tools, notebooks, or any context where a human is present.

Set the grant type in `~/.wiz/wiz.config`:

```yaml
auth:
  grant_type: device_code
  device:
    quiet: true       # auto-authorize without extra prompts (default: true)
    poll_time: 5       # seconds between auth status checks (default: 5)
```

Then use the client normally — the browser will open automatically:

```python
from wizsec import WizClient, Config

Config.load()
client = WizClient(environment="app")  # opens browser for authorization
result = client.create_request(query="...", vars={}).submit()
```

The SDK polls the auth endpoint until the user completes authorization or the request times out.

### Your First Query

```python
from wizsec import WizClient, Config

Config.load()
client = WizClient(environment="app")

response = client.create_request(
    query='{ users(first: 10) { nodes { name email } pageInfo { hasNextPage endCursor } } }',
    vars={}
)
result = response.submit()

if result.success:
    print(result.data)
else:
    print(result.errors)
```

## Usage

### Single Queries

```python
response = client.create_request(query="...", vars={"first": 100})
result = response.submit()
```

Pagination is handled automatically — the SDK detects queries using the [Relay connection pattern](https://relay.dev/graphql/connections.htm) (`nodes` + `pageInfo { hasNextPage endCursor }`) and injects the `$after` cursor variable for you. Results from all pages are merged into `result.data`.

You don't need to declare `$after` in your query — the SDK adds it when:
- The operation is a **query** (not a mutation)
- The query selects both `nodes` and `pageInfo` subfields
- `$after` isn't already declared

Pagination defaults to `api.auto_paginate` in `~/.wiz/wiz.config`. If you set `paginate=False`, no injection occurs and only the first page is returned.

### Query Collections

Organize reusable GraphQL queries in a Python module and reference them by name. This keeps queries out of your application logic and makes them shareable across scripts.

**`queries.py`** — define queries as module-level constants:

```python
ListUsers = """
    query ListUsers($first: Int) {
        users(first: $first) {
            nodes { id name email role }
            pageInfo { hasNextPage endCursor }
        }
    }
"""

GetProject = """
    query GetProject($id: ID!) {
        project(id: $id) { id name slug riskProfile { businessImpact } }
    }
"""
```

**`main.py`** — resolve by name or pass the string directly:

```python
import queries

# Resolve by attribute name from the collection
response = client.create_request(
    queryCollection=queries,
    query="ListUsers",  # resolves to queries.ListUsers
    vars={"first": 50}
)

# Or pass the query string directly (no collection needed)
response = client.create_request(
    query=queries.GetProject,
    vars={"id": "some-project-id"},
    paginate=False,
)
```

`queryCollection` accepts a module, a module name string (auto-imported), or any object whose attributes are the query strings — for example a `types.SimpleNamespace` for one-off scripts that don't warrant a separate `queries.py`:

```python
from types import SimpleNamespace

queries = SimpleNamespace(
    ListUsers="query ListUsers($first: Int) { users(first: $first) { nodes { id } } }",
)

response = client.create_request(queryCollection=queries, query="ListUsers", vars={"first": 50})
```

See [`examples/query_collection/`](examples/query_collection/) for a complete working example.

### Batch Requests (Sync)

```python
batch = client.create_batch_request()
batch.add_request(query="...", vars={"type": "VM"})
batch.add_request(query="...", vars={"type": "CONTAINER"})

batch.set_progress_callback(lambda done, total: print(f"{done}/{total}"))
results = batch.submit(max_concurrent=5)

print(f"{results.success_count()}/{results.total_count()} succeeded")

for request_id, response in results:
    if response.success:
        print(response.data)
```

### Async Requests

```python
import asyncio
from wizsec import WizClient, Config

async def main():
    Config.load()
    client = WizClient(environment="app")

    async with client.async_session() as async_client:
        response = await async_client.create_async_request(
            query="...", vars={"first": 100}
        )
        result = await response.submit()
        print(result.data)

asyncio.run(main())
```

### Async Batch Requests

```python
async with client.async_session() as async_client:
    batch = await async_client.create_async_batch_request()
    batch.add_request(query="...", vars={"type": "VM"})
    batch.add_request(query="...", vars={"type": "CONTAINER"})

    results = await batch.submit(max_concurrent=50)
    print(results.success_rate())
```

Async batch responses preserve one result per submitted request. Failed tasks remain indexed in the batch response with their errors attached.

### Sync vs Async: When to Use Each

The SDK provides both synchronous and asynchronous interfaces. Choose based on your use case:

| Approach  | Best For                                                                 |
| --------- | ------------------------------------------------------------------------ |
| **Sync**  | Simple scripts, single queries, CLI tools, quick prototypes              |
| **Async** | Multiple independent queries, high-throughput applications, web services |

**Performance comparison** (3 queries fetching Projects, Users, and Service Accounts):

```
SYNC (sequential)    2.57s   — queries run one after another
ASYNC (concurrent)   0.60s   — queries run in parallel
```

Async achieves ~4x speedup here because all three API calls happen concurrently instead of waiting for each to complete.

**Use sync when:**

- Running a single query or a few dependent queries
- Writing simple scripts or one-off tools
- Code simplicity matters more than throughput

**Use async when:**

- Fetching data from multiple independent queries
- Building web applications or services that need to handle concurrent requests
- Performance is critical and queries don't depend on each other
- Working with `asyncio`-based frameworks (FastAPI, aiohttp, etc.)

See [`examples/sync_vs_async.py`](examples/sync_vs_async.py) for a runnable comparison.

### Report Generation

```python
response = client.create_request(
    query="mutation { createReport(...) { report { id } } }",
    report_request={"name": "my-report", "stream": True}
)
result = response.submit()

# Report data is automatically polled, downloaded, and attached:
report_rows = result.data.get("report_data", [])
```

### Progress Tracking

Monitor pagination progress with a callback that fires after each page is fetched. The `on_page_event` callback receives a dict with:

| Key | Type | Description |
|-----|------|-------------|
| `page_data` | `dict` | Raw GraphQL data from the current page |
| `page_info` | `dict` | `{"page": int, "per_page": int}` — current page number and page size |
| `errors` | `list` | Any errors accumulated so far |

**Simple progress logging:**

```python
def on_page(event):
    page = event["page_info"]["page"]
    per_page = event["page_info"]["per_page"]
    key = next(iter(event["page_data"]), None)
    count = len(event["page_data"][key]["nodes"]) if key else 0
    print(f"Page {page}: received {count}/{per_page} items")

response = client.create_request(
    query="...",
    vars={"first": 500},
    on_page_event=on_page
)
result = response.submit()
```

**Spinner with live counter** (runs the query on the background thread while animating in the main thread):

```python
import sys, time, threading

progress = {"pages": 0, "items": 0, "done": False}

def on_page(event):
    progress["pages"] = event["page_info"]["page"]
    key = next(iter(event["page_data"]), None)
    if key:
        progress["items"] += len(event["page_data"][key].get("nodes", []))

result_holder = {}

def run_query():
    result_holder["result"] = client.create_request(
        query="...", vars={"first": 100}, on_page_event=on_page
    ).submit()
    progress["done"] = True

thread = threading.Thread(target=run_query)
thread.start()

spinner = "|/-\\"
i = 0
while not progress["done"]:
    sys.stdout.write(f"\r  {spinner[i % 4]} page {progress['pages']}, {progress['items']} items")
    sys.stdout.flush()
    i += 1
    time.sleep(0.15)
thread.join()
```

Progress tracking also works with async requests and report streaming. For reports, the callback receives `{"name", "total_size", "downloaded", "status"}` instead of page data.

See [`examples/progress_tracking.py`](examples/progress_tracking.py) for complete sync, spinner, and async examples.

### Schema Validation

Validate GraphQL queries against the Wiz schema before they hit the API. Catches typos and invalid fields early with helpful suggestions:

```python
from wizsec import WizClient, Config, SchemaValidator, WizSchemaValidationError

Config.load()
Config.set("api", "validate_queries", value=True)  # or set in wiz.config

client = WizClient()

# Typos are caught before the request is sent
try:
    client.create_request(query="query { projectz { nodes { id } } }")
except WizSchemaValidationError as e:
    print(e.validation_errors[0])
    # "Cannot query field 'projectz' on type 'Query'. Did you mean 'project', 'projects', or 'projectTags'?"
```

Validate queries programmatically without creating a request:

```python
try:
    SchemaValidator.validate_query("query { fakeEndpoint { data } }", "app")
except WizSchemaValidationError as e:
    print(e.validation_errors[0])
    # "Cannot query field 'fakeEndpoint' on type 'Query'. Did you mean 'apiEndpoint'?"
```

The schema is cached under the configured `app.wiz_dir` as `schema_<env>.json` and reloaded automatically.

See [`examples/schema_validation.py`](examples/schema_validation.py) for more examples.

## Streaming Pagination

Auto-pagination aggregates every page in memory before returning. For large result sets, iterate instead — `iterate_nodes` fetches pages lazily and never holds more than one page:

```python
for issue in client.iterate_nodes(query=ISSUES_QUERY, vars={"first": 500}):
    process(issue)          # breaking out early stops further API calls

async with client.async_session() as ac:
    async for issue in ac.iterate_nodes_async(query=ISSUES_QUERY, page_size=500):
        process(issue)
```

Each page goes through the normal rate-limiting and retry pipeline; a failed page raises the typed error (`WizAPIError`, `WizRateLimitError`, …).

## Rate Limiting

The SDK automatically enforces Wiz's API rate limits so you don't have to think about throttling. Rate limiters are shared across all `WizClient` instances on the same environment, even across different profiles.

Limits are applied per request type (query vs. mutation) and account type (user vs. service account), based on [Wiz's published rate limits](https://docs.wiz.io/wiz-docs/docs/rate-limiting).

By default the SDK runs at **80% of the published limits**. The headroom keeps normal operation clear of server-side `429`s — the server's quota is per tenant, so network jitter and other consumers (other scripts, CI jobs, integrations) draw from the same budget. Both the headroom and the absolute rates are configurable:

```yaml
rate_limit:
  headroom: 0.8            # fraction of published limits to use (default 0.8)
  overrides:               # optional absolute requests-per-second per limiter key
    query_service: 8       # wins over headroom for that key
```

Local limiter waits and server `429` backoffs are handled automatically and never consume retry attempts. Rate-limit responses reported inside a `200` response's GraphQL `errors` list are detected and backed off the same way. Async clients on the same environment share the same server backoff window.

Note that the local limiter coordinates requests within a single Python process. Separate processes each apply their own budget, so if you run several wizsec processes in parallel against one tenant, lower the headroom (or set explicit overrides) accordingly.

## CLI

Installing the package registers a `wizsec` command for managing configuration and credentials:

```bash
wizsec config init            # create a default config file (~/.wiz/wiz.config)
wizsec config show            # print the current config
wizsec config get api.timeout
wizsec config set rate_limit.headroom 0.7
wizsec config unset query_splitting.enabled

wizsec creds set --profile default --environment gov   # prompts for ID/secret
wizsec creds list             # profiles with masked IDs (secrets never shown)
wizsec creds remove staging
wizsec creds test             # authenticate with stored credentials

wizsec doctor                 # diagnose config, credentials, rate budgets, schema caches
wizsec doctor --auth          # same, plus a live authentication check
```

All commands accept `--file` to operate on a non-default location. `config set` rewrites the YAML file, which drops comments — see the bundled template for documentation of every option.

## Configuration

The SDK reads `~/.wiz/wiz.config` (YAML). Example:

```yaml
app:
  name: wizsec
  release: "1.1.0"
  config_schema: 2

auth:
  grant_type: client_credentials
  credentials:
    storage_method: file
    file_path: ~/.wiz/
  proxy:
    http:
      url: ""
      port: 80
    https:
      url: ""
      port: 80

domain:
  default: gov
  app:
    enabled: false
  gov:
    enabled: true
  fedramp:
    enabled: false

api:
  timeout: 60
  max_retries: 3
  retry_time: 2
  auto_paginate: true
  validate_queries: false

logging:
  enabled: false
  console_handler:
    enabled: true
    logging_level: INFO
```

Blank proxy URLs use environment proxy variables. Config can also be set via `Config.load(overrides=["api.timeout=120"])`.

## Multi-Environment & Multi-Profile

```python
# Different Wiz tenants
app_client = WizClient(environment="app")
gov_client = WizClient(environment="gov")

# Different credential profiles on the same tenant
admin = WizClient(environment="app", profile="admin")
readonly = WizClient(environment="app", profile="readonly")
```

Clients sharing the same environment automatically share a single request queue and rate limiter.
Requested environments must be enabled under `domain.<environment>.enabled`. Auth state is isolated by `(environment, profile)`, so the same profile name can be used safely against different environments.

## Serverless / Lambda

Set `WIZ_SERVERLESS=1` or deploy to an environment with `AWS_LAMBDA_FUNCTION_NAME` set. The SDK adapts automatically:

- Disables background worker threads (executes inline)
- Reads config from `/var/task/.wiz/`
- Call `client.cleanup_for_lambda()` at the end of each invocation

```python
def handler(event, context):
    Config.load()
    client = WizClient(environment="app", serverless=True)
    try:
        result = client.create_request(query="...", vars={}).submit()
        return result.data
    finally:
        client.cleanup_for_lambda()
```

## Error Handling

The SDK provides a structured exception hierarchy:

| Exception                    | When                                                         |
| ---------------------------- | ------------------------------------------------------------ |
| `WizError`                   | Base class for all SDK errors                                |
| `WizAuthenticationError`     | Auth flow fails                                              |
| `WizAPIError`                | API returns an error (includes `status_code`)                |
| `WizCredentialsError`        | Credentials missing or invalid                               |
| `WizConfigurationError`      | Config file missing or malformed                             |
| `WizRateLimitError`          | Rate limit exceeded (includes `retry_after`)                 |
| `WizQueryError`              | Invalid GraphQL query (includes `query`, `errors`)           |
| `WizSchemaValidationError`   | Query fails schema validation (includes `validation_errors`) |
| `WizReportError`             | Report generation/download fails                             |
| `WizTimeoutError`            | Operation timed out                                          |
| `WizFileError`               | File I/O error                                               |
| `WizServerlessError`         | Serverless-specific failure                                  |

`submit()` never raises for API-level failures — inspect the response instead. Failed requests carry a typed exception on `response.error` (`WizRateLimitError` when the server rate limit won, `WizAPIError` with `status_code` otherwise), and `response.raise_on_error()` converts a failure into that exception:

```python
from wizsec import WizRateLimitError

response = client.create_request(query=QUERY, vars={"first": 100})
result = response.submit()

if not response.success:
    if isinstance(response.error, WizRateLimitError):
        print(f"Rate limited — retry after {response.error.retry_after}s")
    else:
        print(f"Failed: {response.errors}")

# or, exception-style:
response.raise_on_error()
```

`iterate_nodes` raises the typed error directly, since an iterator has no response object to inspect.

## Development

```bash
pip install -e ".[dev]"     # install with dev + docs dependencies
python -m pytest tests/ -q  # run tests
```

### Documentation

API docs are built with [MkDocs Material](https://squidfunk.github.io/mkdocs-material/):

```bash
pip install -e ".[docs]"
mkdocs serve                # live preview at http://127.0.0.1:8000
mkdocs build                # static site in site/
```

## License

MIT
