diff --git a/.github/workflows/cloudg.yml b/.github/workflows/cloudg.yml index 836f8ab..3d0ba9c 100644 --- a/.github/workflows/cloudg.yml +++ b/.github/workflows/cloudg.yml @@ -22,13 +22,15 @@ jobs: lint: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v7 - - uses: actions/setup-python@v5 + - uses: actions/setup-python@v7 with: python-version: ${{ env.PYTHON_VERSION }} - - uses: astral-sh/setup-uv@d4b2f3b6ecc6e67c4457f6d3e41ec42d3d0fcb86 # v5 + - uses: astral-sh/setup-uv@c18668ad3cf93ea998bef934396af7bb5c839dc7 # v10.2.0 + with: + cache-dependency-glob: pyproject.toml - name: Install dependencies run: | @@ -45,13 +47,15 @@ jobs: runs-on: ubuntu-latest needs: lint steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v7 - - uses: actions/setup-python@v5 + - uses: actions/setup-python@v7 with: python-version: ${{ env.PYTHON_VERSION }} - - uses: astral-sh/setup-uv@d4b2f3b6ecc6e67c4457f6d3e41ec42d3d0fcb86 # v5 + - uses: astral-sh/setup-uv@c18668ad3cf93ea998bef934396af7bb5c839dc7 # v10.2.0 + with: + cache-dependency-glob: pyproject.toml - name: Install dependencies run: | @@ -67,13 +71,15 @@ jobs: needs: test if: github.event_name == 'schedule' || github.ref == 'refs/heads/main' steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v7 - - uses: actions/setup-python@v5 + - uses: actions/setup-python@v7 with: python-version: ${{ env.PYTHON_VERSION }} - - uses: astral-sh/setup-uv@d4b2f3b6ecc6e67c4457f6d3e41ec42d3d0fcb86 # v5 + - uses: astral-sh/setup-uv@c18668ad3cf93ea998bef934396af7bb5c839dc7 # v10.2.0 + with: + cache-dependency-glob: pyproject.toml - name: Install dependencies run: | @@ -93,7 +99,7 @@ jobs: continue-on-error: true - name: Upload report artifacts - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v7 with: name: cloudg-report path: ${{ env.OUTPUT_DIR }}/ @@ -104,13 +110,15 @@ jobs: runs-on: ubuntu-latest if: github.event_name == 'pull_request' steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v7 - - uses: actions/setup-python@v5 + - uses: actions/setup-python@v7 with: python-version: ${{ env.PYTHON_VERSION }} - - uses: astral-sh/setup-uv@d4b2f3b6ecc6e67c4457f6d3e41ec42d3d0fcb86 # v5 + - uses: astral-sh/setup-uv@c18668ad3cf93ea998bef934396af7bb5c839dc7 # v10.2.0 + with: + cache-dependency-glob: pyproject.toml - name: Install Checkov run: uv pip install --system checkov @@ -121,7 +129,7 @@ jobs: - name: Upload Checkov results - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v7 with: name: checkov-results path: checkov-results.json @@ -132,10 +140,10 @@ jobs: needs: test if: github.ref == 'refs/heads/main' steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v7 - name: Set up Docker Buildx - uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # v3 + uses: docker/setup-buildx-action@f87e5991a6d7451dcb8d9637bfbc97413f497069 # v4.4.1 - name: Build Docker image run: | diff --git a/.github/workflows/pages.yml b/.github/workflows/pages.yml index 6e9489b..659e9b5 100644 --- a/.github/workflows/pages.yml +++ b/.github/workflows/pages.yml @@ -32,20 +32,20 @@ jobs: name: github-pages url: ${{ steps.deployment.outputs.page_url }} steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v7 - - uses: actions/configure-pages@v5 + - uses: actions/configure-pages@v6 # The artifact name carries the run attempt: on a re-run, the previous # attempt's "github-pages" artifact still exists in the run, and # deploy-pages refuses when it finds two artifacts with the same name # ("Multiple artifacts named github-pages were unexpectedly found"). - - uses: actions/upload-pages-artifact@v3 + - uses: actions/upload-pages-artifact@v5 with: path: docs/ name: github-pages-${{ github.run_attempt }} - id: deployment - uses: actions/deploy-pages@v4 + uses: actions/deploy-pages@v5 with: artifact_name: github-pages-${{ github.run_attempt }} diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index 010e487..195fa0d 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -1,7 +1,7 @@ name: Publish # Runs when a GitHub release is published, or manually from the Actions tab. -# Publishes the Python package to PyPI (trusted publishing / OIDC — no API +# Publishes the Python package to PyPI (trusted publishing / OIDC: no API # token secret needed, but the publisher must be registered once on PyPI, # see notes at the bottom) and the Docker image to GitHub Container Registry. @@ -19,9 +19,11 @@ jobs: build: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v7 - - uses: astral-sh/setup-uv@d4b2f3b6ecc6e67c4457f6d3e41ec42d3d0fcb86 # v5 + - uses: astral-sh/setup-uv@c18668ad3cf93ea998bef934396af7bb5c839dc7 # v10.2.0 + with: + cache-dependency-glob: pyproject.toml - name: Set version from release tag # The release tag is the source of truth: v0.4.0 builds 0.4.0 @@ -44,7 +46,7 @@ jobs: /tmp/smoke/bin/python -c "import cloudg; print('cloudg', cloudg.__version__)" /tmp/smoke/bin/cloudg --version - - uses: actions/upload-artifact@v4 + - uses: actions/upload-artifact@v7 with: name: dist path: dist/ @@ -57,7 +59,7 @@ jobs: # Required for PyPI trusted publishing (OIDC) id-token: write # nosemgrep: codacy.yaml.security.hard-coded-tokens -- OIDC permission scope, not a token value steps: - - uses: actions/download-artifact@v4 + - uses: actions/download-artifact@v8 with: name: dist path: dist/ @@ -76,11 +78,11 @@ jobs: # Required to push to GitHub Container Registry packages: write steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v7 - - uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # v3 + - uses: docker/setup-buildx-action@f87e5991a6d7451dcb8d9637bfbc97413f497069 # v4.4.1 - - uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # v3 + - uses: docker/login-action@dbcb813823bdd20940b903addbd779551569679f # v4.6.0 with: registry: ghcr.io username: ${{ github.actor }} @@ -92,7 +94,7 @@ jobs: run: echo "version=${GITHUB_REF_NAME#v}" >> "$GITHUB_OUTPUT" - name: Build and push image - uses: docker/build-push-action@10e90e3645eae34f1e60eeb005ba3a3d33f178e8 # v6 + uses: docker/build-push-action@c3c9e263c25d99ce0380d002d59b67737d91b0dc # v7.4.0 with: context: . push: true diff --git a/CHANGELOG.md b/CHANGELOG.md index 8ee3612..b2fd5af 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,21 +2,46 @@ Notable changes per release. Patch releases are folded into the major entry they belong to. -## Unreleased +## 0.6.0 (2026-10-09) -A behaviour change in `cloudg.graph.reachability`, which is public API. Embedders that store or count reachability findings will see different open-port findings for the same estate. +Two additions and a round of graph fixes. cloudg is now a Model Context Protocol server, so AI agents can query an inventory map, its graph, findings, compliance and ontology under an access and privacy policy. Every cloud API call now goes through a shared rate limiter that slows down when a provider throttles, so an organization-wide map (or an agent asking for one too often) no longer fails or hammers the APIs. -Fixed: +The release also changes what the graph analysis reports for the same input. Reachability follows network-flow edges only, matches ports as numbers and gives findings stable ids; the ontology emits different triples; `GraphBuilder` merges parallel rule edges; the inventory mapper keeps parallel rules when it deduplicates assets; `cloudg.retry.with_retry` retries a different set of errors. The modules concerned are public API: `cloudg.graph.ports` (new), `cloudg.graph.builder`, `cloudg.graph.reachability`, `cloudg.graph.ontology_rules`, the ontology and RAG exporters (`cloudg.graph.ontology`, `cloudg.graph.rag_export`), `cloudg.inventory.dependencies`, `cloudg.inventory.mapper` and `cloudg.retry`. Embedders that store reachability findings, count graph edges or query the ontology should read Changed and Fixed below. -- `ReachabilityAnalyzer` decided whether a sensitive port was open with a substring test on the edge's `port_range` string. A rule open on `0-65535`, which is how the AWS collector writes protocol `-1` (all traffic), produced no SSH or RDP finding, and neither did an Azure NSG rule with `*` or a comma list such as `443,3380-3390`. The same test reported port 22 for `2200-2300` and port 3389 for `33890`. The string is now parsed into numeric ranges by the new `cloudg.graph.ports` module (`parse_port_ranges()`, `edge_port_ranges()`, `port_in_ranges()`) and each of the 11 ports in `SENSITIVE_PORTS` is tested against those ranges. The module reads `a-b`, single ports, comma lists and `*`. An empty `port_range` on a `SECURITY_GROUP_RULE` or `NACL_RULE` edge means every port when the protocol is `ALL` / `-1` / `*` / `Any`, or TCP, UDP or SCTP (a GCP firewall protocol with no ports listed). On other edge types an empty string still opens nothing, so a GCP `INTERNET_EXPOSED` edge for a public bucket gets no port findings. ICMP, ICMPv6, ESP, AH, GRE and IPIP rules open no ports; for ICMP the numbers AWS stores in the port fields are a type and code. -- Egress rules (`direction: "egress"`) no longer produce open-port findings. They point from a group to `0.0.0.0/0` and admit nothing inbound. Before this change they were checked like ingress rules and only escaped because the default AWS egress string `0-65535` contains none of the sensitive port numbers as a substring; with numeric matching every security group's default egress rule would otherwise report all 11 ports. +Added: -Expect more open-port findings on estates with all-traffic or wide-range ingress rules from `0.0.0.0/0`, and fewer on rules whose range only looked like a sensitive port. An edge still yields at most one finding per port, however many of its ranges overlap that port. The analyzer also reads a `ports` list from a graph edge when one is present, as before; `GraphBuilder` still does not copy `NetworkEdge.ports` onto the graph, since the collectors cap that list at 100 entries and GraphML export cannot store lists. -======= -Behaviour changes in public API: `cloudg.graph.reachability`, `cloudg.graph.ontology_rules`, the ontology and RAG exporters, and `cloudg.inventory.dependencies`. Embedders that store reachability findings or query the ontology should read this section. +- The MCP layer, `cloudg.mcp`, built around one class, `CloudGMCPLayer`. Its default catalog has 74 tools in ten categories (graph 17, inventory 11, findings 10, workspace 8, ontology 6, compliance 5, live 5, privacy 5, meta 4, export 3), 13 resources, 11 resource templates and 11 prompts (`security_posture_review`, `attack_surface_report`, `blast_radius_assessment`, `compliance_gap_analysis`, `incident_triage` and six more). The CLI also lists `cloudg://metrics`. Tools return structured content, resource links and progress notifications; lists are paginated and arguments complete. A list cursor belongs to the query and dataset version that produced it: once the dataset is reloaded or replaced, the old cursor is refused with an invalid-arguments error instead of returning the wrong page. +- Graph analysis in the MCP catalog uses the same rules as the library. Attack paths and `find_paths` follow the reachability network-flow hops plus explicit identity pivots (role assumption, IAM trust, grants); `internet_exposure` and the edge search read Azure `*` ports and `Internet` / `Any` / `*` sources, GCP all-port rules and AWS ICMP rules, and skip egress rules; `find_paths(source="internet")` works on an Azure-only dataset. Findings and asset references resolve through `AssetIndex`, so an ambiguous name is never matched to the first asset that has it, and `reachability_findings` returns the analyser's deterministic ids. +- Tool errors do not repeat the caller's value or list close matches in their message. Both are in the error data instead (`data.value`, `data.suggestions`), where the privacy policy can transform them. +- A dependency-free native server (`cloudg.mcp.native`): stdio, streamable HTTP and the deprecated SSE endpoints, bearer tokens mapped to roles, `Origin` and `Host` checks against loopback by default, CORS, `/healthz`. It serves both protocol eras on the same object: the handshake versions 2024-11-05, 2025-03-26, 2025-06-18 and 2025-11-25, and the stateless 2026-07-28 (`server/discover`, `subscriptions/listen`). `HTTPConfig` caps sessions per principal (`max_sessions_per_principal`, 100; at the cap the principal's least recently used idle session is closed) and open connections (`max_connections`, 512), and gives a client 30 seconds to send its headers (`header_timeout`) and 120 seconds for the body (`body_timeout`). +- Flavors and adapters: `--flavor native`, `sdk` (the official SDK's low-level server), `fastmcp` or `auto`. `register_into(server)` mounts the whole catalog into an mcp 2.x `MCPServer`, an mcp 1.x `FastMCP`, a low-level `Server` or a `fastmcp.FastMCP` (2.x, 3.x, 4.x) next to the host's own tools; `build_server()` and `export_definitions()` cover the other cases. Verified with mcp 1.30.0 and 2.3.0 and fastmcp 2.14.7, 3.4.8 and 4.0.11. On the SDK and fastmcp HTTP transports a session id is bound to the principal that opened it. +- Options objects for embedders: `LayerOptions` (`CloudGMCPLayer(options=...)`) and `ServeOptions` (`serve_async(options=...)`). The existing keyword arguments still work; an unknown keyword raises `TypeError`. +- A tool handler can take a single pydantic arguments model, `def tool(ctx, args: Model)`, instead of one parameter per argument. The input schema and the validation are the same as for the flat form. +- Content `_meta` keys `cloudg/data` and `cloudg/render`: a resource or prompt that renders structured data as text passes the data itself, so the privacy transforms see it with its keys before it is rendered. +- Policies: seven built-in profiles (`open`, `standard` (default), `strict`, `read_only`, `airgapped`, `audit`, `soc-analyst`) plus policy files that extend each other, with per-role rules, sensitivity ceilings, capability denials and per-principal rate limits. Hidden primitives look the same as nonexistent ones. The textual graph and ontology exports (`cloudg://graph/{format}`, `cloudg://ontology/{format}`, `cloudg://ontology/turtle`), `sparql_query` and `subgraph_export` are RESTRICTED, because the privacy layer cannot see inside one opaque string; a profile with a `confidential` ceiling hides them (`strict`, and `soc-analyst` for every role except `lead`). +- Transforms, composed per principal into input and output pipelines: `redact`, `sanitize` (fences attacker-controllable text), `project`, `annotate`, `substitute`, `alias`, `regex_replace`, `rename_keys`, `template`, `guard_secrets` (refuses secrets in arguments) and `depseudonymize`, with built-in detectors for secrets and tokens, private keys, cloud credential ids, ARNs and resource ids, account ids, IP addresses, hostnames and e-mail addresses. +- `TokenVault`: format-preserving pseudonyms derived with HMAC-SHA256, stable across restarts with `CLOUDG_MCP_VAULT_KEY`, and saved when the server shuts down when the policy sets `vault.path`. A pseudonym an agent passes back as an argument is reversed before the tool runs; `reveal_token` reverses pseudonyms and the aliases introduced by `substitute` and `alias` transforms. +- Middleware: `AuditLogMiddleware` (JSONL, argument values hashed), `MetricsMiddleware` (`cloudg://metrics`), `CachingMiddleware`, `ConcurrencyLimitMiddleware`, `RetryMiddleware`, and your own `async (info, call_next)` callables. +- CLI: `cloudg mcp serve | tools | resources | prompts | call | read | config`, also installed as the `cloudg-mcp` console script and runnable as `python -m cloudg.mcp`. `cloudg mcp config` prints client configuration for Claude Desktop, Claude Code, Cursor and VS Code. Under `cloudg mcp`, the banner and UI output go to stderr so stdout stays a clean protocol channel. +- The MCP workspace reads only files under its allowed roots. With none configured (`allowed_roots` or `CLOUDG_MCP_ALLOWED_ROOTS`) they are the working directory and the report directory, and either one is left out when it is `/` or the home directory. A load error never quotes the file's contents, and the search for report files in a directory is bounded. SPARQL queries stop at the row limit or after 60 seconds, and a slow analysis build on one dataset no longer blocks the others or the workspace status. +- Live tools (`map_inventory`, `collect_assets`, `run_scanners`, `run_pipeline`) check the configured credentials before any cloud call (`preflight`, on by default, at most 25 seconds) and run under a live operation guard: identical concurrent calls share one run, a scope collected in the last 120 seconds is refused with `retry_after_seconds` and a loaded live dataset that covers the same providers, and live operations are capped at one per scope and two overall. A live operation that runs longer than `ratelimit.live_operation_timeout_seconds` (default 3600) is cancelled, raises `LiveOperationTimeout` and frees its scopes. `force=true` skips the cooldown for the `admin` and `operator` roles. Each live tool reports the throttling of its own run. `run_scanners` refuses container image names that start with `-`. The new `rate_limit_status` tool and `cloudg://ratelimit` resource show the guard and throttling state without calling a cloud API. +- `cloudg.resilience`: provider-aware error classification (`classify`, `classify_strict` for codes, classes and statuses without message text, `is_provider_answer`, `retry_after`, `describe_error`), adaptive token buckets per provider, account, region and service (per action for EC2) that halve on throttling and recover additively, bulkheads, circuit breakers at the granularity of the rate bucket (per action for EC2, per service elsewhere), `call_with_resilience` / `call_with_resilience_sync` / `@resilient` with decorrelated-jitter backoff that never undercuts Retry-After, a per-call deadline and a per-provider retry budget, per-run telemetry (`stats_scope`, `ResilienceStats`), the process-wide governor (`get_governor`, `configure`, `reset_governor`) and `LiveOperationGuard` with its `GuardSettings`. The first two retries of every call do not draw on the shared retry budget, so one busy run cannot starve every other call of its retries. A half-open breaker lets one probe through and takes the probe back after one cooldown if it never reports; a provider's non-retryable answer (an `AccessDenied`, a 404) counts as proof that the service is reachable. Calls refused by an open circuit, and calls given up because the retry budget ran out, do not count as breaker failures. Reconfiguring the limits keeps the breakers, budgets and held concurrency slots of providers whose settings did not change. +- Provider wiring, with rate limits, concurrency caps, breakers and the shared retry budget enforced in every SDK: botocore event hooks on every AWS session cloudg creates (the collectors' aioboto3 sessions, the base and assumed-role sessions behind STS, Organizations and Control Tower discovery and the organization fan-out); an explicit azure-core `RetryPolicy` and a per-attempt policy on every Azure client, which also reads the `x-ms-ratelimit-remaining-subscription-reads` header (pausing the whole subscription) and Resource Graph's `x-ms-user-quota-*` headers; Resource Graph paging that honours Retry-After and quota resets, without retrying a page again on top of azure-core's own 429 retries; a google-api-core `Retry` capped per page, and per-page pacing, for Cloud Asset Inventory. An instrumented AWS client that refuses a call because a circuit is open or the retry budget is spent raises an error that is also a botocore `ClientError` with code `Throttling`, so existing `except ClientError` handlers catch it. A cancelled aioboto3 call frees its concurrency slot at once. +- `ratelimit` config section (`RateLimitConfig`, `ProviderRateLimitConfig`, `ServiceRateLimitConfig`): master and per-provider switches, rates, bursts, account and provider-wide limits, concurrency, adaptive floor, breaker threshold and cooldown, retries, backoff cap, deadline, retry budget, per-service and per-operation overrides, and the `live_*` settings of the guard. Unknown keys are logged and ignored. `config.yaml` ships the section with comments. +- `InventoryResult.throttling`: the run's throttling summary (totals, per-scope counters, messages, skipped scopes, tripped breakers, slowed buckets), or `None` when no API pushed back. It is added to `summary` and `inventory-map.json` only when present, and `load()` reads it back. `MultiAccountCollector.resilience_stats` holds the last run's stats. +- `cloudg.graph.ports`: the port and edge rules that the graph builder, the reachability analysis and the ontology share. `parse_port_ranges()`, `edge_port_ranges()` and `port_in_ranges()` turn `port_range` strings into numeric ranges; `is_internet_source()` is the single check for a source that stands for the internet (`0.0.0.0/0`, `::/0`, and the Azure tags `Internet`, `Any` and `*` in any letter case). +- Network-flow traversal as public API: `ReachabilityAnalyzer.flow_hops()`, `flow_successors()` and `internet_entry_points()`, and `cloudg.graph.reachability.network_flow_graph()`, so other code walks the network the way the exposure analysis does. `finding_id()` and `asset_key()` give the id scheme of reachability findings. +- `cloudg.inventory.dependencies.AssetIndex`, the one asset matcher: asset id, ARN or alias, unique name, unique ARN tail, plus two opt-in tiers (`arn_parts` for ARN resource parts such as `db:orders`, `casefold` for names in any letter case). `AssetIndex.add()` takes extra `aliases`. `DependencyGraph.find`, the ontology and the RAG exporter use it. +- `RelationType.CONTAINS`, a generic containment relation in the `CONTAINMENT` group, for 64 relation types in all. Changed: +- `aws.max_retries` and `aws.retry_mode` are honoured. Earlier releases accepted them and then hard-coded adaptive mode with 10 retries, so the defaults behave as before. botocore counts `max_retries` after the first attempt: the default of 10 allows up to 11 attempts per call. +- Every cloud API call is paced by default at or below the provider's published limits: AWS 20 rps per service and per EC2 action, Route 53 5, tagging API 5, Azure 20 rps per subscription and 3 rps per user for Resource Graph, GCP 1.5 rps for `ListAssets`. Per-item fan-out is held to those rates, so maps of AMI-heavy or bucket-heavy accounts can take longer than in 0.5.x. `ratelimit.enabled: false` turns the pacing off. +- A service that stays throttled after its retries is recorded in coverage as FAILED or PARTIAL with an error starting `throttled:` instead of failing the run. AWS collectors that swallowed per-item errors used to report SUCCESS with missing assets when they were throttled; they now report PARTIAL, telling calls that gave up after retries apart from calls an open circuit skipped. +- `cloudg.retry.with_retry` keeps its signature and still retries everything it retried in 0.5.x. It now also retries provider throttling identified by error code, exception class or HTTP status (Azure 429s, GCP `ResourceExhausted`), and never sleeps less than the server's Retry-After, capped at its `max_delay`. It does not retry cloudg's own resilience errors (an open circuit, a spent retry budget, a missed deadline), a plain HTTP 500 or an error whose message only mentions rate limiting. `cloudg.retry` re-exports the main `cloudg.resilience` names. +- New optional extra `mcp` (`mcp>=1.30,<3`). `all` now includes it, and so does `dev`. The core layer and its native server need no extra. +- The wheel ships `docs/DOCUMENTATION.md` and `docs/INVENTORY_REFERENCE.md` as `cloudg/_docs/`, so the MCP `cloudg://docs/{topic}` resources work from a plain `pip install`. The sdist includes every `docs/*.md`. - `ReachabilityAnalyzer.find_internet_exposed()` walks network-flow edges only, so `generate_findings()` no longer reports resources that are reachable only through an IAM grant. Before, the walk followed every edge type out of an internet-exposed node, including `GRANTS_ACCESS`, `ASSUMES_ROLE`, `IAM_TRUST` and `IAM_POLICY_ATTACHMENT`, and produced findings such as a CRITICAL "Internet-exposed RDS_INSTANCE" for a database whose only link to the internet was a role granted access to it. The rules, documented in the module docstring: - `INTERNET_EXPOSED`, `LOAD_BALANCER_TARGET`, `ROUTE` and `PEERING` are followed source to target. - `SECURITY_GROUP_RULE` and `NACL_RULE` are followed source to target, except egress rules (`direction: "egress"`). @@ -25,18 +50,66 @@ Changed: - `INVOKES`, `REFERENCES`, `USES_IMAGE`, `LOGS_TO`, `PROTECTS`, `MONITORS`, `MANAGES` and `GOVERNS` are not followed. A Lambda function behind a public API Gateway (an `INVOKES` edge) is therefore no longer reported as internet-exposed by this analysis. - `compute_blast_radius()` is unchanged and still follows every edge type. - Internet entry points are nodes named `0.0.0.0/0` or `::/0` and the source of any non-egress edge whose CIDR is `0.0.0.0/0`, `::/0` or an Azure service tag for the internet (`Internet`, `Any`, `*`). Before, every external placeholder node counted as an entry point, including private CIDRs such as `10.0.0.0/8` and unresolved references, and a security group with an egress rule to `0.0.0.0/0` became an entry point itself. -- Reachability finding ids are deterministic: a UUID5 hash of the rule and the asset (`cloudg.graph.reachability.finding_id(rule, asset)`) instead of a random `uuid4`. The same exposure gets the same id on every run. Exposure findings key on the node id, open-port findings on the rule edge and port. `generate_findings()` also returns exposure findings in a stable order. -- Ontology: `CONTAINS` edges get a relation that matches both endpoint types: `VPC_CONTAINS_SUBNET` (VPC or VNet to subnet), `SUBNET_CONTAINS_INSTANCE`, `CLUSTER_CONTAINS_SERVICE` (cluster or Kubernetes namespace to workload) and `ORG_CONTAINS_ACCOUNT` (organization or OU to account). Every other containment (organization to OU, account to VPC, VNet to VM, resource group to resource, unresolved endpoints) uses the new generic `RelationType.CONTAINS` in the `CONTAINMENT` group. Before, these fell back to `VPC_CONTAINS_SUBNET` and produced triples such as `sample-org VPC_CONTAINS_SUBNET Workloads`. A relationship declared on a `CONTAINS` edge now replaces the inferred one instead of being added to it. -- Ontology: `SECURITY_GROUP_RULE` and `NACL_RULE` edges no longer yield `PROTECTED_BY_SG` / `PROTECTED_BY_NACL`. Those edges point from the traffic source to the group, so they produced triples such as `0.0.0.0/0 PROTECTED_BY_SG sg-web`. The protection relations now come from `ATTACHED_TO` edges whose target is a security group or NSG (`PROTECTED_BY_SG`) or a NACL (`PROTECTED_BY_NACL`), giving `web-alb PROTECTED_BY_SG sg-web`. Other `ATTACHED_TO` edges without a declared relationship stay `DEPENDS_ON`. -- Ontology: `ENCRYPTED_BY_KMS` now links an encrypted asset to its key: `orders-db ENCRYPTED_BY_KMS `. Before, `infer_relations()` added it to every edge whose target had `encryption` or `storage_encrypted` plus a `kms_key_id` in its metadata, so `CloudOntology` wrote triples with the edge's source as the subject and the encrypted asset as the object, such as `ec2-1 ENCRYPTED_BY_KMS rds-1` or `subnet-private ENCRYPTED_BY_KMS orders-db`. The relation now comes from `infer_asset_relations()` for the encrypted asset itself. Its `kms_key_id` is matched against the collected `KMS_KEY` assets by asset id, ARN, bare key id, alias name or alias ARN. If no collected key matches, the object is the raw `kms_key_id` string. `infer_relations()` no longer returns `ENCRYPTED_BY_KMS` unless the edge declares it as its `relationship` (the linker's `REFERENCES` edges still do). The RAG export, which builds its chunks from `infer_relations()`, loses the wrong per-edge entries. +- Security groups, NSGs, NACLs and target groups are no longer reported as "unexpected internet-exposed resources". The walk still passes through them and marks them `is_internet_exposed`; a group's open rules are reported by the sensitive-port findings, and a target group's exposure by the findings on the targets behind it. +- Reachability finding ids are deterministic: a UUID5 hash of the rule and a stable asset key (`cloudg.graph.reachability.finding_id(rule, asset_key)`) instead of a random `uuid4`. The key is the asset's ARN, or its asset id when it has no ARN, because collectors give assets a fresh random id on every collection. Exposure findings key on the asset, open-port findings on the keys of the rule's two endpoints and the port, so the same exposure gets the same id on every scan. `generate_findings()` also returns exposure findings in a stable order. +- Ontology: `CONTAINS` edges get a relation only when both endpoint types match it: `VPC_CONTAINS_SUBNET` (VPC or VNet to subnet), `SUBNET_CONTAINS_INSTANCE` (subnet to a resource placed in it, not a route table, NACL, gateway or other network construct), `CLUSTER_CONTAINS_SERVICE` (cluster or Kubernetes namespace to a workload, service, namespace, node group or add-on) and `ORG_CONTAINS_ACCOUNT` (organization or OU to account). Every other containment (organization to OU, account to VPC, VNet to VM, resource group to resource, unresolved endpoints) uses the new generic `RelationType.CONTAINS`. Before, these fell back to `VPC_CONTAINS_SUBNET` and produced triples such as `sample-org VPC_CONTAINS_SUBNET Workloads`. +- Ontology: a relationship declared on an edge (`NetworkEdge.relationship`) replaces the inferred relation on `CONTAINS`, `ATTACHED_TO` and the typed inventory edges (`INVOKES`, `USES_IMAGE`, `ASSUMES_ROLE`, `GRANTS_ACCESS`, `LOGS_TO`, `PROTECTS`, `MONITORS`, `MANAGES`, `GOVERNS`, `REFERENCES`). On `SECURITY_GROUP_RULE`, `IAM_TRUST`, `IAM_POLICY_ATTACHMENT`, `LOAD_BALANCER_TARGET`, `ROUTE`, `PEERING` and `INTERNET_EXPOSED` edges the declared relation comes first and the inferred ones are added after it. +- Ontology: `SECURITY_GROUP_RULE` and `NACL_RULE` edges no longer yield `PROTECTED_BY_SG` / `PROTECTED_BY_NACL`. Those edges point from the traffic source to the group, so they produced triples such as `0.0.0.0/0 PROTECTED_BY_SG sg-web`. The protection relations now come from `ATTACHED_TO` edges whose target is a security group or NSG (`PROTECTED_BY_SG`) or a NACL (`PROTECTED_BY_NACL`), so the triple reads `web-alb PROTECTED_BY_SG sg-web`. Other `ATTACHED_TO` edges without a declared relationship stay `DEPENDS_ON`. +- Ontology: the port relations of a security group rule are read from the parsed port ranges, as the sensitive-port findings are. A rule open on every port (`0-65535`, protocol `ALL` / `-1` / `*`, or an Azure `*` rule for `Tcp`) is `ALL_TRAFFIC`; otherwise each of ports 22, 80, 443 and 3389 inside the ranges gives `ONLY_SSH`, `ONLY_HTTP`, `ONLY_HTTPS` or `ONLY_RDP` (an AWS rule for `1-1024` gives the first three), and a rule that opens none of them is `PORT_RESTRICTED`. A rule from an Azure `Internet`, `Any` or `*` source is `INTERNET_REACHABLE`, like one from `0.0.0.0/0`. +- Ontology: `ENCRYPTED_BY_KMS` now links an encrypted asset to its key: `orders-db ENCRYPTED_BY_KMS `. Before, `infer_relations()` added it to every edge whose target had `encryption` or `storage_encrypted` plus a `kms_key_id` in its metadata, so `CloudOntology` wrote triples with the edge's source as the subject and the encrypted asset as the object, such as `ec2-1 ENCRYPTED_BY_KMS rds-1` or `subnet-private ENCRYPTED_BY_KMS orders-db`. The relation now comes from `infer_asset_relations()` for the encrypted asset itself. Its `kms_key_id` is matched against the collected `KMS_KEY` assets by asset id, ARN, bare key id, alias name or alias ARN. The keys are indexed once per ontology build instead of being scanned for every encrypted asset. If no collected key matches, the object is the raw `kms_key_id` string. `infer_relations()` no longer returns `ENCRYPTED_BY_KMS` unless the edge declares it as its `relationship` (the linker's `REFERENCES` edges still do). - Ontology: containment inferred from `vpc_id` metadata points from the VPC to its members. Before, `infer_asset_relations()` returned `(VPC_CONTAINS_SUBNET, vpc)` for a subnet and `(SUBNET_CONTAINS_INSTANCE, vpc)` for anything else in the VPC, which gave `subnet VPC_CONTAINS_SUBNET vpc` and `instance SUBNET_CONTAINS_INSTANCE vpc`: the wrong direction, and a subnet relation whose object was a VPC. Now a VPC or VNet asset returns `(VPC_CONTAINS_SUBNET, subnet)` for each subnet and `(CONTAINS, member)` for every other asset with the same `vpc_id`, and member assets return no containment. The return shape of `infer_asset_relations()` is unchanged: the asset passed in is still the subject of every tuple. If several VPC or VNet assets share a `vpc_id`, the first one in the list claims the members. `CloudOntology.query_asset_neighbourhood()` with `hops > 1` follows these relations forward, so a walk from a VPC now reaches its members, and a walk from an instance no longer reaches the VPC through them. +- RAG export: relation-group chunks include the asset-level relations (`ENCRYPTED_BY_KMS`, `vpc_id` containment, tag governance, key rotation, EC2 security group membership) next to the edge relations. The per-edge `ENCRYPTED_BY_KMS` entries described above are gone. +- `GraphBuilder` merges parallel rule edges (see Fixed), so `graph.number_of_edges()` and the edge lists in `to_d3_json()`, `to_cytoscape_json()` and GraphML exports shrink by the number of merged rules, and a merged edge's `port_range`, `protocol` and `description` can hold comma- or semicolon-separated lists. Fixed: -- The ontology and RAG exporters attached findings to assets by raw `Finding.resource_id` only. Findings whose `resource_id` is an ARN or a display name, which is common in scanner output, left a dangling resource node (for example `cmr:arn:aws:rds:...`) with no `FINDING_AFFECTS` edge to the real asset, and RAG entity and community chunks reported a `finding_count` of 0 for it. Both exporters now resolve a finding's asset by ID, then ARN (from `resource_id` or `resource_arn`), then unique name, then unique ARN tail, and fall back to the raw `resource_id` only when nothing matches. Ambiguous names never match. -- The matching rules live in a new `cloudg.inventory.dependencies.AssetIndex`, which `DependencyGraph.find` now uses too. Its behaviour is unchanged apart from building the lookup tables once instead of scanning every asset on each call, and an empty reference now returns `None`. -- AWS `SECURITY_GROUP_RULE` edges from `collect_edges()` now use the security group asset's `id` as the security group endpoint instead of the native `sg-…` ID, as Azure NSG rule edges already did. The graph exports no longer contain a disconnected `EXTERNAL` node per security group, and sensitive-port findings carry the security group's asset ID and ARN (their deterministic ids change once, since the rule edge key now holds the asset ID). The internet-exposure walk now continues from an open security group to the resources attached to it, so an instance behind a group that admits `0.0.0.0/0` is reported as internet-exposed; before, the walk stopped at the native-ID placeholder. The CIDR endpoint is unchanged. -- Reachability analysis no longer reports security groups and NSGs themselves as "unexpected internet-exposed resources"; their open rules are still reported by the sensitive-port findings. +- `ReachabilityAnalyzer` decided whether a sensitive port was open with a substring test on the edge's `port_range` string. A rule open on `0-65535`, which is how the AWS collector writes protocol `-1` (all traffic), produced no SSH or RDP finding, and neither did an Azure NSG rule with `*` or a comma list such as `443,3380-3390`. The same test reported port 22 for `2200-2300` and port 3389 for `33890`. The string is now parsed into numeric ranges by `cloudg.graph.ports` and each of the 11 ports in `SENSITIVE_PORTS` is tested against those ranges. The module reads `a-b`, single ports, comma lists and `*`. An empty `port_range` on a `SECURITY_GROUP_RULE` or `NACL_RULE` edge means every port when the protocol is `ALL` / `-1` / `*` / `Any`, or TCP, UDP or SCTP (a GCP firewall protocol with no ports listed). On other edge types an empty string still opens nothing, so a GCP `INTERNET_EXPOSED` edge for a public bucket gets no port findings. ICMP, ICMPv6, ESP, AH, GRE and IPIP rules open no ports; for ICMP the numbers AWS stores in the port fields are a type and code. Expect more open-port findings on estates with all-traffic or wide-range ingress rules from the internet, and fewer on rules whose range only looked like a sensitive port. An edge still yields at most one finding per port, however many of its ranges overlap that port. The analyzer also reads a `ports` list from a graph edge when one is present, as before; `GraphBuilder` still does not copy `NetworkEdge.ports` onto the graph, since the collectors cap that list at 100 entries and GraphML export cannot store lists. +- Egress rules (`direction: "egress"`) no longer produce open-port findings. They point from a group to `0.0.0.0/0` and admit nothing inbound. Before this change they were checked like ingress rules and only escaped because the default AWS egress string `0-65535` contains none of the sensitive port numbers as a substring; with numeric matching every security group's default egress rule would otherwise report all 11 ports. +- Azure NSG rules whose source is the `Internet`, `Any` or `*` service tag now produce open-port findings, and the finding title names the tag. Before, they marked resources internet-exposed but no port finding was ever raised for them. +- `GraphBuilder` lost rules that shared a source and target. The graph is an `nx.DiGraph`, which holds one edge per pair, so the second `NetworkEdge` for a pair overwrote the first one's attributes. A security group with ingress rules from `0.0.0.0/0` for port 22 and port 443 ended up with a single edge for 443, and `ReachabilityAnalyzer` never reported SSH open to the internet. An ICMP rule from `0.0.0.0/0` listed after the SSH rule hid it the same way. AWS security groups, Azure NSG rules and GCP firewall rules all produce such pairs. `SECURITY_GROUP_RULE`, `NACL_RULE` and `INTERNET_EXPOSED` edges on the same pair are now merged into one graph edge when their CIDR and direction match: + - `port_range` lists the ports of every merged rule, comma-separated and without repeats: rules for 22 and 443 give `"22,443"`. + - Rules for different protocols are merged too, and `protocol` lists each one once (`"TCP,ICMP"`). ICMP, ESP, AH, GRE and IPIP rules add no ports, because the numbers AWS stores for ICMP are a type and code, and Azure writes `"*"` for an ICMP rule. A TCP, UDP or all-protocol filter rule with no ports (GCP writes these) adds `0-65535`. With mixed protocols the ports are a union: the edge shows which ports are open from the source, not which protocol each port is open for. Sensitive-port findings never looked at the protocol, so they come out the same as if each rule had its own edge. + - `description` joins the rules' descriptions with `"; "`. + - Rules whose CIDR or direction differ, and edges of any other type, still replace the earlier edge on the pair. +- The inventory mapper's deduplication dropped parallel security group, NACL and internet-exposure rules from `inventory-map.json` whenever it merged a duplicate asset. Rules that differ in ports, protocol, CIDR or direction are now kept. +- The ontology and RAG exporters attached findings to assets by raw `Finding.resource_id` only. Findings whose `resource_id` is an ARN or a display name, which is common in scanner output, left a dangling resource node (for example `cmr:arn:aws:rds:...`) with no `FINDING_AFFECTS` edge to the real asset, and RAG entity and community chunks reported a `finding_count` of 0 for it. Both exporters now resolve a finding's asset with `AssetIndex`: by id, then ARN (from `resource_id` or `resource_arn`), then unique name, then unique ARN tail, falling back to the raw `resource_id` only when nothing matches. Ambiguous names never match. `DependencyGraph.find` behaves as before apart from building its lookup tables once instead of scanning every asset on each call, and an empty reference now returns `None`. +- AWS `SECURITY_GROUP_RULE` edges from `collect_edges()` now use the security group asset's `id` as the security group endpoint instead of the native `sg-...` ID, as Azure NSG rule edges already did. The graph exports no longer contain a disconnected `EXTERNAL` node per security group, and sensitive-port findings carry the security group's asset ID and ARN. The internet-exposure walk now continues from an open security group to the resources attached to it, so an instance behind a group that admits `0.0.0.0/0` is reported as internet-exposed; before, the walk stopped at the native-ID placeholder. The CIDR endpoint is unchanged. +- The multi-account collector called STS (`AssumeRole`, `GetCallerIdentity`) synchronously on the event loop, which stalled every other collection while it waited. Those calls now run in a worker thread. + +Security: + +- The MCP layer never echoes a real value that it restored from a pseudonym in the arguments. It records each restored value with its pseudonym, and error messages and error data, results, resource links, progress messages and log notifications are pseudonymised again with those pairs before they leave, for tools, resources, prompts and completions alike. Text payloads (JSON strings, GraphML, N-Triples) are pseudonymised consistently with structured results, using the values the vault already knows (`TokenVault.replace_known_values`). +- Progress messages, log notifications, resource-link names, titles and descriptions, and the URIs of embedded resources go through the caller's privacy policy. +- An unexpected exception in a resource or prompt handler reaches the client as a generic error on every flavor, never as the raw exception text. +- Percent-encoded and templated resource URIs are put in canonical form before the policy's allow and deny rules are checked, so `cloudg://graph/%64%33` is treated as `cloudg://graph/d3`. +- DNS rebinding: on a non-loopback bind without `--allowed-host`, the native HTTP server no longer trusts an `Origin` just because it equals the `Host` header. Same-origin requests are accepted only when the `Host` header is validated. +- `cloudg mcp serve` refuses an empty `--auth-token` or `CLOUDG_MCP_AUTH_TOKENS`. `--registry` accepts only a `module:attribute` dotted name that resolves to a `Registry`; it loads operator code and must never come from an untrusted source. +- Session caps, connection caps and read timeouts on the native HTTP server, and session ids bound to the principal on the SDK and fastmcp HTTP transports (see Added). +- Policy allow and deny lists and rules also match the concrete resource URI that was requested, both as sent and percent-decoded. +- `preview_transform` pseudonymises with a separate vault under a random key (`Policy.preview_vault()`), so a preview cannot confirm which real value is behind a pseudonym. The `strict` profile hides the tool. +- `privacy_audit_log` returns only the caller's own entries unless the caller has the `admin` or `privacy-admin` role, and its result says which (`scope`: `own` or `all`). +- Audit log lines escape CR, LF and other control characters, and a probe for an unknown resource is recorded with a keyed hash of the URI's variable part instead of the value. +- An invalid policy is reported without echoing the values in it, so a misplaced vault key does not end up in an error message. +- Vault files are written atomically with fsync, with mode 0600 in directories created with mode 0700. +- Regular expressions that could backtrack for a long time on hostile input were replaced with linear patterns or small parsers: the vault's token scanner, the detectors for GCP resource names, e-mail addresses, hostnames and AWS secret access keys, the prompt-injection pattern, the RAG member and triple parsers, fence stripping, the `explain_*` metadata helpers and the literal stripping in `sparql_query` (a 20,000-character query used to take 23 seconds there; 200,000 characters now take 0.06). +- Over HTTP with OAuth, the scopes of the access token become the principal's roles as they are, so a token with the scope `admin` or `privacy-admin` gets what the policy grants those roles, reveal included. Issue such scopes only to people who may see real identifiers, or pass a `principal_resolver` that maps scopes onto roles. + +Notes: + +- Existing imports, signatures and config files keep working, and a config without `ratelimit` gets the built-in limits. The release is not purely additive, though: the reachability, ontology, graph builder, mapper and `with_retry` changes under Changed and Fixed alter results for the same input. Reachability finding ids, open-port findings, ontology triples, graph edge counts and the rules kept in `inventory-map.json` all differ from 0.5.x. +- Several large modules were split into packages or sibling modules (`cloudg.mcp.state` and `cloudg.mcp.policy` are packages, the transforms are spread over smaller modules, `cloudg.api` moved its scanner runs to `cloudg.api_scanners`, the linker rules moved to `cloudg.inventory.linker_rules`), and long functions across the CLI, the collectors, the resilience layer and the MCP server were split up. The documented import paths, public signatures and CLI options are unchanged. The GCP default skip list is now `cloudg.config.GCP_DEFAULT_SKIP_ASSET_TYPES`; `cloudg.collectors.gcp.DEFAULT_SKIP_ASSET_TYPES` stays as an alias. +- The `sanitize` transform is about three times faster on large results, and the redactor no longer scans the same value twice. +- The `mcp` extra allows `mcp>=1.30,<3`, verified with mcp 1.30.0 and 2.3.0. Mounting into an SDK `MCPServer` (mcp 2.x) or `FastMCP` (mcp 1.x) goes through private SDK attributes that a future release can rename, so applications using that path should pin the minor version they tested, for example `mcp>=2.3,<2.4`. Mounting into the low-level `Server` and the native server are unaffected. +- Sessions on the legacy SSE transport of the SDK and fastmcp flavors are not bound to a principal. + +Documentation: + +- [MCP guide](docs/MCP.md): installation, client setup, CLI, transports and protocol versions, flavors, mounting, in-process use, workspace, errors, middleware, live tools, troubleshooting. +- [MCP tool reference](docs/MCP_TOOLS.md): every tool, resource, template and prompt. +- [MCP privacy](docs/MCP_PRIVACY.md): profiles, policy schema, transforms, detectors, the token vault. +- [MCP internals](docs/MCP_INTERNALS.md): package map, core contracts, the layer, the adapters, the native server. +- [Rate limits and throttling](docs/RESILIENCE.md): architecture, provider wiring, built-in limits with sources, every `ratelimit` key, a throttled run's output, the Python API, tuning, limitations. +- The feature reference and the handbook have new MCP server and rate limit chapters, the `ratelimit` configuration and the new Python API. The inventory reference, catalog and internals describe merged rule edges, the port and protocol formats of each provider and the new reachability rules. ## 0.5.3 (2026-10-04) @@ -134,7 +207,7 @@ Fixed (pre-release hardening, folded in before publish): - install.sh no longer uses `A && B || C` as if-then-else anywhere; every step is an explicit if/else, with a `pip_tool` helper for the repeated pip installs. - Dockerfile pins parliament to 1.6.4. - Codacy complexity findings: the long CLI commands (`run`, `scan`, `ingest`), the engine phases (`scan`, `analyze`, `run_pipeline`), the SVG/HTML renderers, the ontology relation inference, the RAG exporters, credentials, registry discovery and the Prowler/IAM-linter scanners were split into focused helpers; behavior is unchanged and the CLI `run` command no longer takes 27 named parameters. -- Codacy file-length findings: the four oversized modules were split along their natural seams, with every import path preserved — the CLI helpers moved to `cli_helpers.py` / `cli_run_helpers.py` and the `run` command to `cli_commands.py`; the AWS per-service collectors moved to `aws_services.py` / `aws_services_extended.py` mixins; the ontology relation-inference layer moved to `ontology_rules.py` (re-exported from `ontology`); the SVG hierarchical layout moved to `svg_layout.py`. The SVG output is byte-identical and the whole suite passes unchanged. +- Codacy file-length findings: the four oversized modules were split along their natural seams, with every import path preserved. The CLI helpers moved to `cli_helpers.py` / `cli_run_helpers.py` and the `run` command to `cli_commands.py`; the AWS per-service collectors moved to `aws_services.py` / `aws_services_extended.py` mixins; the ontology relation-inference layer moved to `ontology_rules.py` (re-exported from `ontology`); the SVG hierarchical layout moved to `svg_layout.py`. The SVG output is byte-identical and the whole suite passes unchanged. ## 0.4.1 (2026-09-29) diff --git a/Dockerfile b/Dockerfile index cc831be..8a5446a 100644 --- a/Dockerfile +++ b/Dockerfile @@ -86,8 +86,8 @@ RUN echo "=== Verifying scanner installations ===" && \ trivy --version || echo "WARNING: trivy not working" && \ echo "=== Scanner verification complete ===" -# Parliament (IAM linting) — small pure-Python package, install directly -RUN pip install --no-cache-dir parliament +# Parliament (IAM linting): small pure-Python package, pinned like the builder stage +RUN pip install --no-cache-dir parliament==1.6.4 # ── Everything below here rebuilds on every code change (fast) ── diff --git a/README.md b/README.md index 10bef89..0dbd130 100644 --- a/README.md +++ b/README.md @@ -5,7 +5,7 @@

License: MIT Python 3.11+ - Version 0.5.3 + Version 0.6.0 uv

@@ -20,7 +20,9 @@ One command collects assets from every configured provider in parallel, feeds th Need the map without the security tooling? `cloudg map` is a scanner-independent inventory mapper: it deep-collects everything deployed (or default) in an account, down to the network fabric, sweeps every service for the rest, and links it all into one asset map you can overlay with scanner findings later. See [Inventory mapping](#inventory-mapping). -**Full documentation:** the rendered handbook lives at [morpheuslord.github.io/cloudg](https://morpheuslord.github.io/cloudg/), with the same content as markdown in the [feature reference](https://github.com/morpheuslord/cloudg/blob/main/docs/DOCUMENTATION.md) and release notes in the [changelog](https://github.com/morpheuslord/cloudg/blob/main/CHANGELOG.md). The inventory mapper has its own deep-dive documents: the [inventory reference](https://github.com/morpheuslord/cloudg/blob/main/docs/INVENTORY_REFERENCE.md) (every return structure and output file, field by field), the [inventory catalog](https://github.com/morpheuslord/cloudg/blob/main/docs/INVENTORY_CATALOG.md) (all 156 asset types, their metadata and relationships) and the [inventory internals](https://github.com/morpheuslord/cloudg/blob/main/docs/INVENTORY_INTERNALS.md) (how it works and how to extend it). +Want an AI agent to work with all of it? `cloudg mcp serve` turns cloudg into an MCP server for Claude Desktop, Claude Code, Cursor or anything else that speaks the Model Context Protocol. See [MCP server](#mcp-server). + +**Full documentation:** the rendered handbook lives at [morpheuslord.github.io/cloudg](https://morpheuslord.github.io/cloudg/), with the same content as markdown in the [feature reference](https://github.com/morpheuslord/cloudg/blob/main/docs/DOCUMENTATION.md) and release notes in the [changelog](https://github.com/morpheuslord/cloudg/blob/main/CHANGELOG.md). The inventory mapper has its own deep-dive documents: the [inventory reference](https://github.com/morpheuslord/cloudg/blob/main/docs/INVENTORY_REFERENCE.md) (every return structure and output file, field by field), the [inventory catalog](https://github.com/morpheuslord/cloudg/blob/main/docs/INVENTORY_CATALOG.md) (all 156 asset types, their metadata and relationships) and the [inventory internals](https://github.com/morpheuslord/cloudg/blob/main/docs/INVENTORY_INTERNALS.md) (how it works and how to extend it). The MCP server is covered by the [MCP guide](https://github.com/morpheuslord/cloudg/blob/main/docs/MCP.md), the [tool reference](https://github.com/morpheuslord/cloudg/blob/main/docs/MCP_TOOLS.md) and the [privacy reference](https://github.com/morpheuslord/cloudg/blob/main/docs/MCP_PRIVACY.md), and throttling by [rate limits and throttling](https://github.com/morpheuslord/cloudg/blob/main/docs/RESILIENCE.md). --- @@ -47,7 +49,7 @@ cloudg run -p aws --regions us-east-1 --terraform Reports land in `./reports`. Open `report.html` first. -Already ran the scanners yourself? Feed cloudg their native output files instead — no cloud credentials, no scanner binaries: +Already ran the scanners yourself? Feed cloudg their native output files instead. No cloud credentials, no scanner binaries: ```bash cloudg ingest --prowler ./prowler-output/ --checkov ./results_json.json --trivy ./trivy.json @@ -115,6 +117,31 @@ print(inventory.summary["assets_by_service"]) --- +## MCP server + +cloudg can hand its inventory, graph, findings, compliance mappings and ontology to AI agents over the Model Context Protocol: 74 tools, 13 resources, 11 resource templates and 11 prompts, from `find_assets` and `attack_paths` to `blast_radius`, `compliance_gaps` and a `security_posture_review` prompt. + +```bash +pip install "cloudg[mcp]" # the native server needs no extra; this adds the official SDK (mcp>=1.30,<3) + +# stdio for a desktop client, with a map you already made +cloudg mcp serve --dataset prod=./reports/inventory-map.json +cloudg mcp config --client claude-desktop --dataset prod=./reports/inventory-map.json + +# streamable HTTP on 127.0.0.1:8765/mcp with a bearer token +cloudg mcp serve --transport http --auth-token env:CLOUDG_MCP_TOKEN:analyst +``` + +Already running an MCP server? Mount cloudg into it (official `mcp` SDK 1.x or 2.x, or fastmcp) with `CloudGMCPLayer(...).register_into(server)`, or skip the wire and `await layer.call_tool(...)` from Python. Every request, whichever way it arrives, goes through one policy. The default `standard` profile redacts secrets and fences attacker-controllable text; `strict` also pseudonymises account ids, ARNs, IPs and names (pseudonyms passed back as arguments are resolved) and hides the tools that touch the cloud. The live tools (`map_inventory`, `collect_assets`, `run_scanners`, `run_pipeline`) check credentials before the first API call and refuse to re-collect a scope that was collected a couple of minutes ago. + +The [MCP guide](https://github.com/morpheuslord/cloudg/blob/main/docs/MCP.md) covers serving, client setup, transports, mounting and in-process use; the [tool reference](https://github.com/morpheuslord/cloudg/blob/main/docs/MCP_TOOLS.md) lists every primitive, and the [privacy reference](https://github.com/morpheuslord/cloudg/blob/main/docs/MCP_PRIVACY.md) every profile and transform. + +### Rate limits + +Mapping a whole AWS Organization means thousands of API calls per account, so every call cloudg makes goes through a shared rate limiter with per-service limits at or below what each provider publishes. A throttled API gets slowed down for every client at once, Retry-After is honoured, requests in flight are capped per account, retries beyond the first two of each call draw from one budget per provider, and an API that stays throttled is skipped for a while by a circuit breaker. The map does not fail over it: coverage says `throttled: ...` and the result carries a `throttling` summary. Limits are tunable under `ratelimit` in `config.yaml`; see [rate limits and throttling](https://github.com/morpheuslord/cloudg/blob/main/docs/RESILIENCE.md). + +--- + ## Authentication Every provider supports several auth methods, resolved in a fixed priority order. The same config works on a laptop, in CI, and on cloud compute. The full set of fields lives in `config.yaml` with comments for each method. @@ -174,9 +201,9 @@ graph LR Collection runs all providers concurrently with asyncio, iterating accounts and regions per provider (regions are auto-discovered when you pass `--regions all`). Assets and network edges go into a directed graph, where BFS from the internet node over network-flow edges finds exposed resources and blast radius scoring estimates what an attacker could reach from each node. -The same inventory feeds three other exports. The ontology module infers about 62 typed relations (`exposed_to_internet`, `assumes_role`, `encrypted_by`, `hosted_in_vpc` and so on) and writes RDF you can query with SPARQL. The RAG exporter chunks the graph three ways (per asset, per Louvain community, per relation domain) into JSONL for retrieval pipelines. The Terraform exporter maps 25+ asset types to `.tf.json` resources with an `import.sh` to adopt them into state. +The same inventory feeds three other exports. The ontology module infers 64 typed relations (`INTERNET_REACHABLE`, `ROLE_ASSUMES_ROLE`, `ENCRYPTED_BY_KMS`, `VPC_CONTAINS_SUBNET` and so on) and writes RDF you can query with SPARQL. The RAG exporter chunks the graph three ways (per asset, per Louvain community, per relation domain) into JSONL for retrieval pipelines. The Terraform exporter maps 25+ asset types to `.tf.json` resources with an `import.sh` to adopt them into state. -Scanner findings are deduplicated in two passes — within a scanner by (scanner, check ID, resource), and across scanners only when both the normalised title and the underlying check semantics (`rules/check_equivalence.yaml`) match — then rescored against CVSS and mapped to compliance controls. +Scanner findings are deduplicated in two passes (within a scanner by scanner, check ID and resource; across scanners only when both the normalised title and the underlying check semantics in `cloudg/rules/check_equivalence.yaml` match), then rescored against CVSS and mapped to compliance controls. --- @@ -313,11 +340,13 @@ Much more of cloudg is public, importable API than the CLI suggests. The [Python |---|---| | `cloudg.inventory.InventoryMapper` / `RelationshipLinker` | scanner-free inventory maps and metadata-derived relationship edges | | `cloudg.graph.builder.GraphBuilder` | NetworkX graph, attack paths, centrality/blast-radius metrics, D3/Cytoscape/GraphML export | -| `cloudg.graph.ontology.CloudOntology` | RDF ontology (~62 typed relations), SPARQL-queryable, Turtle/JSON-LD | +| `cloudg.graph.ontology.CloudOntology` | RDF ontology (64 typed relations), SPARQL-queryable, Turtle/JSON-LD | | `cloudg.graph.rag_export.RAGExporter` | retrieval-ready JSONL chunks of the infrastructure for LLM pipelines | | `cloudg.renderers.terraform_export.TerraformExporter` | `.tf.json` recreation of live infrastructure plus `import.sh` | | `cloudg.ingest.parse_report` and the scanner classes | every scanner's parser, usable standalone | | `cloudg.normaliser.FindingsNormaliser` | cross-scanner dedupe, CVSS rescoring, compliance mapping | +| `cloudg.mcp.CloudGMCPLayer` | the MCP layer, in-process or mounted into your own MCP server | +| `cloudg.resilience` | shared rate limits, circuit breakers and retries (`call_with_resilience`, `@resilient`) for your own cloud calls | Custom collectors and scanners register through entry points, no core changes needed: @@ -332,7 +361,7 @@ mycloud = "my_package.collector:MyCollector" ```bash uv pip install -e ".[all,dev]" -pytest # 152 tests, moto-mocked AWS included +pytest # about 1,650 tests, moto-mocked AWS included ruff check cloudg/ tests/ uv build # wheel + sdist for PyPI ``` diff --git a/REMAINING_WORK.md b/REMAINING_WORK.md new file mode 100644 index 0000000..8b280b4 --- /dev/null +++ b/REMAINING_WORK.md @@ -0,0 +1,189 @@ +# Remaining work on feature-mcp-layer (PR #26) + +Snapshot of 2026-10-08. The branch is synced with origin/main at 368efce and carries the +undelivered fix/parallel-rule-edges commit. Delete this file before the PR merges. + +## Where things stand + +- The last commit is a work-in-progress checkpoint. Five fix groups were stopped partway through + their work, and their uncommitted edits were committed as they were. +- `ruff check cloudg tests` passes. `ruff format --check` fails on 56 files (the new MCP and + resilience code); CI's lint job fails until `ruff format cloudg tests` is run. +- Tests outside `tests/mcp`: 702 passed. +- `tests/mcp`: 815 passed, 16 failed, 6 skipped. The failures are tests that still assert the old + "did you mean" error text, which G2 had started moving into error `data`, and the half-split + `state` package: + - test_catalog_export_live_meta.py::test_terraform_preview + - test_catalog_findings.py::test_list_findings_projection_pagination_errors, test_findings_for_asset + - test_catalog_graph.py::test_neighbors_errors, test_find_paths + - test_catalog_inventory.py::test_find_assets_bad_filters, test_get_asset_not_found_suggests + - test_catalog_ontology.py::test_ontology_neighbourhood + - test_catalog_prompts.py::test_prompt_errors + - test_catalog_regressions.py::test_snapshot_refuses_existing_name + - test_catalog_resources.py::test_asset_templates + - test_catalog_workspace.py::test_select_unknown, test_snapshot_and_diff_tools, test_dataset_summary + - test_state.py::test_resolve_asset_suggests_close_matches, test_workspace_load_select_remove_and_events +- Codacy on PR #26 reports 99 new issues (2 critical, 14 high, 83 medium, mostly complexity) and + 24 duplicated blocks. CodeQL reports two ReDoS alerts. + +## Already done and pushed + +- Repaired main's merge leftovers: the pasted-together `tests/test_reachability.py` (now + `tests/test_ports.py` plus `tests/test_reachability.py`, 41 tests running again) and the + duplicate import in `cloudg/graph/reachability.py`. +- Delivered fix/parallel-rule-edges, which was never pushed: GraphBuilder merges parallel rule + edges, so SSH open to the world is reported again. +- Strict-profile leak fixes (org ids, dependency ids, ontology IRIs, RAG text, account keys). +- Pre-existing Codacy issues from main: duplicate CloudGConfig import, mixin stub + "assignment from no return", unpinned parliament in the Dockerfile. +- CI actions moved to their Node 24 releases; the uv cache now keys on pyproject.toml. +- Each parallel session's goal was checked against the merged code: all six fix sessions + delivered; nothing was overwritten. + +## Groups needed: 7 + +Five code groups and two documentation groups. Each group owns its files outright, so they can +run in parallel without touching each other's work. The lead commits and pushes each group as it +finishes. + +### G1. Graph core + +Files: cloudg/graph/{reachability,ports,builder,ontology,ontology_rules,rag_export}.py, +cloudg/inventory/dependencies.py, tests/test_{ports,reachability,graph_builder,ontology, +rag_export,normaliser}.py. Status: partly done, uncommitted work included in the checkpoint. + +- Finding ids are not stable across scans: they hash the random per-collection asset uuid. Hash + the ARN (or id) instead, for both exposure and open-port findings. +- Three different "is this the internet" checks disagree; Azure NSG sources `Internet`, `*` and + `Any` produce exposure findings but no open-port findings. One shared helper everywhere. +- Duplicated port and rule-edge constants across ports.py, reachability.py, builder.py and + ontology_rules.py (the ontology still uses the old `ports` list and `"0-65535"` string test). +- Public network-flow API (flow successors / flow graph) for the MCP layer to reuse. +- AssetIndex as the single asset matcher: ARN-tail tiers, optional case-insensitive names, a + finding resolver; the KMS key matcher scans every asset per encrypted asset (index it once). +- Containment relations check only the source type; check the target too. +- Target groups are reported as "unexpected internet-exposed resources". +- RAG lost ENCRYPTED_BY_KMS and VPC containment entirely. +- Relation count in docstrings (64). +- Test cleanups: private-method tests, a test of an impossible path, a duplicated normaliser + test, and missing interaction tests (multi-rule SG end to end, Azure Internet sources, id + stability, ontology triples for SG edges). +- Codacy complexity on these files. + +### G2. MCP catalog and workspace + +Files: cloudg/mcp/state.py (being split into the cloudg/mcp/state/ package), cloudg/mcp/catalog/* +except privacy.py, tests/mcp/{conftest.py,fixtures/*,test_state.py,test_catalog_*.py}. +Status: partly done; the state package split and error-message changes are half applied (the 16 +failing tests). + +- Codacy: every catalog `register()` is too complex because handlers are nested inside it; move + handlers to module level. Long and parameter-heavy functions (find_assets, organization_topology, + filter_assets, count_assets, get_edges, ingest_reports, diff_datasets). Files over 500 lines + (state.py, graph.py, inventory.py). `format` argument shadows a builtin (keep the wire name). + Unsafe yaml.load. B105 false positive. Duplicate depends_on/dependents bodies. +- CodeQL: the SPARQL UPDATE detection regex can backtrack exponentially; replace with a scanner. +- flow_graph reports false internet paths; use G1's network-flow API. +- Finding-to-asset matching disagrees with AssetIndex; use it. +- graph.py has its own port and internet parsing (misses Azure `*`, GCP all-ports, ICMP, egress). +- Security: allowed roots default to the working directory (can be `/`); load errors echo file + contents; unbounded rglob. +- Security (strict profile): "did you mean" suggestions and handler errors echo real names; + resources and prompts return pre-serialised text that skips key-based redaction (prompt data + goes through the agreed `cloudg/data` + `cloudg/render` meta marker); textual graph and + ontology exports, subgraph_export and sparql_query become RESTRICTED. +- Lock nesting lets one slow ontology build stall every dataset. +- reachability_findings overrides main's deterministic ids; cursors survive dataset replacement; + trivy image arguments starting with `-`; SPARQL has no cost limit. +- explain_edge_type relation text is stale; blast_radius output order is random. + +### G3. MCP privacy + +Files: cloudg/mcp/policy.py, cloudg/mcp/transforms/*, cloudg/mcp/policies/*.yaml, +cloudg/mcp/catalog/privacy.py, tests/mcp/{test_transforms,test_policy,test_policy_regressions, +test_privacy_tools}.py. Status: partly done. + +- CodeQL/semgrep critical: vault `_SHAPE_RE` can backtrack exponentially; check the other + transform regexes too. +- Codacy: B105 false positives on label strings, an `assert`, a silent try/except, and complexity + or length in policy.py, detectors.py, redaction.py, vault.py, projection.py, annotation.py and + privacy.py; duplicated branches in projection._walk. +- Security: record restored (real, token) pairs so every outgoing string of a call can be + re-pseudonymised (errors were a reveal oracle); redact pre-serialised text using the vault's + known values; let deny/allow rules match the concrete resource URI; preview_transform must not + confirm pseudonym guesses; privacy_audit_log needs an admin gate; escape CR/LF in audit lines. +- Performance: strict mode on 50k assets takes about 21 s cold. + +### G4. MCP core, layer and servers + +Files: cloudg/mcp/{__init__,__main__,core,layer,schema,context,middleware,server,cli}.py, +cloudg/mcp/adapters/*, cloudg/mcp/native/*, the mcp lines of cloudg/cli.py, +tests/mcp/{_adapter_testkit.py,test_adapters_*,test_native_*,test_middleware,test_cli_mcp}.py. +Status: barely started (new content.py and layer_output.py modules from the core and layer split). + +- Codacy critical: `importlib.import_module` on the `--registry` value; validate it. +- Codacy: an `assert`, silent try/except blocks, many long or parameter-heavy functions + (server.py serve paths, cli.py serve_cmd, native http/protocol/stdio, fastmcp adapter), files + over 500 lines (core.py, layer.py, native/http.py, native/protocol.py). +- Duplication: protocol tools/call vs prompts/get, http GET vs DELETE and body parsing, adapter + notify and principal extraction, change dispatch, fastmcp run_tool vs render, Registry + resource vs resource_template. +- Security: transform errors from resources, prompts and completions; never forward raw handler + exception text; re-pseudonymise restored values; apply the `cloudg/data` marker; transform + link titles and embedded resource URIs; send progress and log messages through the policy; + DNS rebinding on non-loopback binds; an empty `--auth-token` must fail; canonicalise resource + URIs before policy checks; per-principal session caps, a connection cap and header timeout; + session binding on the SDK and fastmcp flavors. + +### G5. Resilience and older main issues + +Files: cloudg/resilience/*, cloudg/retry.py, cloudg/config.py (rate-limit parts), +cloudg/collectors/{gcp,azure,multi}.py, cloudg/credentials.py, cloudg/inventory/{linker,mapper, +mapper_result,gcp_deep}.py, cloudg/inventory/azure_graph/query.py, cloudg/api.py, +cloudg/renderers/svg_layout.py, cloudg/cli_run_helpers.py, cloudg/cli_commands.py, +tests/{test_resilience,conftest,test_api,test_collectors}.py, tests/test_inventory*.py. +Status: partly done (new linker_rules.py, resilience refactors). + +- High: a half-open circuit breaker can stay stuck forever after a fatal error, cancellation or + deadline (process-wide, every call in that service scope rejected). Timestamp the probe and + treat a fatal provider answer as a health signal. +- High: a tripped GCP Cloud Asset Inventory breaker never recovers on multi-page listings (page 2 + of the probe is rejected by its own breaker). Rate-limit only in `paced()`; never count calls + that were rejected before sending as failures (same bug in the Azure collector). +- `with_retry` now retries far more than before (any HTTP 500, text matches, its own circuit and + budget errors); restrict it. +- Instrumented boto3 sessions raise errors that are not botocore exceptions (breaks + `except ClientError` in embedders). +- A hung live operation pins its scope as busy forever; add a timeout. +- The process-wide retry budget gates botocore's own retries for the life of the process. +- Governor reconfiguration cross-talk between concurrent runs; bulkhead slots on cancelled + aioboto3 calls; sync hooks sleeping on the event loop; nested Resource Graph retries. +- mapper.deduplicate collapses parallel security group rule edges whenever an asset was merged. +- Codacy: silent try/except blocks, Azure policy class complexity 33, functions with 10 + parameters (guard, call_with_resilience), duplicated needs_retry, an unused re-export in + cloudg/retry.py, one issue in collectors/gcp.py, and main's older findings + (linker._link_rules complexity 28, api.py 612 lines, svg_layout and cli_run_helpers parameter + counts). + +### D1. Documentation refresh (after G1 to G5 land) + +docs/MCP.md, MCP_TOOLS.md, MCP_PRIVACY.md, MCP_INTERNALS.md, RESILIENCE.md, DOCUMENTATION.md, +INVENTORY_REFERENCE.md, INVENTORY_CATALOG.md, INVENTORY_INTERNALS.md, docs/index.html, README.md. + +- Regenerate every MCP_TOOLS.md example from the final code (the reachability, ontology and RAG + examples went stale with main's fixes) and fix the prose that describes the old behaviour. +- PyPI page: eight relative links that 404 on PyPI; reachability, relation count, Finding.id and + is_internet_exposed wording. +- Inventory docs: graph exports are no longer one-to-one with edges, port and protocol formats per + provider, the throttling field, AWS CONTAINS edges, GCP internet edges, 0.5.x references. +- MCP_PRIVACY.md and RESILIENCE.md: fixed known gaps, new detectors and key rules, UUID ids kept + under strict, small resilience inaccuracies, plus whatever G3 and G5 change. +- Handbook changelog entry for main's merged fixes. + +### D2. CHANGELOG rewrite and dash sweep (lead, last) + +- Rewrite the 0.6.0 CHANGELOG entry as one coherent release note: it is currently three entries + stuck together with a stray `=======` marker that renders a paragraph as a heading, and it says + "all public API changes are additive" while listing behaviour changes. +- Remove the 140 remaining em dashes in older code, help text, CLI output and rule files. +- Run `ruff format cloudg tests` and the full CI-equivalent check, then push. diff --git a/cloudg/__init__.py b/cloudg/__init__.py index 8380bcf..22da0bb 100644 --- a/cloudg/__init__.py +++ b/cloudg/__init__.py @@ -1,6 +1,6 @@ -"""cloudg — cloud graphing: multi-cloud mapping and security intelligence.""" +"""cloudg (cloud graphing): multi-cloud mapping and security intelligence.""" -__version__ = "0.5.3" +__version__ = "0.6.0" from cloudg.api import ( AnalysisResult, diff --git a/cloudg/api.py b/cloudg/api.py index cf303f0..06f9f99 100644 --- a/cloudg/api.py +++ b/cloudg/api.py @@ -1,4 +1,4 @@ -"""CloudG programmatic API — integration-ready entry point. +"""CloudG programmatic API: the integration-ready entry point. Designed for embedding CloudG in larger systems. Provides: - Single `CloudGEngine` class for full pipeline control @@ -40,35 +40,18 @@ ScanResult, ) -logger = logging.getLogger(__name__) - - -def resolve_iac_dirs( - iac_dir: str | None, - config_dirs: list[str] | None, - terraform_dir: str | Path | None = None, -) -> tuple[list[str], str | None]: - """Resolve which directories the IaC scanners (Checkov, Trivy fs) target. +# resolve_iac_dirs lives with the scanner runs; re-exported here (public API) +from cloudg.api_scanners import ScannerRunsMixin, resolve_iac_dirs - Priority: explicit dir -> configured dirs -> the generated Terraform - recreation of the live infrastructure. There is deliberately no fallback - to "." any more: scanning whatever directory cloudg happens to run from - is not a scan of the cloud, and its zero findings read like a clean bill - of health. +__all__ = [ + "AnalysisResult", + "CloudGEngine", + "CollectionResult", + "PipelineResult", + "resolve_iac_dirs", +] - Returns: - (dirs, source) where source is "cli", "config", "terraform", or - None when there is nothing for the IaC scanners to target. - """ - if iac_dir: - return [iac_dir], "cli" - if config_dirs: - return list(config_dirs), "config" - if terraform_dir: - tf_path = Path(terraform_dir) - if tf_path.is_dir() and any(tf_path.glob("*.tf.json")): - return [str(tf_path)], "terraform" - return [], None +logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- @@ -105,7 +88,7 @@ class AnalysisResult: @dataclass class PipelineResult: - """Complete pipeline output — single object for downstream consumption. + """Complete pipeline output: a single object for downstream consumption. This is the primary return type for `CloudGEngine.run_pipeline()`. Designed for easy serialisation and integration with larger systems. @@ -184,7 +167,7 @@ def to_summary(self) -> dict[str, Any]: # --------------------------------------------------------------------------- -class CloudGEngine: +class CloudGEngine(ScannerRunsMixin): """Integration-ready programmatic API for CloudG. Provides fine-grained control over the pipeline for embedding in @@ -199,7 +182,7 @@ class CloudGEngine: def __init__(self, config: CloudGConfig) -> None: self.config = config - # Event hooks — set these before calling run_pipeline() + # Event hooks: set these before calling run_pipeline() self.on_collection_complete: OnCollectionComplete | None = None self.on_scan_complete: OnScanComplete | None = None self.on_finding: OnFinding | None = None @@ -272,7 +255,7 @@ async def map_inventory( findings: list[Finding] | None = None, tagging_sweep: bool | None = None, ) -> "Any": - """Map the complete infrastructure inventory — no scanners involved. + """Map the complete infrastructure inventory; no scanners are involved. Runs the deep inventory collectors (full network fabric plus catch-all sweeps: AWS Resource Groups Tagging API, Azure ARM @@ -371,204 +354,6 @@ async def scan( return all_findings - def _run_reachability_analysis( - self, assets: list[CloudAsset], edges: list[NetworkEdge] - ) -> list[Finding]: - """Graph-based reachability analysis for the scanning phase.""" - try: - from cloudg.graph.builder import GraphBuilder - from cloudg.graph.reachability import ReachabilityAnalyzer - - builder = GraphBuilder() - graph = builder.build(assets, edges) - analyzer = ReachabilityAnalyzer(graph) - return analyzer.generate_findings() - except Exception as exc: - logger.error("Graph analysis failed: %s", exc) - self._emit_error("graph_analysis", exc) - return [] - - def _resolve_scan_iac_dirs( - self, - iac_dir: str | None, - scanner_list: list[str], - assets: list[CloudAsset], - edges: list[NetworkEdge], - out: Path, - ) -> list[str]: - """Resolve the IaC directories the scanners should target. - - When nothing is configured but Terraform recreation is enabled, - generate it now and scan that: the IaC scanners then audit the live - infrastructure via its Terraform representation instead of an - unrelated local directory. - """ - tf_dir: str | None = None - if ( - not iac_dir - and not self.config.scanners.iac_directories - and self.config.terraform.enabled - ): - tf_dir = self.config.terraform.output_dir or str(out / "terraform") - try: - from cloudg.renderers.terraform_export import TerraformExporter - - TerraformExporter(output_dir=tf_dir).export(assets, edges) - except Exception as exc: - logger.error("Terraform export for IaC scanning failed: %s", exc) - self._emit_error("terraform", exc) - tf_dir = None - - resolved_iac_dirs, iac_source = resolve_iac_dirs( - iac_dir, self.config.scanners.iac_directories, tf_dir - ) - if iac_source == "terraform": - logger.info("IaC scanners target the Terraform recreation at %s", resolved_iac_dirs[0]) - elif not resolved_iac_dirs and "checkov" in scanner_list: - logger.warning( - "Skipping Checkov: no IaC directory configured. " - "Pass iac_dir, set scanners.iac_directories, or enable terraform " - "so the recreated infrastructure can be scanned." - ) - return resolved_iac_dirs - - def _run_prowler(self, prov: str, profile: str | None, out: Path) -> list[Finding]: - from cloudg.scanners.prowler import ProwlerScanner - - region = self.config.aws.regions[0] if self.config.aws.regions else None - scanner = ProwlerScanner( - provider=prov, - profile=profile, - output_dir=str(out / "prowler" / prov), - extra_args=self.config.scanners.prowler_extra_args or [], - aws_access_key_id=self.config.aws.access_key_id if prov == "aws" else None, - aws_secret_access_key=self.config.aws.secret_access_key if prov == "aws" else None, - aws_region=region if prov == "aws" else None, - ) - return scanner.run() - - def _run_scoutsuite(self, prov: str, profile: str | None, out: Path) -> list[Finding]: - from cloudg.scanners.scoutsuite import ScoutSuiteScanner - - scanner = ScoutSuiteScanner( - provider=prov, - profile=profile if prov == "aws" else None, - report_dir=str(out / "scoutsuite" / prov), - extra_args=self.config.scanners.scoutsuite_extra_args or [], - ) - return scanner.run() - - def _run_checkov(self, target_dir: str) -> list[Finding]: - from cloudg.scanners.checkov import CheckovScanner - - scanner = CheckovScanner( - target_dir=target_dir, - frameworks=self.config.scanners.checkov_frameworks or None, - extra_args=self.config.scanners.checkov_extra_args or [], - ) - return scanner.run() - - def _run_trivy(self, image_list: list[str]) -> list[Finding]: - from cloudg.scanners.trivy import TrivyScanner - - scanner = TrivyScanner(extra_args=self.config.scanners.trivy_extra_args or []) - return scanner.scan_images(image_list) - - def _run_iam_linter(self, assets: list[CloudAsset]) -> list[Finding]: - from cloudg.scanners.iam_linter import IAMLinter - - linter = IAMLinter() - return linter.analyze_policies(assets) - - def _submit_scanner_jobs( - self, - executor: Any, - scanner_list: list[str], - assets: list[CloudAsset], - resolved_iac_dirs: list[str], - resolved_images: list[str], - profile: str | None, - out: Path, - ) -> dict[Any, str]: - """Submit one job per enabled scanner target; returns future -> name.""" - future_to_name: dict[Any, str] = {} - - if "prowler" in scanner_list: - for prov in self.config.providers: - future_to_name[executor.submit(self._run_prowler, prov, profile, out)] = ( - f"prowler-{prov}" - ) - - if "scoutsuite" in scanner_list: - for prov in self.config.providers: - future_to_name[executor.submit(self._run_scoutsuite, prov, profile, out)] = ( - f"scoutsuite-{prov}" - ) - - if "checkov" in scanner_list: - for d in resolved_iac_dirs: - future_to_name[executor.submit(self._run_checkov, d)] = f"checkov-{d}" - - if "trivy" in scanner_list and resolved_images: - future_to_name[executor.submit(self._run_trivy, resolved_images)] = "trivy" - - # IAM linter always runs if assets exist - if assets: - future_to_name[executor.submit(self._run_iam_linter, assets)] = "iam" - - return future_to_name - - def _run_scanners_parallel( - self, - scanner_list: list[str], - assets: list[CloudAsset], - resolved_iac_dirs: list[str], - resolved_images: list[str], - profile: str | None, - out: Path, - ) -> list[Finding]: - """Run all enabled scanners concurrently and collect their findings.""" - import concurrent.futures - - scanner_findings: list[Finding] = [] - max_workers = len(scanner_list) + len(self.config.providers) + 1 - with concurrent.futures.ThreadPoolExecutor(max_workers=max(max_workers, 2)) as executor: - future_to_name = self._submit_scanner_jobs( - executor, scanner_list, assets, resolved_iac_dirs, resolved_images, profile, out - ) - - for future in concurrent.futures.as_completed(future_to_name): - name = future_to_name[future] - try: - findings = future.result(timeout=self.config.scanners.timeout_seconds) - scanner_findings.extend(findings) - logger.info("[%s] %d findings", name, len(findings)) - except concurrent.futures.TimeoutError: - logger.error( - "[%s] timed out after %ds", name, self.config.scanners.timeout_seconds - ) - self._emit_error(name, TimeoutError(f"{name} timed out")) - except Exception as exc: - logger.error("[%s] failed: %s", name, exc) - self._emit_error(name, exc) - - return scanner_findings - - def _emit_scan_results(self, all_findings: list[Finding]) -> None: - """Emit per-finding and scan-complete callbacks.""" - if self.on_finding: - for f in all_findings: - try: - self.on_finding(f) - except Exception: - logger.debug("Event hook raised; ignoring", exc_info=True) - - if self.on_scan_complete: - try: - self.on_scan_complete(all_findings) - except Exception: - logger.debug("Event hook raised; ignoring", exc_info=True) - # ------------------------------------------------------------------ # Ingest: use existing scanner outputs instead of running scanners # ------------------------------------------------------------------ @@ -592,18 +377,7 @@ def ingest_reports(self, reports: dict[str, list[str | Path]]) -> list[Finding]: from cloudg.ingest import ingest_reports as _ingest findings = _ingest(reports) - - if self.on_finding: - for f in findings: - try: - self.on_finding(f) - except Exception: - logger.debug("Event hook raised; ignoring", exc_info=True) - if self.on_scan_complete: - try: - self.on_scan_complete(findings) - except Exception: - logger.debug("Event hook raised; ignoring", exc_info=True) + self._emit_scan_results(findings) return findings def normalise_findings( @@ -627,13 +401,13 @@ async def run_from_reports( reports: dict[str, list[str | Path]], output_dir: str | Path = "./reports", ) -> PipelineResult: - """Run the cloudg pipeline on existing scanner outputs — no cloud - access and no scanner binaries needed. + """Run the cloudg pipeline on existing scanner outputs. - Phases: ingest -> normalise -> reports (JSON + HTML). Collection - does not run, so graph/ontology/Terraform outputs that need live - assets are empty; combine with `collect()` + `analyze()` when - cloud credentials are available. + Needs no cloud access and no scanner binaries. Phases: ingest, + normalise, reports (JSON + HTML). Collection does not run, so + graph/ontology/Terraform outputs that need live assets are empty; + combine with `collect()` + `analyze()` when cloud credentials are + available. Returns: PipelineResult with findings, scan_result, and report paths. @@ -652,27 +426,10 @@ async def run_from_reports( scan_result = self.normalise_findings(findings) all_findings = scan_result.findings except Exception as exc: - logger.error("Normalisation failed: %s", exc) - errors.append(f"normalisation: {exc}") - self._emit_error("normalisation", exc) + self._record_phase_error("normalisation", "Normalisation failed", exc, errors) self._emit_phase_start("reporting") - report_paths: dict[str, Path] = {} - if scan_result: - try: - from cloudg.renderers.json_export import JSONExporter - - exporter = JSONExporter(output_dir=str(out)) - report_paths["json"] = Path(exporter.export(scan_result)) - - from cloudg.renderers.html_report import HTMLReportGenerator - - html_gen = HTMLReportGenerator(output_dir=str(out)) - report_paths["html"] = Path(html_gen.generate(scan_result)) - except Exception as exc: - logger.error("Report generation failed: %s", exc) - errors.append(f"reporting: {exc}") - self._emit_error("reporting", exc) + report_paths = self._write_reports(scan_result, out, errors) if scan_result else {} return PipelineResult( findings=all_findings, @@ -682,6 +439,37 @@ async def run_from_reports( errors=errors, ) + def _record_phase_error( + self, phase: str, label: str, exc: Exception, errors: list[str] + ) -> None: + """Log a failed pipeline phase, keep it in ``errors`` and emit it.""" + logger.error("%s: %s", label, exc) + errors.append(f"{phase}: {exc}") + self._emit_error(phase, exc) + + def _write_reports( + self, + scan_result: ScanResult, + out: Path, + errors: list[str], + graph_json: dict[str, Any] | None = None, + ) -> dict[str, Path]: + """Write the JSON and HTML reports; a failure is recorded in ``errors``.""" + report_paths: dict[str, Path] = {} + try: + from cloudg.renderers.json_export import JSONExporter + + exporter = JSONExporter(output_dir=str(out)) + report_paths["json"] = Path(exporter.export(scan_result, graph_json=graph_json)) + + from cloudg.renderers.html_report import HTMLReportGenerator + + html_gen = HTMLReportGenerator(output_dir=str(out)) + report_paths["html"] = Path(html_gen.generate(scan_result, graph_json=graph_json)) + except Exception as exc: + self._record_phase_error("reporting", "Report generation failed", exc, errors) + return report_paths + def run_from_reports_sync( self, reports: dict[str, list[str | Path]], @@ -840,12 +628,9 @@ def _analyze_terraform( async def run_pipeline(self, output_dir: str | Path = "./reports") -> PipelineResult: """Run the complete CloudG pipeline. - Phases: - 1. Collection (multi-provider, multi-region) - 2. Scanning (reachability, IAM linting) - 3. Analysis (ontology, RAG, Terraform) - 4. Normalisation - 5. Report generation + Phases: 1. collection (multi-provider, multi-region), 2. scanning + (reachability, IAM linting), 3. analysis (ontology, RAG, Terraform), + 4. normalisation, 5. report generation. Returns: PipelineResult with all outputs. @@ -877,20 +662,27 @@ async def run_pipeline(self, output_dir: str | Path = "./reports") -> PipelineRe edges=collection.edges, findings=all_findings, scan_result=scan_result, - graph_nodes=analysis.graph_nodes, - graph_edges=analysis.graph_edges, - ontology_triples=analysis.ontology_triples, - rag_chunks_path=analysis.rag_chunks_path, - terraform_paths=analysis.terraform_paths, - attack_paths=analysis.attack_paths, - providers_scanned=collection.providers_scanned, - regions_scanned=collection.regions_scanned, - coverage=collection.coverage, report_paths=report_paths, duration_ms=int((time.time() - start) * 1000), errors=errors, + **self._pipeline_outputs(collection, analysis), ) + @staticmethod + def _pipeline_outputs(collection: CollectionResult, analysis: AnalysisResult) -> dict[str, Any]: + """Collection coverage and analysis outputs carried into a PipelineResult.""" + return { + "graph_nodes": analysis.graph_nodes, + "graph_edges": analysis.graph_edges, + "ontology_triples": analysis.ontology_triples, + "rag_chunks_path": analysis.rag_chunks_path, + "terraform_paths": analysis.terraform_paths, + "attack_paths": analysis.attack_paths, + "providers_scanned": collection.providers_scanned, + "regions_scanned": collection.regions_scanned, + "coverage": collection.coverage, + } + def _normalise_pipeline_findings( self, collection: CollectionResult, @@ -915,9 +707,7 @@ def _normalise_pipeline_findings( scan_result.edges = collection.edges all_findings = scan_result.findings except Exception as exc: - logger.error("Normalisation failed: %s", exc) - errors.append(f"normalisation: {exc}") - self._emit_error("normalisation", exc) + self._record_phase_error("normalisation", "Normalisation failed", exc, errors) return scan_result, all_findings def _generate_pipeline_reports( @@ -929,29 +719,18 @@ def _generate_pipeline_reports( ) -> dict[str, Path]: """Generate JSON and HTML reports (with graph JSON) for the pipeline.""" self._emit_phase_start("reporting") - report_paths: dict[str, Path] = {} - if scan_result: - try: - from cloudg.graph.builder import GraphBuilder - - builder = GraphBuilder() - builder.build(collection.assets, collection.edges) - graph_json = builder.to_d3_json() - - from cloudg.renderers.json_export import JSONExporter - - exporter = JSONExporter(output_dir=str(out)) - report_paths["json"] = Path(exporter.export(scan_result, graph_json=graph_json)) - - from cloudg.renderers.html_report import HTMLReportGenerator + if not scan_result: + return {} + try: + from cloudg.graph.builder import GraphBuilder - html_gen = HTMLReportGenerator(output_dir=str(out)) - report_paths["html"] = Path(html_gen.generate(scan_result, graph_json=graph_json)) - except Exception as exc: - logger.error("Report generation failed: %s", exc) - errors.append(f"reporting: {exc}") - self._emit_error("reporting", exc) - return report_paths + builder = GraphBuilder() + builder.build(collection.assets, collection.edges) + graph_json = builder.to_d3_json() + except Exception as exc: + self._record_phase_error("reporting", "Report generation failed", exc, errors) + return {} + return self._write_reports(scan_result, out, errors, graph_json=graph_json) # ------------------------------------------------------------------ # Sync wrapper diff --git a/cloudg/api_scanners.py b/cloudg/api_scanners.py new file mode 100644 index 0000000..99a81d9 --- /dev/null +++ b/cloudg/api_scanners.py @@ -0,0 +1,255 @@ +"""Scanner runs of :class:`cloudg.api.CloudGEngine` (the ``scan`` phase). + +Split out of :mod:`cloudg.api` to keep that module readable; the engine +inherits these methods from :class:`ScannerRunsMixin`, and +:func:`resolve_iac_dirs` stays importable from ``cloudg.api``. +""" + +from __future__ import annotations + +import logging +from pathlib import Path +from typing import Any, Callable + +from cloudg.schema.models import CloudAsset, Finding, NetworkEdge + +__all__ = ["ScannerRunsMixin", "resolve_iac_dirs"] + +logger = logging.getLogger(__name__) + + +def resolve_iac_dirs( + iac_dir: str | None, + config_dirs: list[str] | None, + terraform_dir: str | Path | None = None, +) -> tuple[list[str], str | None]: + """Resolve which directories the IaC scanners (Checkov, Trivy fs) target. + + Priority: explicit dir -> configured dirs -> the generated Terraform + recreation of the live infrastructure. There is deliberately no fallback + to "." any more: scanning whatever directory cloudg happens to run from + is not a scan of the cloud, and its zero findings read like a clean bill + of health. + + Returns: + (dirs, source) where source is "cli", "config", "terraform", or + None when there is nothing for the IaC scanners to target. + """ + if iac_dir: + return [iac_dir], "cli" + if config_dirs: + return list(config_dirs), "config" + if terraform_dir: + tf_path = Path(terraform_dir) + if tf_path.is_dir() and any(tf_path.glob("*.tf.json")): + return [str(tf_path)], "terraform" + return [], None + + +class ScannerRunsMixin: + """Reachability analysis, IaC target resolution and the parallel scanner + runs; mixed into CloudGEngine, which provides ``config``, the event + hooks and ``_emit_error``.""" + + config: Any + on_finding: Callable[[Finding], None] | None + on_scan_complete: Callable[[list[Finding]], None] | None + _emit_error: Callable[[str, Exception], None] + + def _run_reachability_analysis( + self, assets: list[CloudAsset], edges: list[NetworkEdge] + ) -> list[Finding]: + """Graph-based reachability analysis for the scanning phase.""" + try: + from cloudg.graph.builder import GraphBuilder + from cloudg.graph.reachability import ReachabilityAnalyzer + + builder = GraphBuilder() + graph = builder.build(assets, edges) + analyzer = ReachabilityAnalyzer(graph) + return analyzer.generate_findings() + except Exception as exc: + logger.error("Graph analysis failed: %s", exc) + self._emit_error("graph_analysis", exc) + return [] + + def _resolve_scan_iac_dirs( + self, + iac_dir: str | None, + scanner_list: list[str], + assets: list[CloudAsset], + edges: list[NetworkEdge], + out: Path, + ) -> list[str]: + """Resolve the IaC directories the scanners should target. + + When nothing is configured but Terraform recreation is enabled, + generate it now and scan that: the IaC scanners then audit the live + infrastructure via its Terraform representation instead of an + unrelated local directory. + """ + tf_dir: str | None = None + if ( + not iac_dir + and not self.config.scanners.iac_directories + and self.config.terraform.enabled + ): + tf_dir = self.config.terraform.output_dir or str(out / "terraform") + try: + from cloudg.renderers.terraform_export import TerraformExporter + + TerraformExporter(output_dir=tf_dir).export(assets, edges) + except Exception as exc: + logger.error("Terraform export for IaC scanning failed: %s", exc) + self._emit_error("terraform", exc) + tf_dir = None + + resolved_iac_dirs, iac_source = resolve_iac_dirs( + iac_dir, self.config.scanners.iac_directories, tf_dir + ) + if iac_source == "terraform": + logger.info("IaC scanners target the Terraform recreation at %s", resolved_iac_dirs[0]) + elif not resolved_iac_dirs and "checkov" in scanner_list: + logger.warning( + "Skipping Checkov: no IaC directory configured. " + "Pass iac_dir, set scanners.iac_directories, or enable terraform " + "so the recreated infrastructure can be scanned." + ) + return resolved_iac_dirs + + def _run_prowler(self, prov: str, profile: str | None, out: Path) -> list[Finding]: + from cloudg.scanners.prowler import ProwlerScanner + + region = self.config.aws.regions[0] if self.config.aws.regions else None + scanner = ProwlerScanner( + provider=prov, + profile=profile, + output_dir=str(out / "prowler" / prov), + extra_args=self.config.scanners.prowler_extra_args or [], + aws_access_key_id=self.config.aws.access_key_id if prov == "aws" else None, + aws_secret_access_key=self.config.aws.secret_access_key if prov == "aws" else None, + aws_region=region if prov == "aws" else None, + ) + return scanner.run() + + def _run_scoutsuite(self, prov: str, profile: str | None, out: Path) -> list[Finding]: + from cloudg.scanners.scoutsuite import ScoutSuiteScanner + + scanner = ScoutSuiteScanner( + provider=prov, + profile=profile if prov == "aws" else None, + report_dir=str(out / "scoutsuite" / prov), + extra_args=self.config.scanners.scoutsuite_extra_args or [], + ) + return scanner.run() + + def _run_checkov(self, target_dir: str) -> list[Finding]: + from cloudg.scanners.checkov import CheckovScanner + + scanner = CheckovScanner( + target_dir=target_dir, + frameworks=self.config.scanners.checkov_frameworks or None, + extra_args=self.config.scanners.checkov_extra_args or [], + ) + return scanner.run() + + def _run_trivy(self, image_list: list[str]) -> list[Finding]: + from cloudg.scanners.trivy import TrivyScanner + + scanner = TrivyScanner(extra_args=self.config.scanners.trivy_extra_args or []) + return scanner.scan_images(image_list) + + def _run_iam_linter(self, assets: list[CloudAsset]) -> list[Finding]: + from cloudg.scanners.iam_linter import IAMLinter + + linter = IAMLinter() + return linter.analyze_policies(assets) + + def _submit_scanner_jobs( + self, + executor: Any, + scanner_list: list[str], + assets: list[CloudAsset], + resolved_iac_dirs: list[str], + resolved_images: list[str], + profile: str | None, + out: Path, + ) -> dict[Any, str]: + """Submit one job per enabled scanner target; returns future -> name.""" + future_to_name: dict[Any, str] = {} + + if "prowler" in scanner_list: + for prov in self.config.providers: + future_to_name[executor.submit(self._run_prowler, prov, profile, out)] = ( + f"prowler-{prov}" + ) + + if "scoutsuite" in scanner_list: + for prov in self.config.providers: + future_to_name[executor.submit(self._run_scoutsuite, prov, profile, out)] = ( + f"scoutsuite-{prov}" + ) + + if "checkov" in scanner_list: + for d in resolved_iac_dirs: + future_to_name[executor.submit(self._run_checkov, d)] = f"checkov-{d}" + + if "trivy" in scanner_list and resolved_images: + future_to_name[executor.submit(self._run_trivy, resolved_images)] = "trivy" + + # IAM linter always runs if assets exist + if assets: + future_to_name[executor.submit(self._run_iam_linter, assets)] = "iam" + + return future_to_name + + def _run_scanners_parallel( + self, + scanner_list: list[str], + assets: list[CloudAsset], + resolved_iac_dirs: list[str], + resolved_images: list[str], + profile: str | None, + out: Path, + ) -> list[Finding]: + """Run all enabled scanners concurrently and collect their findings.""" + import concurrent.futures + + scanner_findings: list[Finding] = [] + max_workers = len(scanner_list) + len(self.config.providers) + 1 + with concurrent.futures.ThreadPoolExecutor(max_workers=max(max_workers, 2)) as executor: + future_to_name = self._submit_scanner_jobs( + executor, scanner_list, assets, resolved_iac_dirs, resolved_images, profile, out + ) + + for future in concurrent.futures.as_completed(future_to_name): + name = future_to_name[future] + try: + findings = future.result(timeout=self.config.scanners.timeout_seconds) + scanner_findings.extend(findings) + logger.info("[%s] %d findings", name, len(findings)) + except concurrent.futures.TimeoutError: + logger.error( + "[%s] timed out after %ds", name, self.config.scanners.timeout_seconds + ) + self._emit_error(name, TimeoutError(f"{name} timed out")) + except Exception as exc: + logger.error("[%s] failed: %s", name, exc) + self._emit_error(name, exc) + + return scanner_findings + + def _emit_scan_results(self, all_findings: list[Finding]) -> None: + """Emit per-finding and scan-complete callbacks.""" + if self.on_finding: + for f in all_findings: + try: + self.on_finding(f) + except Exception: + logger.debug("Event hook raised; ignoring", exc_info=True) + + if self.on_scan_complete: + try: + self.on_scan_complete(all_findings) + except Exception: + logger.debug("Event hook raised; ignoring", exc_info=True) diff --git a/cloudg/cli.py b/cloudg/cli.py index f1269db..d6324c5 100644 --- a/cloudg/cli.py +++ b/cloudg/cli.py @@ -1,4 +1,4 @@ -"""CloudG CLI — Click-based command-line interface. +"""CloudG CLI: Click-based command-line interface. All terminal rendering goes through :mod:`cloudg.ui`, the Rich UI layer. """ @@ -25,6 +25,7 @@ _run_scan_jobs, ) from cloudg.cli_inventory import deps, map_inventory +from cloudg.mcp.cli import mcp_group from cloudg.ui import console # Global config reference (set by CLI group) @@ -65,10 +66,14 @@ def get_help(self, ctx: click.Context) -> str: @click.option("--log-file", default=None, help="Path to log file") @click.pass_context def cli(ctx: click.Context, verbose: bool, config_path: str | None, log_file: str | None) -> None: - """☁️ cloudg — cloud graphing: infrastructure mapping and security intelligence.""" + """☁️ cloudg (cloud graphing): infrastructure mapping and security intelligence.""" global _config from cloudg.config import load_config + if ctx.invoked_subcommand == "mcp": + # stdout is the MCP stdio protocol channel / machine-readable output + console.stderr = True + ctx.call_on_close(lambda: setattr(console, "stderr", False)) ui.print_banner(__version__) _config = load_config(config_path) # Expose the loaded config to subcommands defined outside this module @@ -83,6 +88,51 @@ def cli(ctx: click.Context, verbose: bool, config_path: str | None, log_file: st # ───────────────────────────────────────────────────────────────────── +def _build_collector( + provider: str, + resolver: Any, + *, + profile: str | None, + region: str, + subscription_id: str | None, + project_id: str | None, +) -> Any: + """Resolve credentials for ``provider`` and return its collector.""" + if provider == "aws": + from cloudg.collectors.aws import AsyncAWSCollector + + creds = resolver.resolve_aws(profile=profile, region=region) + return AsyncAWSCollector( + session=creds.session, region=creds.region, account_id=creds.account_id + ) + if provider == "azure": + from cloudg.collectors.azure import AzureCollector + + creds = resolver.resolve_azure(subscription_id=subscription_id) + return AzureCollector( + credential=creds.credential, + subscription_id=creds.subscription_id, + ) + if provider == "gcp": + from cloudg.collectors.gcp import GCPCollector + + creds = resolver.resolve_gcp(project_id=project_id) + return GCPCollector(project_id=creds.project_id, credentials=creds.credentials) + raise click.BadParameter(f"Unknown provider: {provider}") + + +async def _collect_inventory( + provider: str, resolver: Any, collector_kwargs: dict[str, Any] +) -> dict[str, Any]: + """Run the provider's collector and return its assets and edges as JSON data.""" + collector = _build_collector(provider, resolver, **collector_kwargs) + assets, edges = await collector.run() + return { + "assets": [a.model_dump(mode="json") for a in assets], + "edges": [e.model_dump(mode="json") for e in edges], + } + + @cli.command() @click.option( "-p", @@ -110,7 +160,7 @@ def collect( output: str, ) -> None: """Collect cloud assets from the specified provider.""" - ui.section(f"Asset Collection — {provider.upper()}") + ui.section(f"Asset Collection: {provider.upper()}") from cloudg.credentials import CredentialResolver @@ -118,41 +168,17 @@ def collect( output_dir = Path(output) output_dir.mkdir(parents=True, exist_ok=True) - async def _collect() -> dict[str, Any]: - if provider == "aws": - from cloudg.collectors.aws import AsyncAWSCollector - - creds = resolver.resolve_aws(profile=profile, region=region) - collector = AsyncAWSCollector( - session=creds.session, region=creds.region, account_id=creds.account_id - ) - elif provider == "azure": - from cloudg.collectors.azure import AzureCollector - - creds = resolver.resolve_azure(subscription_id=subscription_id) - collector = AzureCollector( - credential=creds.credential, - subscription_id=creds.subscription_id, - ) - elif provider == "gcp": - from cloudg.collectors.gcp import GCPCollector - - creds = resolver.resolve_gcp(project_id=project_id) - collector = GCPCollector(project_id=creds.project_id, credentials=creds.credentials) - else: - raise click.BadParameter(f"Unknown provider: {provider}") - - assets, edges = await collector.run() - return { - "assets": [a.model_dump(mode="json") for a in assets], - "edges": [e.model_dump(mode="json") for e in edges], - } - + collector_kwargs = { + "profile": profile, + "region": region, + "subscription_id": subscription_id, + "project_id": project_id, + } try: with console.status( f"[accent]Collecting assets from {provider.upper()}…[/]", spinner="dots" ): - result = asyncio.run(_collect()) + result = asyncio.run(_collect_inventory(provider, resolver, collector_kwargs)) except Exception as exc: ui.error_panel("Collection failed", exc) sys.exit(1) @@ -310,6 +336,37 @@ def report(input_file: str, output: str, fmt: str) -> None: # ───────────────────────────────────────────────────────────────────── +def _normalise_and_report_ingest( + all_findings: list[Any], tool_count: int, output: str, fmt: str +) -> None: + """Normalise ingested findings, then write the raw findings and the reports.""" + # Normalise: dedupe within and across scanners, score, map compliance + from cloudg.normaliser import FindingsNormaliser + + cfg = _config + normaliser = FindingsNormaliser(rules_dir=cfg.rulesets.rules_dir if cfg else None) + scan_result = normaliser.normalise(all_findings) + + output_dir = Path(output) + output_dir.mkdir(parents=True, exist_ok=True) + + # Save raw (pre-normalisation) findings alongside the reports + raw_path = output_dir / "raw-findings.json" + with open(raw_path, "w") as f: + dumped = [fi.model_dump(mode="json") for fi in all_findings] + json.dump(dumped, f, indent=2, default=str) + + _render_ingest_reports(scan_result, fmt, output_dir) + + ui.artifact("Raw findings", raw_path) + + console.print() + ui.success( + f"Ingested [metric]{len(all_findings)}[/] findings from {tool_count} tool(s), " + f"[metric]{len(scan_result.findings)}[/] after deduplication" + ) + + @cli.command() @click.option( "--prowler", @@ -356,10 +413,10 @@ def ingest( output: str, fmt: str, ) -> None: - """Aggregate existing scanner outputs — no scanners are executed. + """Aggregate existing scanner outputs; no scanners are executed. Feed cloudg the native output files of scans you already ran - (Prowler, ScoutSuite, Checkov, Trivy — any combination) and it + (Prowler, ScoutSuite, Checkov, Trivy, in any combination) and it normalises, deduplicates across scanners via the check-equivalence rulesets, maps compliance frameworks, and generates reports. @@ -385,41 +442,18 @@ def ingest( if not all_findings: ui.warn("No findings parsed from the given reports") - # Normalise: dedupe within and across scanners, score, map compliance - from cloudg.normaliser import FindingsNormaliser - - cfg = _config - normaliser = FindingsNormaliser(rules_dir=cfg.rulesets.rules_dir if cfg else None) - scan_result = normaliser.normalise(all_findings) - - output_dir = Path(output) - output_dir.mkdir(parents=True, exist_ok=True) - - # Save raw (pre-normalisation) findings alongside the reports - raw_path = output_dir / "raw-findings.json" - with open(raw_path, "w") as f: - dumped = [fi.model_dump(mode="json") for fi in all_findings] - json.dump(dumped, f, indent=2, default=str) - - _render_ingest_reports(scan_result, fmt, output_dir) - - ui.artifact("Raw findings", raw_path) - - console.print() - ui.success( - f"Ingested [metric]{len(all_findings)}[/] findings from {len(per_tool)} tool(s) — " - f"[metric]{len(scan_result.findings)}[/] after deduplication" - ) + _normalise_and_report_ingest(all_findings, len(per_tool), output, fmt) # ───────────────────────────────────────────────────────────────────── -# RUN command (full pipeline) — defined in cloudg.cli_commands -# MAP / DEPS commands (inventory mapping) — defined in cloudg.cli_inventory +# RUN command (full pipeline): defined in cloudg.cli_commands +# MAP / DEPS commands (inventory mapping): defined in cloudg.cli_inventory # ───────────────────────────────────────────────────────────────────── cli.add_command(run) cli.add_command(map_inventory) cli.add_command(deps) +cli.add_command(mcp_group) if __name__ == "__main__": diff --git a/cloudg/cli_commands.py b/cloudg/cli_commands.py index f8f21e8..6fc9534 100644 --- a/cloudg/cli_commands.py +++ b/cloudg/cli_commands.py @@ -7,7 +7,7 @@ from __future__ import annotations from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import Any import click @@ -22,14 +22,12 @@ _collect_assets, _export_rag_phase, _graph_phase, + RunProducts, _post_scan_phases, _scanner_phase, _terraform_phase, ) -if TYPE_CHECKING: - from cloudg.config import CloudGConfig - @click.command() @click.option( @@ -139,13 +137,38 @@ def run(ctx: click.Context, **kwargs: Any) -> None: _show_run_config(cfg, scanner_list, output_dir) resolved_images = _resolve_run_images(kwargs["images"], cfg) + products, tf_dir = _pre_scan_phases(cfg, kwargs, output_dir) + + # Phase 3: Security Scanning (all scanners in parallel) + scanner_findings, iam_findings = _scanner_phase( + cfg, + scanner_list, + kwargs["profile"], + kwargs["iac_dir"], + resolved_images, + tf_dir, + products.assets, + output_dir, + ) + + # Phases 3b to 5: ontology, RAG update, normalisation, reports, summary + _post_scan_phases(cfg, kwargs, products, scanner_findings, iam_findings, output_dir) + + +def _pre_scan_phases( + cfg: Any, kwargs: dict[str, Any], output_dir: Path +) -> tuple[RunProducts, str | None]: + """Phases 1 to 2d of ``cloudg run``: collection, graph, RAG and Terraform. + + Returns what those phases produced and the Terraform output directory. + """ # Phase 1: Asset Collection (always uses multi-provider orchestrator) assets, edges, coverage_records = _collect_assets(cfg) # Phase 2: Graph Analysis graph, graph_json, reachability_findings = _graph_phase(cfg, assets, edges, output_dir) - # Phase 2b: Semantic Ontology — deferred to after scanner phase + # Phase 2b: Semantic Ontology, deferred to after the scanner phase # (so security/compliance findings can be included in the ontology) # Phase 2c: RAG Export @@ -155,29 +178,7 @@ def run(ctx: click.Context, **kwargs: Any) -> None: # Phase 2d: Terraform Recreation tf_dir = _terraform_phase(cfg, kwargs["terraform"], assets, edges, output_dir) - # Phase 3: Security Scanning (all scanners in parallel) - scanner_findings, iam_findings = _scanner_phase( - cfg, - scanner_list, - kwargs["profile"], - kwargs["iac_dir"], - resolved_images, - tf_dir, - assets, - output_dir, - ) - - # Phases 3b–5: ontology, RAG update, normalisation, reports, summary - _post_scan_phases( - cfg, - kwargs, - assets, - edges, - graph, - graph_json, - reachability_findings, - scanner_findings, - iam_findings, - coverage_records, - output_dir, + products = RunProducts( + assets, edges, graph, graph_json, reachability_findings, coverage_records ) + return products, tf_dir diff --git a/cloudg/cli_helpers.py b/cloudg/cli_helpers.py index c25681a..cb07251 100644 --- a/cloudg/cli_helpers.py +++ b/cloudg/cli_helpers.py @@ -123,7 +123,7 @@ def _run_scan_jobs(jobs: list[tuple[str, str, Any]]) -> list[Any]: try: findings = future.result(timeout=3600) all_findings.extend(findings) - ui.task_done(progress, task_id, f"{name} — {len(findings)} findings") + ui.task_done(progress, task_id, f"{name}: {len(findings)} findings") except Exception as exc: ui.task_failed(progress, task_id, f"{name} failed: {exc}") return all_findings @@ -165,7 +165,7 @@ def _parse_ingest_reports(reports: dict[str, list[str]]) -> tuple[list[Any], dic try: findings = parse_report(tool, path) except (ValueError, FileNotFoundError) as exc: - ui.warn(f"{tool}: skipping {path} — {exc}") + ui.warn(f"{tool}: skipping {path} ({exc})") continue count += len(findings) all_findings.extend(findings) diff --git a/cloudg/cli_inventory.py b/cloudg/cli_inventory.py index 4145227..00ecc01 100644 --- a/cloudg/cli_inventory.py +++ b/cloudg/cli_inventory.py @@ -122,7 +122,7 @@ def _show_map_summary(summary: dict[str, Any]) -> None: "Accounts": org_summary["accounts"], "Organizational units": org_summary["ous"], "Control Tower": "yes" if org_summary["control_tower"] else "no", - "Governed regions": ", ".join(org_summary["governed_regions"]) or "—", + "Governed regions": ", ".join(org_summary["governed_regions"]) or "-", }, ) ui.stats_table( @@ -317,7 +317,7 @@ def _merge_findings( @click.option("-o", "--output", default="./reports", help="Output directory") @click.pass_context def map_inventory(ctx: click.Context, **kwargs: Any) -> None: - """Map the complete infrastructure inventory — no scanners involved. + """Map the complete infrastructure inventory, with no scanners involved. Deep-collects everything deployed (or default) across the configured providers: the network fabric, compute, containers (ECR, ECS, EKS and diff --git a/cloudg/cli_run_helpers.py b/cloudg/cli_run_helpers.py index cfd1c66..34699f3 100644 --- a/cloudg/cli_run_helpers.py +++ b/cloudg/cli_run_helpers.py @@ -10,6 +10,7 @@ import asyncio import json import logging +from dataclasses import dataclass, field from pathlib import Path from typing import TYPE_CHECKING, Any @@ -225,7 +226,7 @@ def _submit_provider_scanners( profile: str | None, output_dir: Path, ) -> None: - """Submit Prowler and ScoutSuite — one instance per provider.""" + """Submit Prowler and ScoutSuite, one instance per provider.""" if "prowler" in scanner_list: for prov in cfg.providers: if prov in ("aws", "azure", "gcp"): @@ -266,7 +267,7 @@ def _submit_checkov( return if not iac_dirs: ui.warn( - "Checkov: nothing to scan — pass --iac-dir, set " + "Checkov: nothing to scan; pass --iac-dir, set " "scanners.iac_directories, or enable --terraform to scan the " "recreated infrastructure" ) @@ -311,11 +312,11 @@ def _submit_trivy_and_iam( iac_dirs, ) else: - ui.warn("Trivy: nothing to scan — configure --images, --iac-dir, or enable --terraform") + ui.warn("Trivy: nothing to scan; configure --images, --iac-dir, or enable --terraform") else: ui.skip("Trivy: not enabled") - # IAM Linter — always runs internally to analyze collected assets + # IAM Linter: always runs internally to analyze collected assets if "iam" in scanner_list or assets: submit("IAM Linter", "IAM Lint", _scan_iam, assets) @@ -338,7 +339,7 @@ def _collect_scanner_results( iam_findings.extend(findings) else: scanner_findings.extend(findings) - ui.task_done(progress, task_id, f"{scanner_name} — {len(findings)} findings") + ui.task_done(progress, task_id, f"{scanner_name}: {len(findings)} findings") except concurrent.futures.TimeoutError: ui.task_failed( progress, @@ -349,27 +350,12 @@ def _collect_scanner_results( ui.task_failed(progress, task_id, f"{scanner_name} failed: {exc}") -def _scanner_phase( - cfg: CloudGConfig, - scanner_list: list[str], - profile: str | None, - iac_dir: str | None, - resolved_images: list[str], - tf_dir: str | None, - assets: list[Any], - output_dir: Path, -) -> tuple[list[Any], list[Any]]: - """Phase 3: run all enabled scanners in parallel. +def _resolve_scan_iac_dirs(cfg: CloudGConfig, iac_dir: str | None, tf_dir: str | None) -> list[str]: + """Resolve IaC scan targets: CLI flag, then config, then the Terraform recreation. - Returns (scanner findings, IAM linter findings). + There is no fallback to ".": scanning the directory cloudg runs from is + not a scan of the cloud, and its zero findings look like a clean result. """ - import concurrent.futures - - ui.phase("Phase 3 · Security Scanning", note="running scanners in parallel") - - # ── Resolve IaC scan targets: CLI flag → config → Terraform recreation ── - # No fallback to "." — scanning the directory cloudg runs from is not a - # scan of the cloud, and its zero findings look like a clean result. from cloudg.api import resolve_iac_dirs resolved_iac_dirs, iac_source = resolve_iac_dirs(iac_dir, cfg.scanners.iac_directories, tf_dir) @@ -378,6 +364,20 @@ def _scanner_phase( "IaC scanners target the Terraform recreation of the live " f"infrastructure ({resolved_iac_dirs[0]})" ) + return resolved_iac_dirs + + +def _run_scanners_parallel( + cfg: CloudGConfig, + scanner_list: list[str], + profile: str | None, + resolved_iac_dirs: list[str], + resolved_images: list[str], + assets: list[Any], + output_dir: Path, +) -> tuple[list[Any], list[Any]]: + """Submit every enabled scanner to a thread pool and gather their findings.""" + import concurrent.futures scanner_findings: list[Any] = [] iam_findings: list[Any] = [] @@ -403,6 +403,27 @@ def submit(name: str, description: str, fn: Any, *args: Any) -> None: return scanner_findings, iam_findings +def _scanner_phase( + cfg: CloudGConfig, + scanner_list: list[str], + profile: str | None, + iac_dir: str | None, + resolved_images: list[str], + tf_dir: str | None, + assets: list[Any], + output_dir: Path, +) -> tuple[list[Any], list[Any]]: + """Phase 3: run all enabled scanners in parallel. + + Returns (scanner findings, IAM linter findings). + """ + ui.phase("Phase 3 · Security Scanning", note="running scanners in parallel") + resolved_iac_dirs = _resolve_scan_iac_dirs(cfg, iac_dir, tf_dir) + return _run_scanners_parallel( + cfg, scanner_list, profile, resolved_iac_dirs, resolved_images, assets, output_dir + ) + + def _ontology_phase( cfg: CloudGConfig, assets: list[Any], @@ -465,23 +486,32 @@ def _render_run_reports( ui.artifact("HTML", html_path) +@dataclass +class RunProducts: + """What the collection and graph phases of ``cloudg run`` produced.""" + + assets: list[Any] + edges: list[Any] + graph: Any + graph_json: dict[str, Any] + reachability_findings: list[Any] + coverage_records: list[Any] = field(default_factory=list) + + def _post_scan_phases( cfg: CloudGConfig, kwargs: dict[str, Any], - assets: list[Any], - edges: list[Any], - graph: Any, - graph_json: dict[str, Any], - reachability_findings: list[Any], + products: RunProducts, scanner_findings: list[Any], iam_findings: list[Any], - coverage_records: list[Any], output_dir: Path, ) -> None: """Run everything after the scanner phase: ontology, RAG update, normalisation, report rendering and the results summary.""" from cloudg.normaliser import FindingsNormaliser + assets, edges = products.assets, products.edges + reachability_findings = products.reachability_findings # Combine all findings for downstream phases all_security_findings = scanner_findings + iam_findings + reachability_findings console.print() @@ -493,7 +523,7 @@ def _post_scan_phases( # Also update RAG export with all findings if kwargs["rag_export"] and cfg.rag.enabled: - _update_rag_export(cfg, assets, edges, graph, all_security_findings, output_dir) + _update_rag_export(cfg, assets, edges, products.graph, all_security_findings, output_dir) # Phase 4: Normalise (with external rulesets) ui.phase("Phase 4 · Normalisation") @@ -505,15 +535,15 @@ def _post_scan_phases( ui.success(f"[metric]{len(scan_result.findings)}[/] normalised findings") # Phase 5: Render - _render_run_reports(scan_result, graph_json, assets, edges, output_dir) + _render_run_reports(scan_result, products.graph_json, assets, edges, output_dir) # Summary ui.section("Results") ui.summary_table(scan_result.summary) # Coverage summary - if coverage_records: - ui.coverage_table(coverage_records) + if products.coverage_records: + ui.coverage_table(products.coverage_records) console.print() ui.success(f"All reports saved to [path]{output_dir}[/]") diff --git a/cloudg/collectors/__init__.py b/cloudg/collectors/__init__.py index eac0925..d7bb07f 100644 --- a/cloudg/collectors/__init__.py +++ b/cloudg/collectors/__init__.py @@ -1,4 +1,4 @@ -"""Collectors package — async cloud resource collectors.""" +"""Collectors package: async cloud resource collectors.""" from cloudg.collectors.base import BaseCollector diff --git a/cloudg/collectors/aws.py b/cloudg/collectors/aws.py index 2cf2187..4ba3fbe 100644 --- a/cloudg/collectors/aws.py +++ b/cloudg/collectors/aws.py @@ -1,4 +1,10 @@ -"""Async AWS resource collector using aioboto3 with adaptive retries.""" +"""Async AWS resource collector using aioboto3 with adaptive retries. + +Every aioboto3 session is wired into :mod:`cloudg.resilience` (shared +adaptive rate limits, circuit breakers, throttle telemetry); a service whose +calls stay throttled after retries is recorded FAILED / PARTIAL with reason +"throttled" instead of failing the run. +""" from __future__ import annotations @@ -11,6 +17,8 @@ from cloudg.collectors.aws_services_extended import ExtendedServiceCollectorsMixin from cloudg.collectors.base import BaseCollector from cloudg.coverage import CollectionCoverage, ServiceStatus +from cloudg.resilience.errors import describe_error +from cloudg.resilience.stats import throttle_ledger from cloudg.schema.models import ( AssetType, CloudAsset, @@ -21,22 +29,45 @@ logger = logging.getLogger(__name__) -def _get_aio_config() -> Any: - """Return AioConfig with adaptive retries.""" +def _retries(max_attempts: int | None = None, mode: str | None = None) -> dict[str, Any]: + """botocore ``retries`` from ``aws.max_retries`` / ``aws.retry_mode`` + (applied to the resilience governor by the orchestrator; adaptive / 10).""" + from cloudg.resilience.aws import botocore_retries + + retries = botocore_retries() + if max_attempts is not None: + retries["max_attempts"] = max_attempts + if mode is not None: + retries["mode"] = mode + return retries + + +def _get_aio_config(max_attempts: int | None = None, mode: str | None = None) -> Any: + """Return AioConfig with the configured retries (default: adaptive mode, 10 retries).""" + retries = _retries(max_attempts, mode) try: from aiobotocore.config import AioConfig return AioConfig( - retries={"mode": "adaptive", "max_attempts": 10}, + retries=retries, connect_timeout=10, read_timeout=30, ) except ImportError: from botocore.config import Config - return Config( - retries={"mode": "adaptive", "max_attempts": 10}, - ) + return Config(retries=retries) + + +def _default_client_config(retries: dict[str, Any] | None = None) -> Any: + """Session-wide default for clients created without ``config=``: the + configured retries only (botocore's default timeouts are kept).""" + try: + from aiobotocore.config import AioConfig + + return AioConfig(retries=dict(retries) if retries else _retries()) + except ImportError: # pragma: no cover (aiobotocore always ships with aioboto3) + return None class AsyncAWSCollector(CoreServiceCollectorsMixin, ExtendedServiceCollectorsMixin, BaseCollector): @@ -61,7 +92,13 @@ def __init__( self._region = region self._account_id = account_id self._aioboto3_session: Any = None - self._aio_config = _get_aio_config() + # Read the governor's retry settings once: a later reconfiguration + # (another run with a different config) must not change the retries + # of a collector that is already running. + self._retry_settings = _retries() + self._aio_config = _get_aio_config( + self._retry_settings["max_attempts"], self._retry_settings["mode"] + ) self.coverage = CollectionCoverage(provider="aws", region=region, account_id=account_id) def _get_aio_session(self) -> Any: @@ -92,6 +129,16 @@ def _get_aio_session(self) -> Any: session_kwargs["profile_name"] = self._boto3_session.profile_name self._aioboto3_session = aioboto3.Session(**session_kwargs) + # Every client of this session shares cloudg's adaptive rate + # limiter, circuit breakers and throttle telemetry. + from cloudg.resilience.aws import install_aws_hooks + + install_aws_hooks( + self._aioboto3_session, + account_id=self._account_id, + region=self._region, + default_config=_default_client_config(self._retry_settings), + ) return self._aioboto3_session # ------------------------------------------------------------------ @@ -106,54 +153,8 @@ async def _collect_sg_edges(self, assets: list[CloudAsset]) -> list[NetworkEdge] for sg in sg_assets: # Endpoint is the SG asset's id (like Azure NSG rule edges) so the # edge attaches to the SECURITY_GROUP node; the CIDR side stays external. - sg_id = sg.id - - # Ingress rules - for rule in sg.metadata.get("ingress_rules", []): - for ip_range in rule.get("IpRanges", []): - cidr = ip_range.get("CidrIp", "") - from_port = rule.get("FromPort", 0) - to_port = rule.get("ToPort", 65535) - protocol = rule.get("IpProtocol", "-1") - - port_range = ( - f"{from_port}-{to_port}" if from_port != to_port else str(from_port) - ) - - # Cap port list to avoid OOM on wide ranges (e.g., 0-65535) - port_count = min(to_port - from_port + 1, 100) - edges.append( - NetworkEdge( - source_id=cidr, - target_id=sg_id, - edge_type=EdgeType.SECURITY_GROUP_RULE, - ports=list(range(from_port, from_port + port_count)), - port_range=port_range, - protocol="ALL" if protocol == "-1" else protocol.upper(), - cidr=cidr, - direction="ingress", - ) - ) - - # Egress rules - for rule in sg.metadata.get("egress_rules", []): - for ip_range in rule.get("IpRanges", []): - cidr = ip_range.get("CidrIp", "") - from_port = rule.get("FromPort", 0) - to_port = rule.get("ToPort", 65535) - protocol = rule.get("IpProtocol", "-1") - - edges.append( - NetworkEdge( - source_id=sg_id, - target_id=cidr, - edge_type=EdgeType.SECURITY_GROUP_RULE, - port_range=f"{from_port}-{to_port}", - protocol="ALL" if protocol == "-1" else protocol.upper(), - cidr=cidr, - direction="egress", - ) - ) + edges.extend(_sg_ingress_edges(sg.id, sg.metadata.get("ingress_rules", []))) + edges.extend(_sg_egress_edges(sg.id, sg.metadata.get("egress_rules", []))) return edges @@ -184,20 +185,34 @@ async def _collect_containment_edges(self, assets: list[CloudAsset]) -> list[Net async def _run_service_collector(self, name: str, coro: Any) -> list[CloudAsset]: """Run a service collector with coverage tracking.""" start = time.time() - try: - result = await coro - duration_ms = int((time.time() - start) * 1000) + # The ledger collects throttling that per-item error handling inside + # the collector swallowed, so the service shows up PARTIAL + # ("throttled: ...") instead of a silently incomplete SUCCESS. + with throttle_ledger() as ledger: + try: + result = await coro + except Exception as exc: + duration_ms = int((time.time() - start) * 1000) + logger.error("Collector %s failed: %s", name, exc) + self.coverage.record( + name, ServiceStatus.FAILED, error=describe_error(exc), duration_ms=duration_ms + ) + return [] + duration_ms = int((time.time() - start) * 1000) + if ledger.degraded: + logger.warning("Collector %s: %s", name, ledger.describe()) self.coverage.record( - name, ServiceStatus.SUCCESS, asset_count=len(result), duration_ms=duration_ms + name, + ServiceStatus.PARTIAL, + asset_count=len(result), + error=ledger.describe(), + duration_ms=duration_ms, ) - return result - except Exception as exc: - duration_ms = int((time.time() - start) * 1000) - logger.error("Collector %s failed: %s", name, exc) + else: self.coverage.record( - name, ServiceStatus.FAILED, error=str(exc), duration_ms=duration_ms + name, ServiceStatus.SUCCESS, asset_count=len(result), duration_ms=duration_ms ) - return [] + return result def _service_tasks(self) -> dict[str, Any]: """Collector callables per service, invoked lazily by :meth:`collect`. @@ -267,3 +282,57 @@ async def collect_edges(self) -> list[NetworkEdge]: all_edges = sg_edges + containment_edges logger.info("Collected %d AWS edges", len(all_edges)) return all_edges + + +def _sg_rule_ports(rule: dict[str, Any]) -> tuple[int, int, str]: + """Return (from port, to port, protocol label) for one security group rule.""" + from_port = rule.get("FromPort", 0) + to_port = rule.get("ToPort", 65535) + protocol = rule.get("IpProtocol", "-1") + return from_port, to_port, "ALL" if protocol == "-1" else protocol.upper() + + +def _sg_ingress_edges(sg_id: str, rules: list[dict[str, Any]]) -> list[NetworkEdge]: + """Edges from each ingress rule's CIDR ranges into the security group.""" + edges: list[NetworkEdge] = [] + for rule in rules: + for ip_range in rule.get("IpRanges", []): + cidr = ip_range.get("CidrIp", "") + from_port, to_port, protocol = _sg_rule_ports(rule) + port_range = f"{from_port}-{to_port}" if from_port != to_port else str(from_port) + # Cap port list to avoid OOM on wide ranges (e.g., 0-65535) + port_count = min(to_port - from_port + 1, 100) + edges.append( + NetworkEdge( + source_id=cidr, + target_id=sg_id, + edge_type=EdgeType.SECURITY_GROUP_RULE, + ports=list(range(from_port, from_port + port_count)), + port_range=port_range, + protocol=protocol, + cidr=cidr, + direction="ingress", + ) + ) + return edges + + +def _sg_egress_edges(sg_id: str, rules: list[dict[str, Any]]) -> list[NetworkEdge]: + """Edges from the security group out to each egress rule's CIDR ranges.""" + edges: list[NetworkEdge] = [] + for rule in rules: + for ip_range in rule.get("IpRanges", []): + cidr = ip_range.get("CidrIp", "") + from_port, to_port, protocol = _sg_rule_ports(rule) + edges.append( + NetworkEdge( + source_id=sg_id, + target_id=cidr, + edge_type=EdgeType.SECURITY_GROUP_RULE, + port_range=f"{from_port}-{to_port}", + protocol=protocol, + cidr=cidr, + direction="egress", + ) + ) + return edges diff --git a/cloudg/collectors/aws_services.py b/cloudg/collectors/aws_services.py index 8ddd286..7e0e20d 100644 --- a/cloudg/collectors/aws_services.py +++ b/cloudg/collectors/aws_services.py @@ -43,7 +43,9 @@ class CoreServiceCollectorsMixin: _account_id: str | None coverage: CollectionCoverage - def _get_aio_session(self) -> Any: ... + def _get_aio_session(self) -> Any: + # Provided by AsyncAWSCollector; this stub only types the mixin + raise NotImplementedError async def _collect_ec2(self) -> list[CloudAsset]: """Collect EC2 instances.""" diff --git a/cloudg/collectors/aws_services_extended.py b/cloudg/collectors/aws_services_extended.py index 4269f78..e17e512 100644 --- a/cloudg/collectors/aws_services_extended.py +++ b/cloudg/collectors/aws_services_extended.py @@ -45,7 +45,9 @@ class ExtendedServiceCollectorsMixin: _aio_config: Any coverage: CollectionCoverage - def _get_aio_session(self) -> Any: ... + def _get_aio_session(self) -> Any: + # Provided by AsyncAWSCollector; this stub only types the mixin + raise NotImplementedError async def _collect_ecs(self) -> list[CloudAsset]: """Collect ECS clusters and services.""" @@ -221,8 +223,8 @@ async def _collect_secrets_manager(self) -> list[CloudAsset]: duration_ms=int((time.time() - start) * 1000), ) except Exception as exc: - # nosemgrep: the message names the AWS service, no credential is logged - logger.error("secretsmanager service collection failed: %s", exc) + # The service name is an argument so the format string names no secret + logger.error("%s service collection failed: %s", "secretsmanager", exc) self.coverage.record("secretsmanager", ServiceStatus.FAILED, error=str(exc)) return assets diff --git a/cloudg/collectors/azure.py b/cloudg/collectors/azure.py index f36801c..a269c47 100644 --- a/cloudg/collectors/azure.py +++ b/cloudg/collectors/azure.py @@ -17,6 +17,9 @@ from typing import Any, Callable from cloudg.collectors.base import BaseCollector +from cloudg.resilience.azure import azure_scope_of_error, build_azure_client +from cloudg.resilience.errors import describe_error, is_throttle +from cloudg.resilience.governor import get_governor from cloudg.schema.models import ( AssetType, CloudAsset, @@ -52,11 +55,13 @@ def _subscription_client_class() -> Any: def resource_management_client(credential: Any, subscription_id: str) -> Any: - return _resource_client_class()(credential, subscription_id) + return build_azure_client( + _resource_client_class(), credential, subscription_id, subscription_id=subscription_id + ) def subscription_client(credential: Any) -> Any: - return _subscription_client_class()(credential) + return build_azure_client(_subscription_client_class(), credential) def _subscription_dict(sub: Any) -> dict[str, Any]: @@ -122,6 +127,15 @@ def first_subscription_id(credential: Any) -> str | None: return subs[0]["subscription_id"] if subs else None +def _rule_values(rule: dict[str, Any], plural: str, singular: str | None = None) -> list[str]: + """Distinct non-empty values of an NSG rule's list field, falling back to + its single-value field (``source_address_prefixes`` / ``..._prefix``).""" + values = [str(v) for v in rule.get(plural) or [] if v] + if not values and singular and rule.get(singular): + values = [str(rule[singular])] + return list(dict.fromkeys(values)) + + def _resource_group_name(resource_id: str | None) -> str | None: parts = (resource_id or "").split("/") for i, part in enumerate(parts[:-1]): @@ -176,7 +190,11 @@ def _client(self, kind: str) -> Any: return client else: # pragma: no cover - programming error raise ValueError(kind) - client = cls(self._credential, self._subscription_id) + # Shared rate limits / breakers / throttle feedback and the + # configured azure-core retries (cloudg.resilience.azure) + client = build_azure_client( + cls, self._credential, self._subscription_id, subscription_id=self._subscription_id + ) self._clients[kind] = client return client @@ -204,7 +222,13 @@ async def _run_service( logger.warning("Azure %s collection skipped: SDK not installed (%s)", name, exc) except Exception as exc: logger.error("Failed to collect Azure %s: %s", name, exc) - self.service_errors[name] = str(exc) + self.service_errors[name] = describe_error(exc) + if is_throttle(exc): + # on_final_error skips calls that were never sent and prefers + # the scope carried by cloudg's own errors + get_governor().on_final_error( + azure_scope_of_error(exc, self._subscription_id, name), exc, f"Azure {name}" + ) return [] async def _gather_rows( @@ -338,24 +362,20 @@ def _rule_edges( ) -> list[NetworkEdge]: if str(rule.get("access") or "").lower() != "allow": return [] - sources = [str(s) for s in rule.get("source_address_prefixes") or [] if s] - if not sources and rule.get("source_address_prefix"): - sources = [str(rule["source_address_prefix"])] - asgs = [str(a) for a in rule.get("source_application_security_groups") or [] if a] - ports = [str(p) for p in rule.get("destination_port_ranges") or [] if p] - if not ports and rule.get("destination_port_range"): - ports = [str(rule["destination_port_range"])] + sources = _rule_values(rule, "source_address_prefixes", "source_address_prefix") + asgs = _rule_values(rule, "source_application_security_groups") + ports = _rule_values(rule, "destination_port_ranges", "destination_port_range") + suffix = " (default)" if rule.get("default") else "" common: dict[str, Any] = { "target_id": nsg.id, "edge_type": EdgeType.SECURITY_GROUP_RULE, - "port_range": ",".join(dict.fromkeys(ports)) or None, + "port_range": ",".join(ports) or None, "protocol": str(rule.get("protocol") or "ALL"), "direction": "ingress", - "description": f"NSG rule {rule.get('name')}" - + (" (default)" if rule.get("default") else ""), + "description": f"NSG rule {rule.get('name')}{suffix}", } - edges = [NetworkEdge(source_id=src, cidr=src, **common) for src in dict.fromkeys(sources)] - for asg in dict.fromkeys(asgs): + edges = [NetworkEdge(source_id=src, cidr=src, **common) for src in sources] + for asg in asgs: edges.append( NetworkEdge( source_id=by_arn.get(asg.lower(), asg), @@ -367,23 +387,25 @@ def _rule_edges( edges.append(NetworkEdge(source_id="*", cidr="*", **common)) return edges - async def collect_edges(self) -> list[NetworkEdge]: - """Collect Azure network edges from NSG rules and VNet containment.""" - assets = self._cached_assets or await self.collect() - edges: list[NetworkEdge] = [] + def _nsg_edges(self, assets: list[CloudAsset]) -> list[NetworkEdge]: + """SECURITY_GROUP_RULE edges of every allow rule of every NSG.""" by_arn = {a.arn.lower(): a.id for a in assets if a.arn} - + edges: list[NetworkEdge] = [] for nsg in (a for a in assets if a.asset_type == AssetType.NSG): for rule in nsg.metadata.get("ingress_rules", []) or []: try: edges.extend(self._rule_edges(nsg, rule, by_arn)) except Exception: logger.debug("Skipping NSG rule %s on %s", rule, nsg.name, exc_info=True) + return edges - # Containment edges: VNet contains Subnet + @staticmethod + def _vnet_edges(assets: list[CloudAsset]) -> list[NetworkEdge]: + """Containment edges: VNet contains Subnet.""" vnet_assets = { a.arn.lower(): a.id for a in assets if a.asset_type == AssetType.VNET and a.arn } + edges: list[NetworkEdge] = [] for asset in assets: vnet_id = asset.metadata.get("vnet_id") if isinstance(vnet_id, str) and vnet_id.lower() in vnet_assets: @@ -395,6 +417,11 @@ async def collect_edges(self) -> list[NetworkEdge]: relationship="VPC_CONTAINS_SUBNET", ) ) + return edges + async def collect_edges(self) -> list[NetworkEdge]: + """Collect Azure network edges from NSG rules and VNet containment.""" + assets = self._cached_assets or await self.collect() + edges = self._nsg_edges(assets) + self._vnet_edges(assets) logger.info("Collected %d Azure edges", len(edges)) return edges diff --git a/cloudg/collectors/gcp.py b/cloudg/collectors/gcp.py index f00cf67..e891ee6 100644 --- a/cloudg/collectors/gcp.py +++ b/cloudg/collectors/gcp.py @@ -20,6 +20,7 @@ import logging import re import time +from contextlib import nullcontext from dataclasses import dataclass from typing import Any, Iterable @@ -40,7 +41,11 @@ state_str, to_plain, ) +from cloudg.config import GCP_DEFAULT_SKIP_ASSET_TYPES from cloudg.coverage import CollectionCoverage, ServiceStatus +from cloudg.resilience.errors import describe_error +from cloudg.resilience.gcp import RetryCounter, gcp_retry, gcp_scope, paced +from cloudg.resilience.governor import get_governor from cloudg.inventory.catalogs import asset_type_map, load_catalog from cloudg.schema.models import ( AssetType, @@ -58,20 +63,7 @@ ) # High-churn types skipped unless configured otherwise (GCPConfig.skip_asset_types) -DEFAULT_SKIP_ASSET_TYPES: tuple[str, ...] = ( - "k8s.io/Pod", - "k8s.io/Node", - "k8s.io/Event", - "events.k8s.io/Event", - "k8s.io/Endpoints", - "discovery.k8s.io/EndpointSlice", - "apps.k8s.io/ReplicaSet", - "apps.k8s.io/ControllerRevision", - "run.googleapis.com/Revision", - "cloudkms.googleapis.com/CryptoKeyVersion", - "secretmanager.googleapis.com/SecretVersion", - "serviceusage.googleapis.com/Service", -) +DEFAULT_SKIP_ASSET_TYPES: tuple[str, ...] = GCP_DEFAULT_SKIP_ASSET_TYPES _RESOURCE_MANAGER = "//cloudresourcemanager.googleapis.com/" @@ -85,24 +77,12 @@ def _content_type(name: str) -> Any: return name -def _retry() -> Any: - try: - from google.api_core import exceptions as gexc - from google.api_core import retry as retries - except ImportError: - return None - return retries.Retry( - predicate=retries.if_exception_type( - gexc.ResourceExhausted, - gexc.ServiceUnavailable, - gexc.DeadlineExceeded, - gexc.InternalServerError, - ), - initial=1.0, - maximum=60.0, - multiplier=2.0, - timeout=900.0, - ) +def _retry(scope: Any = None, counter: Any = None) -> Any: + """api_core Retry on quota / availability errors (initial 1s, x2, capped + at ratelimit.gcp.max_backoff_seconds (60), overall deadline_seconds + (900), at most max_retries consecutive retries, retry budget); retried + errors feed cloudg's adaptive rate limiter for ``scope``.""" + return gcp_retry(scope, counter=counter) @dataclass @@ -212,9 +192,9 @@ def _get_client(self) -> Any: self._client = asset_v1.AssetServiceClient(credentials=self._credentials) return self._client - def _call_kwargs(self) -> dict[str, Any]: + def _call_kwargs(self, scope: Any = None, counter: Any = None) -> dict[str, Any]: kwargs: dict[str, Any] = {"timeout": self._timeout} - retry = _retry() + retry = _retry(scope, counter) if retry is not None: kwargs["retry"] = retry return kwargs @@ -222,13 +202,32 @@ def _call_kwargs(self) -> dict[str, Any]: def _iterate( self, method: Any, request: dict[str, Any] ) -> tuple[list[dict[str, Any]], Exception | None]: - """Drain a paged CAI call in the calling (worker) thread.""" + """Drain a paged CAI call in the calling (worker) thread. + + Each page takes a token from the shared limiter of this scope x RPC + (Cloud Asset Inventory quotas are per minute per project / org), an + open circuit fails fast, and throttling that outlasts the retries + is reported as such (the caller records it in coverage). + """ + gov = get_governor() + scope = gcp_scope(self._scope, "cloudasset", method) items: list[dict[str, Any]] = [] + counter = RetryCounter() + enabled = gov.provider_enabled("gcp") try: - for item in method(request=request, **self._call_kwargs()): - items.append(to_plain(item)) + if enabled: + gov.check(scope) + gov.limiter.acquire_sync(scope) + gov.record_call(scope) + # Concurrent listings per project / org (ratelimit.gcp.max_concurrency) + with gov.limiter.bulkhead(scope).hold_sync() if enabled else nullcontext(): + result = method(request=request, **self._call_kwargs(scope, counter)) + for item in paced(result, scope, request.get("page_size"), counter=counter): + items.append(to_plain(item)) except Exception as exc: + gov.on_final_error(scope, exc, str(scope.operation)) return items, exc + gov.on_success(scope) return items, None # ------------------------------------------------------------------ @@ -276,7 +275,7 @@ async def _collect_records(self) -> list[dict[str, Any]]: "gcp_list_assets", ServiceStatus.PARTIAL, len(raw), - f"listing truncated: {error}", + f"listing truncated: {describe_error(error)}", start=start, ) return [self._record_from_asset(r) for r in raw] @@ -286,27 +285,33 @@ async def _collect_records(self) -> list[dict[str, Any]]: self._scope, error, ) + return await self._collect_from_search(error, start) + + async def _collect_from_search(self, error: Exception, start: float) -> list[dict[str, Any]]: + """Summary records from SearchAllResources after ListAssets failed with ``error``.""" search_start = time.time() found, search_error = await asyncio.to_thread(self._search_resources) + self._record( + "gcp_list_assets", ServiceStatus.FAILED, error=describe_error(error), start=start + ) if search_error is not None and not found: - self._record("gcp_list_assets", ServiceStatus.FAILED, error=str(error), start=start) self._record( "gcp_search_resources", ServiceStatus.FAILED, - error=str(search_error), + error=describe_error(search_error), start=search_start, ) raise RuntimeError( - f"GCP asset enumeration failed for {self._scope}: ListAssets: {error}; " - f"SearchAllResources: {search_error}" + f"GCP asset enumeration failed for {self._scope}: " + f"ListAssets: {describe_error(error)}; " + f"SearchAllResources: {describe_error(search_error)}" ) - self._record("gcp_list_assets", ServiceStatus.FAILED, error=str(error), start=start) self._record( "gcp_search_resources", ServiceStatus.PARTIAL, len(found), "summary data only (ListAssets unavailable)" - + (f"; listing truncated: {search_error}" if search_error else ""), + + (f"; listing truncated: {describe_error(search_error)}" if search_error else ""), start=search_start, ) return [self._record_from_search(r) for r in found] @@ -317,16 +322,8 @@ async def _collect_records(self) -> list[dict[str, Any]]: def _learn_projects(self, records: list[dict[str, Any]]) -> None: for rec in records: - data = rec.get("data") or {} if rec["asset_type"] == "cloudresourcemanager.googleapis.com/Project": - number = _num(data.get("projectNumber")) - if not number and str(data.get("name", "")).startswith("projects/"): - number = data["name"].split("/", 1)[1] - if not number: - number = rec["name"].rsplit("/", 1)[-1] - pid = data.get("projectId") - if number and pid: - self._number_to_id.setdefault(number, pid) + self._learn_project_asset(rec) for rec in records: number = self._ancestor_number(rec) if not number or number in self._number_to_id: @@ -340,6 +337,18 @@ def _learn_projects(self, records: list[dict[str, Any]]) -> None: self._number_to_id.setdefault(numbers.pop(), self._project_id) self._id_to_number = {v: k for k, v in self._number_to_id.items()} + def _learn_project_asset(self, rec: dict[str, Any]) -> None: + """Project number -> id from a cloudresourcemanager Project asset.""" + data = rec.get("data") or {} + number = _num(data.get("projectNumber")) + if not number and str(data.get("name", "")).startswith("projects/"): + number = data["name"].split("/", 1)[1] + if not number: + number = rec["name"].rsplit("/", 1)[-1] + pid = data.get("projectId") + if number and pid: + self._number_to_id.setdefault(number, pid) + @staticmethod def _ancestor_number(rec: dict[str, Any]) -> str | None: for a in rec.get("ancestors") or []: diff --git a/cloudg/collectors/multi.py b/cloudg/collectors/multi.py index daeb33b..e3e4bf7 100644 --- a/cloudg/collectors/multi.py +++ b/cloudg/collectors/multi.py @@ -16,6 +16,7 @@ from cloudg.config import CloudGConfig from cloudg.coverage import CollectionCoverage, ServiceStatus +from cloudg.resilience.errors import describe_error from cloudg.region_discovery import RegionDiscovery, is_all_regions from cloudg.schema.models import CloudAsset, NetworkEdge @@ -54,6 +55,13 @@ def __init__( # the standard collector for that provider. self._collector_overrides = collector_overrides or {} self._caller_account: str | None = None + # Shared throttling resilience (rate limits, breakers, retries) from + # config.ratelimit / aws.max_retries / aws.retry_mode + from cloudg.resilience import configure + + configure(config) + #: Throttling telemetry of the last collect_all() run + self.resilience_stats: Any = None async def collect_all( self, @@ -63,6 +71,19 @@ async def collect_all( Returns: Tuple of (all_assets, all_edges, coverage_records). """ + from cloudg.resilience import current_stats, stats_scope + + bound = current_stats() + if bound is not None: # the caller (e.g. InventoryMapper) owns the run's stats + self.resilience_stats = bound + return await self._collect_all() + with stats_scope() as stats: + self.resilience_stats = stats + return await self._collect_all() + + async def _collect_all( + self, + ) -> tuple[list[CloudAsset], list[NetworkEdge], list[CollectionCoverage]]: providers = self._config.providers all_assets: list[CloudAsset] = [] all_edges: list[NetworkEdge] = [] @@ -95,7 +116,7 @@ async def collect_all( if isinstance(result, Exception): logger.error("Provider %s collection failed: %s", label, result) coverage = CollectionCoverage(provider=label) - coverage.record(f"{label}_full", ServiceStatus.FAILED, error=str(result)) + coverage.record(f"{label}_full", ServiceStatus.FAILED, error=describe_error(result)) self._coverage.append(coverage) elif isinstance(result, list): for assets, edges in result: @@ -182,16 +203,17 @@ async def _collect_aws_single( async with self._semaphore: try: - session = self._aws_session(account_id, region, cfg, coverage) + # Blocking boto3 calls (AssumeRole, GetCallerIdentity) run in a + # worker thread: their rate-limit waits must not stall the loop + session = await asyncio.to_thread( + self._aws_session, account_id, region, cfg, coverage + ) if session is None: return [], [] # Resolve account ID if not provided if not account_id: - try: - account_id = session.client("sts").get_caller_identity().get("Account") - except Exception: - account_id = "unknown" + account_id = await asyncio.to_thread(_caller_account, session) collector = self._build_aws_collector(session, region, account_id, is_primary) assets, edges, duration_ms = await self._timed_collect(collector) @@ -200,7 +222,7 @@ async def _collect_aws_single( except Exception as exc: logger.error("AWS collection failed for %s/%s: %s", account_id, region, exc) - coverage.record("aws_full", ServiceStatus.FAILED, error=str(exc)) + coverage.record("aws_full", ServiceStatus.FAILED, error=describe_error(exc)) return [], [] @staticmethod @@ -220,7 +242,7 @@ def _aws_session( # Supports direct keys, OIDC web identity, profiles, the # default chain (instance/task roles), and AssumeRole with - # optional ExternalId — see cloudg.credentials. + # optional ExternalId; see cloudg.credentials. assume_into = account_id if account_id and account_id == self._caller_account: assume_into = None @@ -228,7 +250,7 @@ def _aws_session( return build_aws_session(cfg, region, account_id=assume_into) except RuntimeError as exc: logger.error("AWS auth failed for %s/%s: %s", account_id, region, exc) - coverage.record("sts_assume_role", ServiceStatus.FAILED, error=str(exc)) + coverage.record("sts_assume_role", ServiceStatus.FAILED, error=describe_error(exc)) return None def _build_aws_collector( @@ -343,25 +365,11 @@ async def _collect_azure_single( assets = await collector.collect() except Exception as exc: logger.error("Azure collection failed for %s: %s", subscription_id, exc) - coverage.record("azure_full", ServiceStatus.FAILED, error=str(exc)) + coverage.record("azure_full", ServiceStatus.FAILED, error=describe_error(exc)) return [], [] - edges: list[NetworkEdge] = [] - try: - edges = await collector.collect_edges() - except Exception as exc: - logger.error("Azure edge collection failed for %s: %s", subscription_id, exc) - coverage.record("azure_edges", ServiceStatus.FAILED, error=str(exc)) - - duration_ms = int((time.time() - start) * 1000) - service_errors = dict(getattr(collector, "service_errors", None) or {}) - for service, error in service_errors.items(): - coverage.record(f"azure_{service}", ServiceStatus.FAILED, error=error) - coverage.record( - "azure_full", - ServiceStatus.PARTIAL if service_errors else ServiceStatus.SUCCESS, - asset_count=len(assets), - duration_ms=duration_ms, + edges = await _azure_edges_and_coverage( + collector, len(assets), coverage, subscription_id, start ) return assets, edges @@ -390,20 +398,7 @@ async def _collect_gcp_multi(self) -> list[tuple[list[CloudAsset], list[NetworkE if not project_ids: project_ids = [None] - # Region resolution - gcp_regions = self._config.gcp.regions - if is_all_regions(gcp_regions): - logger.info("GCP: discovering all regions...") - try: - from cloudg.credentials import build_gcp_credentials - - creds, default_project = build_gcp_credentials(self._config.gcp) - pid = project_ids[0] or default_project - except Exception: - creds, pid = None, None - gcp_regions = await self._region_discovery.discover_gcp(creds, pid) - - self._resolved_regions["gcp"] = gcp_regions + self._resolved_regions["gcp"] = await self._resolve_gcp_regions(project_ids[0]) # Cloud Asset Inventory returns resources across ALL regions, so no # per-region iteration is needed. @@ -421,6 +416,21 @@ async def _collect_gcp_multi(self) -> list[tuple[list[CloudAsset], list[NetworkE tasks = [self._collect_gcp_single(pid) for pid in project_ids] return await asyncio.gather(*tasks, return_exceptions=False) + async def _resolve_gcp_regions(self, first_project_id: str | None) -> Any: + """Return the configured GCP regions, discovering them when set to all.""" + gcp_regions = self._config.gcp.regions + if is_all_regions(gcp_regions): + logger.info("GCP: discovering all regions...") + try: + from cloudg.credentials import build_gcp_credentials + + creds, default_project = build_gcp_credentials(self._config.gcp) + pid = first_project_id or default_project + except Exception: + creds, pid = None, None + gcp_regions = await self._region_discovery.discover_gcp(creds, pid) + return gcp_regions + async def _collect_gcp_single( self, project_id: str | None, @@ -463,7 +473,7 @@ async def _collect_gcp_single( except Exception as exc: logger.error("GCP collection failed for %s: %s", label, exc) - coverage.record("gcp_full", ServiceStatus.FAILED, error=str(exc)) + coverage.record("gcp_full", ServiceStatus.FAILED, error=describe_error(exc)) return [], [] @staticmethod @@ -505,3 +515,44 @@ def _record_gcp_coverage( duration_ms=duration_ms, error="; ".join(f"{s.service}: {s.error}" for s in degraded) or None, ) + + +def _caller_account(session: Any) -> str | None: + """The account of ``session``'s credentials ("unknown" when STS fails).""" + try: + return session.client("sts").get_caller_identity().get("Account") + except Exception as exc: + logger.debug("GetCallerIdentity failed: %s", exc) + return "unknown" + + +async def _azure_edges_and_coverage( + collector: Any, + asset_count: int, + coverage: CollectionCoverage, + subscription_id: str | None, + start: float, +) -> list[NetworkEdge]: + """Collect one Azure subscription's edges and record its coverage. + + An edge failure is recorded as ``azure_edges`` and never discards the + assets; per-service collector errors become ``azure_`` entries. + """ + edges: list[NetworkEdge] = [] + try: + edges = await collector.collect_edges() + except Exception as exc: + logger.error("Azure edge collection failed for %s: %s", subscription_id, exc) + coverage.record("azure_edges", ServiceStatus.FAILED, error=describe_error(exc)) + + duration_ms = int((time.time() - start) * 1000) + service_errors = dict(getattr(collector, "service_errors", None) or {}) + for service, error in service_errors.items(): + coverage.record(f"azure_{service}", ServiceStatus.FAILED, error=error) + coverage.record( + "azure_full", + ServiceStatus.PARTIAL if service_errors else ServiceStatus.SUCCESS, + asset_count=asset_count, + duration_ms=duration_ms, + ) + return edges diff --git a/cloudg/config.py b/cloudg/config.py index 4b34995..610939f 100644 --- a/cloudg/config.py +++ b/cloudg/config.py @@ -6,7 +6,7 @@ from pathlib import Path from typing import Any -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, model_validator logger = logging.getLogger(__name__) @@ -140,6 +140,24 @@ class AzureConfig(BaseModel): ) +# High-churn Cloud Asset Inventory types the GCP collector skips by default +# (GCPConfig.skip_asset_types; cloudg.collectors.gcp.DEFAULT_SKIP_ASSET_TYPES) +GCP_DEFAULT_SKIP_ASSET_TYPES: tuple[str, ...] = ( + "k8s.io/Pod", + "k8s.io/Node", + "k8s.io/Event", + "events.k8s.io/Event", + "k8s.io/Endpoints", + "discovery.k8s.io/EndpointSlice", + "apps.k8s.io/ReplicaSet", + "apps.k8s.io/ControllerRevision", + "run.googleapis.com/Revision", + "cloudkms.googleapis.com/CryptoKeyVersion", + "secretmanager.googleapis.com/SecretVersion", + "serviceusage.googleapis.com/Service", +) + + class GCPConfig(BaseModel): """GCP-specific configuration. @@ -170,20 +188,7 @@ class GCPConfig(BaseModel): "organization / project force either mode.", ) skip_asset_types: list[str] = Field( - default_factory=lambda: [ - "k8s.io/Pod", - "k8s.io/Node", - "k8s.io/Event", - "events.k8s.io/Event", - "k8s.io/Endpoints", - "discovery.k8s.io/EndpointSlice", - "apps.k8s.io/ReplicaSet", - "apps.k8s.io/ControllerRevision", - "run.googleapis.com/Revision", - "cloudkms.googleapis.com/CryptoKeyVersion", - "secretmanager.googleapis.com/SecretVersion", - "serviceusage.googleapis.com/Service", - ], + default_factory=lambda: list(GCP_DEFAULT_SKIP_ASSET_TYPES), description="Cloud Asset Inventory types left out of the map (high-churn objects)", ) asset_page_size: int = Field(default=1000, ge=1, le=1000) @@ -331,6 +336,197 @@ class RulesetConfig(BaseModel): load_external: bool = Field(default=True) +def _warn_unknown_keys(data: Any, model: type[BaseModel], path: str) -> None: + """Log (not raise: CloudGConfig is public API) keys ``model`` does not define.""" + if not isinstance(data, dict): + return + for key in data: + if key not in model.model_fields: + logger.warning("Unknown config key %s.%s is ignored", path, key) + + +class ServiceRateLimitConfig(BaseModel): + """Rate limit of one service (or ``service.Operation``) leaf bucket.""" + + max_rps: float = Field(gt=0, description="Sustained requests per second") + burst: float | None = Field( + default=None, gt=0, description="Bucket capacity (default: 2 x max_rps)" + ) + per_operation: bool | None = Field( + default=None, + description="One bucket (and circuit breaker) per API operation instead of per " + "service (default: the built-in choice, true for AWS ec2)", + ) + + +class ProviderRateLimitConfig(BaseModel): + """Throttling behaviour for one provider. + + Unset (None) limits fall back to cloudg's built-in per-provider and + per-service defaults (documented with sources in + ``cloudg/resilience/limiter.py``), which sit at or below each + provider's published limits. Each description says which call paths + the key affects: "AWS hooks" are the botocore event hooks on every + boto3 / aioboto3 session cloudg builds, "Azure policy" is the + per-attempt pipeline policy on every Azure SDK client cloudg builds, + "GCP listings" are Cloud Asset Inventory calls, and + "call_with_resilience" is the generic wrapper for embedders' own calls. + """ + + enabled: bool = Field( + default=True, description="Rate limit / circuit-break this provider (all paths)" + ) + max_rps: float | None = Field( + default=None, + gt=0, + description="Default requests/s of a leaf bucket (account x region x service, or " + "operation) on all paths (AWS 20, Azure 15, GCP 10)", + ) + burst: float | None = Field(default=None, gt=0, description="Default leaf bucket capacity") + account_max_rps: float | None = Field( + default=None, + gt=0, + description="Aggregate requests/s per account / subscription / project, all paths " + "(Azure 20; off for AWS and GCP)", + ) + account_burst: float | None = Field( + default=None, gt=0, description="Capacity of the account bucket (Azure 200)" + ) + global_max_rps: float | None = Field( + default=None, + gt=0, + description="Aggregate requests/s of the whole provider, all paths (off by default)", + ) + global_burst: float | None = Field( + default=None, gt=0, description="Capacity of the provider-wide bucket" + ) + max_concurrency: int | None = Field( + default=None, + ge=1, + le=1024, + description="In-flight requests per account (bulkhead): AWS hooks and Azure policy per " + "HTTP request, GCP per concurrent listing, call_with_resilience per call " + "(AWS 128, Azure 16, GCP 16)", + ) + adaptive: bool = Field( + default=True, + description="Halve a scope's rate when throttled, recover additively (AIMD), all paths", + ) + min_rps: float | None = Field( + default=None, gt=0, description="Floor the adaptive rate never drops below" + ) + breaker_threshold: int = Field( + default=5, + ge=1, + le=1000, + description="Consecutive calls failing on throttling (after retries) that open the " + "circuit of a scope (per operation where the bucket is per operation), all paths", + ) + breaker_cooldown_seconds: float = Field( + default=60.0, + ge=0, + description="Seconds an open circuit rejects calls before a probe, all paths", + ) + max_retries: int = Field( + default=8, + ge=0, + le=50, + description="Retries of a throttled / transient call: Azure SDK retry_total and " + "retry_status, GCP consecutive retries per page, call_with_resilience. Not used " + "by the AWS hooks: botocore retries come from aws.max_retries", + ) + max_backoff_seconds: float = Field( + default=60.0, + gt=0, + description="Cap of one backoff sleep: Azure retry_backoff_max, GCP api_core Retry " + "maximum, call_with_resilience. AWS: botocore's own backoff (not configurable here)", + ) + deadline_seconds: float = Field( + default=900.0, + gt=0, + description="Overall time budget of one call including retries: Azure RetryPolicy " + "timeout, GCP api_core Retry timeout, call_with_resilience. AWS: botocore " + "timeouts and aws.max_retries bound a call instead", + ) + retry_budget: int = Field( + default=500, + ge=0, + description="Retries the provider may make before cloudg stops retrying (each " + "success refunds 0.1): AWS hooks, Azure policy, GCP listings, call_with_resilience", + ) + services: dict[str, ServiceRateLimitConfig] = Field( + default_factory=dict, + description="Per-service overrides keyed by SDK service name (AWS: ec2, iam, ...; " + "Azure: network, compute, resourcegraph, ...; GCP: cloudasset) or " + "'.' (e.g. cloudasset.ListAssets)", + ) + live_cooldown_seconds: float | None = Field( + default=None, + ge=0, + description="Overrides ratelimit.live_cooldown_seconds for this provider " + "(LiveOperationGuard only)", + ) + + +def _provider_ratelimit() -> ProviderRateLimitConfig: + return ProviderRateLimitConfig() + + +class RateLimitConfig(BaseModel): + """Cloud API throttling resilience (rate limits, retries, circuit breakers) + and guard rails for live collections triggered on demand (MCP). + + Unknown keys are logged as warnings and ignored. + """ + + enabled: bool = Field(default=True, description="Master switch for every provider") + aws: ProviderRateLimitConfig = Field(default_factory=_provider_ratelimit) + azure: ProviderRateLimitConfig = Field(default_factory=_provider_ratelimit) + gcp: ProviderRateLimitConfig = Field(default_factory=_provider_ratelimit) + live_cooldown_seconds: float = Field( + default=120.0, + ge=0, + description="Minimum seconds between two live collections of the same scope " + "(LiveOperationGuard)", + ) + live_max_concurrent: int = Field( + default=1, ge=1, le=64, description="Concurrent live collections per scope" + ) + live_max_concurrent_total: int = Field( + default=2, ge=1, le=256, description="Concurrent live collections overall" + ) + live_caller_max_operations: int = Field( + default=0, ge=0, description="Live collections per caller per window (0 = unlimited)" + ) + live_caller_window_seconds: float = Field( + default=3600.0, gt=0, description="Window of live_caller_max_operations" + ) + live_operation_timeout_seconds: float | None = Field( + default=3600.0, + gt=0, + description="A live collection running longer is cancelled and its scopes freed " + "(None: no limit)", + ) + + @model_validator(mode="before") + @classmethod + def _unknown_keys(cls, data: Any) -> Any: + _warn_unknown_keys(data, cls, "ratelimit") + if isinstance(data, dict): + for provider in ("aws", "azure", "gcp"): + pdata = data.get(provider) + _warn_unknown_keys(pdata, ProviderRateLimitConfig, f"ratelimit.{provider}") + services = pdata.get("services") if isinstance(pdata, dict) else None + if isinstance(services, dict): + for name, sdata in services.items(): + _warn_unknown_keys( + sdata, + ServiceRateLimitConfig, + f"ratelimit.{provider}.services.{name}", + ) + return data + + class CloudGConfig(BaseModel): """Root configuration model.""" @@ -339,7 +535,7 @@ class CloudGConfig(BaseModel): default=["aws"], description="Providers to scan: aws, azure, gcp. Use multiple for simultaneous scanning.", ) - # Backward-compat alias — single provider string is auto-wrapped + # Backward-compat alias: a single provider string is auto-wrapped provider: str | None = Field( default=None, exclude=True, @@ -356,6 +552,10 @@ class CloudGConfig(BaseModel): terraform: TerraformConfig = Field(default_factory=TerraformConfig) report: ReportConfig = Field(default_factory=ReportConfig) rulesets: RulesetConfig = Field(default_factory=RulesetConfig) + ratelimit: RateLimitConfig = Field( + default_factory=RateLimitConfig, + description="Throttling resilience: rate limits, retries, circuit breakers, live-run guard", + ) log_file: str | None = None verbose: bool = False concurrency_limit: int = Field(default=5, ge=1, le=50, description="Max concurrent API calls") diff --git a/cloudg/coverage.py b/cloudg/coverage.py index 47c5fab..4349117 100644 --- a/cloudg/coverage.py +++ b/cloudg/coverage.py @@ -1,4 +1,4 @@ -"""Collection coverage tracking — records which services succeeded/failed.""" +"""Collection coverage tracking: records which services succeeded/failed.""" from __future__ import annotations diff --git a/cloudg/credentials.py b/cloudg/credentials.py index 92d71e3..6653576 100644 --- a/cloudg/credentials.py +++ b/cloudg/credentials.py @@ -16,10 +16,10 @@ Azure (build_azure_credential): 1. Workload identity federation (tenant_id + client_id + - federated_token_file — AKS / GitHub OIDC) + federated_token_file; AKS / GitHub OIDC) 2. Service principal with client secret 3. Service principal with client certificate - 4. Managed identity (system- or user-assigned — inherited identity) + 4. Managed identity (system- or user-assigned, inherited identity) 5. DefaultAzureCredential chain (env, CLI, PowerShell, managed identity) GCP (build_gcp_credentials): @@ -208,20 +208,51 @@ def build_aws_session( session, assumed_via_oidc = _aws_base_session( boto3, cfg, region, session_name, role_arn, token_file ) + _instrument_aws_session(session, None, region) # 5. Optional role assumption on top of the base credentials - target_role: str | None = None - if role_arn and not assumed_via_oidc: - target_role = role_arn - elif account_id and getattr(cfg, "role_name", None): - target_role = f"arn:aws:iam::{account_id}:role/{cfg.role_name}" - + target_role = _aws_target_role(cfg, role_arn, assumed_via_oidc, account_id) if target_role: session = _aws_assume_role(boto3, session, target_role, session_name, external_id, region) + _instrument_aws_session(session, account_id, region) return session +def _aws_target_role( + cfg: Any, role_arn: str | None, assumed_via_oidc: bool, account_id: str | None +) -> str | None: + """The role to assume on top of the base session, if any. + + An explicit ``role_arn`` wins unless OIDC federation already assumed it; + otherwise ``account_id`` plus ``cfg.role_name`` names a cross-account role. + """ + if role_arn and not assumed_via_oidc: + return role_arn + if account_id and getattr(cfg, "role_name", None): + return f"arn:aws:iam::{account_id}:role/{cfg.role_name}" + return None + + +def _instrument_aws_session(session: Any, account_id: str | None, region: str) -> None: + """Route the session's (STS, Organizations, Control Tower, ...) calls + through cloudg's shared rate limiter / circuit breakers, with the + configured botocore retries as the default client config.""" + try: + from botocore.config import Config + + from cloudg.resilience.aws import botocore_retries, install_aws_hooks + + install_aws_hooks( + session, + account_id=account_id, + region=region, + default_config=Config(retries=botocore_retries()), + ) + except Exception as exc: # instrumentation must never break authentication + logger.debug("AWS session instrumentation skipped: %s", exc) + + # ───────────────────────────────────────────────────────────────────── # Azure # ───────────────────────────────────────────────────────────────────── @@ -243,16 +274,47 @@ def build_azure_credential(cfg: Any) -> Any: "azure-identity is required for Azure. Install with: pip install cloudg[azure]" ) from exc - tenant_id = getattr(cfg, "tenant_id", None) or os.environ.get("AZURE_TENANT_ID") - client_id = getattr(cfg, "client_id", None) or os.environ.get("AZURE_CLIENT_ID") - client_secret = getattr(cfg, "client_secret", None) or os.environ.get("AZURE_CLIENT_SECRET") - certificate_path = getattr(cfg, "certificate_path", None) - federated_token_file = getattr(cfg, "federated_token_file", None) or os.environ.get( - "AZURE_FEDERATED_TOKEN_FILE" - ) + explicit = _azure_service_principal(identity, cfg) + if explicit is not None: + return explicit + + # 4. Managed identity (inherited from the VM / App Service / AKS node) + if getattr(cfg, "use_managed_identity", False): + mi_client_id = getattr(cfg, "managed_identity_client_id", None) + logger.info( + "Azure auth: managed identity (%s)", + mi_client_id or "system-assigned", + ) + if mi_client_id: + return identity.ManagedIdentityCredential(client_id=mi_client_id) + return identity.ManagedIdentityCredential() + + # 5. Default chain (env, managed identity, CLI, PowerShell) + logger.info("Azure auth: DefaultAzureCredential chain") + return identity.DefaultAzureCredential() + + +def _azure_setting(cfg: Any, name: str, env: str | None = None) -> Any: + """``cfg.``, else the environment variable ``env``.""" + value = getattr(cfg, name, None) + if not value and env: + value = os.environ.get(env) + return value + + +def _azure_service_principal(identity: Any, cfg: Any) -> Any: + """Workload identity, secret or certificate credential of the configured + app registration (None when the settings for none of them are complete).""" + tenant_id = _azure_setting(cfg, "tenant_id", "AZURE_TENANT_ID") + client_id = _azure_setting(cfg, "client_id", "AZURE_CLIENT_ID") + if not (tenant_id and client_id): + return None + client_secret = _azure_setting(cfg, "client_secret", "AZURE_CLIENT_SECRET") + certificate_path = _azure_setting(cfg, "certificate_path") + federated_token_file = _azure_setting(cfg, "federated_token_file", "AZURE_FEDERATED_TOKEN_FILE") # 1. Workload identity federation (AKS, GitHub Actions OIDC) - if tenant_id and client_id and federated_token_file: + if federated_token_file: logger.info("Azure auth: workload identity federation (client %s)", client_id) return identity.WorkloadIdentityCredential( tenant_id=tenant_id, @@ -261,7 +323,7 @@ def build_azure_credential(cfg: Any) -> Any: ) # 2. Service principal with secret - if tenant_id and client_id and client_secret: + if client_secret: logger.info("Azure auth: service principal (client %s)", client_id) return identity.ClientSecretCredential( tenant_id=tenant_id, @@ -270,28 +332,14 @@ def build_azure_credential(cfg: Any) -> Any: ) # 3. Service principal with certificate - if tenant_id and client_id and certificate_path: + if certificate_path: logger.info("Azure auth: service principal certificate (client %s)", client_id) return identity.CertificateCredential( tenant_id=tenant_id, client_id=client_id, certificate_path=certificate_path, ) - - # 4. Managed identity (inherited from the VM / App Service / AKS node) - if getattr(cfg, "use_managed_identity", False): - mi_client_id = getattr(cfg, "managed_identity_client_id", None) - logger.info( - "Azure auth: managed identity (%s)", - mi_client_id or "system-assigned", - ) - if mi_client_id: - return identity.ManagedIdentityCredential(client_id=mi_client_id) - return identity.ManagedIdentityCredential() - - # 5. Default chain (env, managed identity, CLI, PowerShell) - logger.info("Azure auth: DefaultAzureCredential chain") - return identity.DefaultAzureCredential() + return None # ───────────────────────────────────────────────────────────────────── @@ -303,7 +351,7 @@ def build_gcp_credentials(cfg: Any) -> tuple[Any, str | None]: """Build google-auth credentials from a GCPConfig. Returns: - (credentials, default_project_id) — project may be None when the + (credentials, default_project_id); project may be None when the auth method carries no project (e.g. impersonation). Raises: @@ -322,7 +370,7 @@ def build_gcp_credentials(cfg: Any) -> tuple[Any, str | None]: if credentials_file and getattr(cfg, "credentials_file", None): # 1. Explicit file: service account key OR workload identity - # federation (external_account) config — load handles both. + # federation (external_account) config; load handles both. logger.info("GCP auth: file-based identity (%s)", credentials_file) credentials, project = google.auth.load_credentials_from_file( credentials_file, scopes=_GCP_SCOPES diff --git a/cloudg/graph/builder.py b/cloudg/graph/builder.py index 690f9fe..668b9e9 100644 --- a/cloudg/graph/builder.py +++ b/cloudg/graph/builder.py @@ -9,16 +9,115 @@ import networkx as nx -from cloudg.schema.models import CloudAsset, NetworkEdge +from cloudg.graph.ports import ( + ALL_PORTS_RANGE, + ALL_PROTOCOLS, + FILTER_RULE_EDGES, + PORTED_PROTOCOLS, + PORTLESS_PROTOCOLS, +) +from cloudg.schema.models import CloudAsset, EdgeType, NetworkEdge logger = logging.getLogger(__name__) +# Edge types that carry ports. When two of them share a source and target +# they are merged into one graph edge (see GraphBuilder._add_edge). +_MERGEABLE_EDGES = FILTER_RULE_EDGES | {EdgeType.INTERNET_EXPOSED.value} + +# Protocols whose empty port list on a filter rule allows every port: the +# port-carrying protocols plus the spellings of "every protocol". The +# port-less protocols (ICMP and the like, see cloudg.graph.ports) add no +# ports when such a rule is merged with a rule for another protocol: AWS +# stores the ICMP type and code in the port fields and Azure writes "*". +_EMPTY_MEANS_ALL_PROTOCOLS = PORTED_PROTOCOLS | ALL_PROTOCOLS + + +def _split(value: str) -> list[str]: + return [part.strip() for part in value.split(",") if part.strip()] + + +def _join_unique(values: list[str], *, ignore_case: bool = False) -> list[str]: + """``values`` without repeats, first spelling kept, in order.""" + seen: dict[str, str] = {} + for value in values: + seen.setdefault(value.upper() if ignore_case else value, value) + return list(seen.values()) + + +def _can_merge(existing: dict[str, Any], new: dict[str, Any]) -> bool: + """Whether two edges on the same source and target can be merged. + + They must be the same port-carrying edge type with the same CIDR and + direction. Protocols may differ (see :func:`_merge_edge_attrs`). + """ + return ( + existing.get("edge_type") == new["edge_type"] + and new["edge_type"] in _MERGEABLE_EDGES + and existing.get("cidr", "") == new["cidr"] + and str(existing.get("direction", "")).lower() == new["direction"].lower() + ) + + +def _rule_ports(attrs: dict[str, Any], mixed_protocols: bool) -> list[str]: + """The ``port_range`` tokens one edge contributes to a merged edge.""" + protocols = {p.upper() for p in _split(str(attrs.get("protocol", "")))} + if mixed_protocols and protocols and protocols <= PORTLESS_PROTOCOLS: + return [] + tokens = _split(str(attrs.get("port_range", ""))) + if tokens: + return tokens + # An empty port list on a TCP / UDP / all-protocol filter rule allows + # every port (the GCP collector writes it that way). Spell it out so the + # other rule's ports do not narrow it. + if attrs.get("edge_type") in FILTER_RULE_EDGES and protocols & _EMPTY_MEANS_ALL_PROTOCOLS: + return [ALL_PORTS_RANGE] + return [] + + +def _merge_edge_attrs(existing: dict[str, Any], new: dict[str, Any]) -> None: + """Fold the attributes of ``new`` into the graph edge ``existing``. + + ``port_range`` becomes the comma-separated union of both edges' ports. + When the protocols match it is the two strings joined as written. When + they differ, ``protocol`` lists both (``"TCP,ICMP"``), ICMP and other + port-less rules add no ports, and an empty TCP / UDP / all-protocol rule + adds ``0-65535``. The ports are then a union across protocols: the edge + still says which ports are open from the source, but not which protocol + each port is open for. Descriptions are joined with ``"; "``. + """ + protocols = _join_unique( + _split(str(existing.get("protocol", ""))) + _split(new["protocol"]), ignore_case=True + ) + mixed = len(protocols) > 1 + if not mixed and not existing.get("port_range") and not new["port_range"]: + port_range = "" + else: + port_range = ",".join(_join_unique(_rule_ports(existing, mixed) + _rule_ports(new, mixed))) + descriptions = _join_unique( + [d for d in str(existing.get("description", "")).split("; ") if d] + + ([new["description"]] if new["description"] else []) + ) + existing["protocol"] = ",".join(protocols) + existing["port_range"] = port_range + existing["description"] = "; ".join(descriptions) + if not existing.get("relationship"): + existing["relationship"] = new["relationship"] + class GraphBuilder: """Builds a NetworkX directed graph from collected cloud assets and edges. Nodes represent cloud resources with full CloudAsset data as attributes. Edges represent network connectivity, IAM trust, or containment relationships. + + The graph is a ``nx.DiGraph``, so it holds one edge per source and target. + Collectors often emit several rules for one pair, for example two + security group ingress rules from ``0.0.0.0/0`` for ports 22 and 443. + Such ``SECURITY_GROUP_RULE``, ``NACL_RULE`` and ``INTERNET_EXPOSED`` edges + are merged into one graph edge whose ``port_range`` lists the ports of + all of them (``"22,443"``), as long as their CIDR and direction match. + Any other edge sharing a source and target with an earlier one replaces + its attributes, as NetworkX does. """ def __init__(self) -> None: @@ -93,22 +192,38 @@ def _ensure_node(self, node_id: str) -> None: ) def _add_edge(self, edge: NetworkEdge) -> None: - """Add an edge, creating placeholder endpoints as needed.""" + """Add an edge, creating placeholder endpoints as needed. + + A rule edge whose source and target already have a compatible rule + edge is merged into it instead of overwriting it; see the class + docstring. + """ # Ensure source/target nodes exist (create placeholder if needed) self._ensure_node(edge.source_id) self._ensure_node(edge.target_id) - self._graph.add_edge( - edge.source_id, - edge.target_id, - edge_type=edge.edge_type.value, - port_range=edge.port_range or "", - protocol=edge.protocol or "", - cidr=edge.cidr or "", - direction=edge.direction or "", - description=edge.description or "", - relationship=edge.relationship or "", - ) + attrs = { + "edge_type": edge.edge_type.value, + "port_range": edge.port_range or "", + "protocol": edge.protocol or "", + "cidr": edge.cidr or "", + "direction": edge.direction or "", + "description": edge.description or "", + "relationship": edge.relationship or "", + } + existing = self._graph.get_edge_data(edge.source_id, edge.target_id) + if existing is not None: + if _can_merge(existing, attrs): + _merge_edge_attrs(existing, attrs) + return + logger.debug( + "Edge %s -> %s (%s) replaces a %s edge on the same pair", + edge.source_id, + edge.target_id, + attrs["edge_type"], + existing.get("edge_type"), + ) + self._graph.add_edge(edge.source_id, edge.target_id, **attrs) def compute_centrality(self) -> dict[str, dict[str, float]]: """Compute graph centrality metrics for blast-radius scoring. diff --git a/cloudg/graph/ontology.py b/cloudg/graph/ontology.py index 1a1e814..4151724 100644 --- a/cloudg/graph/ontology.py +++ b/cloudg/graph/ontology.py @@ -1,7 +1,7 @@ """Semantic ontology engine for cloud infrastructure. Builds an RDF/OWL knowledge graph from CloudG assets, edges, and findings -using rdflib. Provides ~60 typed semantic relations across 7 domain groups, +using rdflib. Provides 64 typed semantic relations across 7 domain groups, SPARQL query interface, and multi-format export (Turtle, JSON-LD, RDF/XML). The ontology is designed to be RAG-ready: every triple carries rich metadata @@ -26,6 +26,7 @@ get_relations_for_group, infer_asset_relations, infer_relations, + iter_asset_relations, ) from cloudg.inventory.dependencies import AssetIndex from cloudg.schema.models import ( @@ -46,6 +47,7 @@ "get_relations_for_group", "infer_asset_relations", "infer_relations", + "iter_asset_relations", ] logger = logging.getLogger(__name__) @@ -351,16 +353,15 @@ def build( for edge in edges: self._add_edge(edge, assets_by_id) - # Add inferred asset-level relations - for asset in assets: - for rel_type, target_id in infer_asset_relations(asset, assets): - src_uri = CMR[asset.id] - tgt_uri = ( - CMR[target_id] - if not target_id.startswith("tag:") - else CMR[target_id.replace(":", "_")] - ) - self._graph.add((src_uri, CMP[rel_type.value], tgt_uri)) + # Add inferred asset-level relations (KMS key index built once) + for asset, rel_type, target_id in iter_asset_relations(assets): + src_uri = CMR[asset.id] + tgt_uri = ( + CMR[target_id] + if not target_id.startswith("tag:") + else CMR[target_id.replace(":", "_")] + ) + self._graph.add((src_uri, CMP[rel_type.value], tgt_uri)) # Add findings, attached to the asset they resolve to if findings: @@ -436,8 +437,9 @@ def _add_finding(self, finding: Finding, index: AssetIndex | None = None) -> Non """Add a Finding as an OWL individual with relations. The affected resource is the asset ``index`` resolves the finding - to (by ID, ARN, unique name or unique ARN tail), falling back to - the raw ``resource_id`` when nothing matches. + to (by ID, ARN, unique name, unique ARN tail or unique ARN + resource-part tail, see :meth:`AssetIndex.resolve_finding`), + falling back to the raw ``resource_id`` when nothing matches. """ uri = CMR[f"finding_{finding.id}"] self._graph.add((uri, RDF.type, CM["SecurityFinding"])) diff --git a/cloudg/graph/ontology_rules.py b/cloudg/graph/ontology_rules.py index 34444f5..d24da02 100644 --- a/cloudg/graph/ontology_rules.py +++ b/cloudg/graph/ontology_rules.py @@ -1,15 +1,24 @@ """Relation taxonomy and inference rules for the cloud ontology. -Defines the ~63 typed semantic relations across 7 domain groups and the +Defines the 64 typed semantic relations across 7 domain groups and the inference layer that derives them from raw CloudG edges and asset metadata. The OWL graph builder itself lives in :mod:`cloudg.graph.ontology`. """ from __future__ import annotations +from collections.abc import Iterable, Iterator from enum import Enum from typing import Callable +from cloudg.graph.ports import ( + PortRange, + covers_all_ports, + edge_port_ranges, + is_internet_source, + port_in_ranges, +) +from cloudg.inventory.dependencies import AssetIndex from cloudg.schema.models import ( AssetType, CloudAsset, @@ -18,7 +27,7 @@ ) # --------------------------------------------------------------------------- -# Relation taxonomy: ~63 semantic relation types across 7 groups +# Relation taxonomy: 64 semantic relation types across 7 groups # --------------------------------------------------------------------------- @@ -37,7 +46,7 @@ class RelationGroup(str, Enum): class RelationType(str, Enum): """Fine-grained semantic relation types for cloud asset mapping.""" - # ── Network (15) ── + # ── Network (16) ── INGRESS_ALLOWED = "INGRESS_ALLOWED" INGRESS_DENIED = "INGRESS_DENIED" EGRESS_ALLOWED = "EGRESS_ALLOWED" @@ -226,27 +235,43 @@ def get_relations_for_group(group: RelationGroup) -> list[RelationType]: } -def _infer_port_relations(edge: NetworkEdge) -> list[RelationType]: - """Infer port-specific relations for a security-group rule.""" - relations: list[RelationType] = [] - ports = edge.ports or [] - protocol = (edge.protocol or "").upper() - - if len(ports) == 1: - port_rel = _PORT_RELATIONS.get(ports[0]) - relations.append(port_rel if port_rel else RelationType.PORT_RESTRICTED) - elif len(ports) > 1 and len(ports) <= 20: - # Specific port set (not wide open) - for p in ports: - port_rel = _PORT_RELATIONS.get(p) - if port_rel and port_rel not in relations: - relations.append(port_rel) - if not any(r in relations for r in _PORT_RELATIONS.values()): - relations.append(RelationType.PORT_RESTRICTED) - elif protocol == "ALL" or (edge.port_range and edge.port_range == "0-65535"): - relations.append(RelationType.ALL_TRAFFIC) +def _rule_port_ranges(edge: NetworkEdge) -> tuple[PortRange, ...]: + """The ports a rule edge opens, parsed as the reachability analysis does. - return relations + ``port_range`` and ``protocol`` go through + :func:`cloudg.graph.ports.edge_port_ranges`. An edge with no + ``port_range`` but a ``ports`` list uses that list instead. + """ + if not (edge.port_range or "").strip() and edge.ports: + return tuple((p, p) for p in edge.ports) + return edge_port_ranges( + { + "edge_type": edge.edge_type.value, + "port_range": edge.port_range, + "protocol": edge.protocol, + } + ) + + +def _infer_port_relations(edge: NetworkEdge) -> list[RelationType]: + """Infer port-specific relations for a security-group rule. + + The rule's ports are parsed numerically (see :func:`_rule_port_ranges`), + so the relations agree with the sensitive-port findings: an all-ports + rule (``0-65535``, protocol ``ALL`` / ``-1`` / ``*``, Azure ``*`` with + ``Tcp``) is ``ALL_TRAFFIC``; otherwise each of 22, 80, 443 and 3389 + inside the ranges gives its ``ONLY_*`` relation (AWS ``1-1024`` gives + ``ONLY_SSH``, ``ONLY_HTTP`` and ``ONLY_HTTPS``), and ports that include + none of them give ``PORT_RESTRICTED``. A rule with no ports (ICMP, or + no port data) gives nothing. + """ + ranges = _rule_port_ranges(edge) + if not ranges: + return [] + if covers_all_ports(ranges): + return [RelationType.ALL_TRAFFIC] + relations = [rel for port, rel in _PORT_RELATIONS.items() if port_in_ranges(port, ranges)] + return relations or [RelationType.PORT_RESTRICTED] def _infer_sg_rule_relations( @@ -256,7 +281,7 @@ def _infer_sg_rule_relations( relations: list[RelationType] = [] cidr = edge.cidr or "" direction = (edge.direction or "ingress").lower() - is_internet = cidr in ("0.0.0.0/0", "::/0") + is_internet = is_internet_source(cidr) if direction == "ingress": relations.append(RelationType.INGRESS_ALLOWED) @@ -281,33 +306,77 @@ def _infer_sg_rule_relations( AssetType.GKE_CLUSTER, AssetType.K8S_NAMESPACE, } + +# What a cluster or namespace runs: the targets of CLUSTER_CONTAINS_SERVICE +_CLUSTER_MEMBER_TYPES = { + AssetType.K8S_NAMESPACE, + AssetType.K8S_WORKLOAD, + AssetType.K8S_SERVICE, + AssetType.K8S_INGRESS, + AssetType.K8S_SERVICE_ACCOUNT, + AssetType.CONTAINER_SERVICE, + AssetType.TASK_DEFINITION, + AssetType.NODE_GROUP, + AssetType.FARGATE_PROFILE, + AssetType.CLUSTER_ADDON, + AssetType.CAPACITY_PROVIDER, +} + +# Network structure and hierarchy: never the "instance" a subnet contains +_NOT_SUBNET_MEMBER_TYPES = { + AssetType.VPC, + AssetType.VNET, + AssetType.SUBNET, + AssetType.ROUTE_TABLE, + AssetType.NACL, + AssetType.SECURITY_GROUP, + AssetType.NSG, + AssetType.INTERNET_GATEWAY, + AssetType.TRANSIT_GATEWAY, + AssetType.PEERING_CONNECTION, + AssetType.PREFIX_LIST, + AssetType.ROUTER, + AssetType.VPN_GATEWAY, + AssetType.CUSTOMER_GATEWAY, + AssetType.ORGANIZATION, + AssetType.ORG_UNIT, + AssetType.CLOUD_ACCOUNT, + AssetType.RESOURCE_GROUP, +} _ORG_TYPES = {AssetType.ORGANIZATION, AssetType.ORG_UNIT} +def _specific_containment(src: AssetType, tgt: AssetType) -> RelationType | None: + """The containment sub-type whose two endpoint types match, if any.""" + if src in (AssetType.VPC, AssetType.VNET) and tgt == AssetType.SUBNET: + return RelationType.VPC_CONTAINS_SUBNET + if src == AssetType.SUBNET and tgt not in _NOT_SUBNET_MEMBER_TYPES: + return RelationType.SUBNET_CONTAINS_INSTANCE + if src in _CLUSTER_TYPES and tgt in _CLUSTER_MEMBER_TYPES: + return RelationType.CLUSTER_CONTAINS_SERVICE + if src in _ORG_TYPES and tgt == AssetType.CLOUD_ACCOUNT: + return RelationType.ORG_CONTAINS_ACCOUNT + return None + + def _infer_containment_relations( edge: NetworkEdge, assets_by_id: dict[str, CloudAsset] ) -> list[RelationType]: """Infer the containment sub-type from the source and target asset types. A specific relation is used only when both ends match it: VPC / VNet to - subnet, subnet to placed resource, cluster or namespace to workload, - organization or OU to account. Anything else (organization to OU, - account to VPC, VNet to VM, resource group to resource, unresolved - endpoints) is plain ``CONTAINS``. + subnet, subnet to a resource placed in it (not another network + construct such as a route table, NACL or gateway), cluster or namespace + to a workload, service, namespace, node group or add-on, organization + or OU to account. Anything else (organization to OU, account to VPC, + VNet to VM, cluster to subnet, subnet to route table, resource group to + resource, unresolved endpoints) is plain ``CONTAINS``. """ src = assets_by_id.get(edge.source_id) tgt = assets_by_id.get(edge.target_id) if not (src and tgt): return [RelationType.CONTAINS] - if src.asset_type in (AssetType.VPC, AssetType.VNET) and tgt.asset_type == AssetType.SUBNET: - return [RelationType.VPC_CONTAINS_SUBNET] - if src.asset_type == AssetType.SUBNET: - return [RelationType.SUBNET_CONTAINS_INSTANCE] - if src.asset_type in _CLUSTER_TYPES: - return [RelationType.CLUSTER_CONTAINS_SERVICE] - if src.asset_type in _ORG_TYPES and tgt.asset_type == AssetType.CLOUD_ACCOUNT: - return [RelationType.ORG_CONTAINS_ACCOUNT] - return [RelationType.CONTAINS] + return [_specific_containment(src.asset_type, tgt.asset_type) or RelationType.CONTAINS] def _infer_iam_trust_relations( @@ -502,37 +571,39 @@ def _infer_vpc_containment( return relations -def _kms_key_identifiers(key: CloudAsset) -> set[str]: - """Every string a ``kms_key_id`` reference may use for this key.""" - names = {key.id, key.name} +def _kms_key_aliases(key: CloudAsset) -> list[str]: + """Every string other than the ARN a ``kms_key_id`` reference may use.""" + names = [key.name] if key.arn: - names.add(key.arn) - names.add(key.arn.rsplit("/", 1)[-1]) # bare key id - names.update(a for a in key.metadata.get("aliases") or [] if isinstance(a, str)) + names.append(key.arn.rsplit("/", 1)[-1]) # bare key id + names.extend(a for a in key.metadata.get("aliases") or [] if isinstance(a, str)) return names -def _resolve_kms_key(ref: str, all_assets: list[CloudAsset]) -> str: - """Asset id of the KMS key ``ref`` names, or ``ref`` itself if not collected. +def kms_key_index(assets: Iterable[CloudAsset]) -> AssetIndex: + """An :class:`AssetIndex` over the KMS keys in ``assets``. - ``ref`` may be the key's asset id, ARN, bare key id, alias name - (``alias/app``) or alias ARN. + A key is found by its asset id, ARN, name, bare key id, alias name + (``alias/app``) or alias ARN; when two keys share a reference the + first one wins. Build it once per inventory and pass it to + :func:`infer_asset_relations` (``kms_index=``). """ - for other in all_assets: - if other.asset_type == AssetType.KMS_KEY and ref in _kms_key_identifiers(other): - return other.id - return ref + index = AssetIndex() + for key in assets: + if key.asset_type == AssetType.KMS_KEY: + index.add(key.id, key.arn, aliases=_kms_key_aliases(key)) + return index def _infer_kms_encryption( - asset: CloudAsset, all_assets: list[CloudAsset] + asset: CloudAsset, kms_index: AssetIndex ) -> list[tuple[RelationType, str]]: """Infer ``asset ENCRYPTED_BY_KMS key`` for an asset encrypted with a named key. The asset needs ``encryption`` or ``storage_encrypted`` set in its metadata and a ``kms_key_id``. The key reference is resolved to the - collected KMS key asset when there is one; otherwise the raw reference - is the object. + collected KMS key asset through ``kms_index`` when there is one; + otherwise the raw reference is the object. """ md = asset.metadata kms_ref = md.get("kms_key_id") @@ -540,7 +611,8 @@ def _infer_kms_encryption( return [] if not (md.get("encryption") or md.get("storage_encrypted")): return [] - return [(RelationType.ENCRYPTED_BY_KMS, _resolve_kms_key(kms_ref, all_assets))] + key_id = kms_index.find(kms_ref, tails=False) + return [(RelationType.ENCRYPTED_BY_KMS, key_id or kms_ref)] def _infer_metadata_relations(asset: CloudAsset) -> list[tuple[RelationType, str]]: @@ -567,7 +639,10 @@ def _infer_metadata_relations(asset: CloudAsset) -> list[tuple[RelationType, str def infer_asset_relations( - asset: CloudAsset, all_assets: list[CloudAsset] + asset: CloudAsset, + all_assets: list[CloudAsset], + *, + kms_index: AssetIndex | None = None, ) -> list[tuple[RelationType, str]]: """Infer additional semantic relations from asset metadata alone. @@ -575,10 +650,32 @@ def infer_asset_relations( always the subject, so each tuple reads ``asset REL object``. Relations that point at an asset from elsewhere (a VPC containing it) are returned when that other asset is passed in instead. + + ``kms_index`` is :func:`kms_key_index` of ``all_assets``. Pass it when + calling this for many assets of one inventory (or use + :func:`iter_asset_relations`); without it the index is rebuilt on every + call. """ + if kms_index is None: + kms_index = kms_key_index(all_assets) relations: list[tuple[RelationType, str]] = [] relations.extend(_infer_tag_relations(asset)) relations.extend(_infer_vpc_containment(asset, all_assets)) - relations.extend(_infer_kms_encryption(asset, all_assets)) + relations.extend(_infer_kms_encryption(asset, kms_index)) relations.extend(_infer_metadata_relations(asset)) return relations + + +def iter_asset_relations( + assets: list[CloudAsset], +) -> Iterator[tuple[CloudAsset, RelationType, str]]: + """``(asset, relation, object_id)`` for every asset-level relation. + + The asset-level relations of a whole inventory, the way the ontology + and the RAG export both add them: :func:`infer_asset_relations` for + each asset, with the KMS key index built once. + """ + kms_index = kms_key_index(assets) + for asset in assets: + for rel_type, target_id in infer_asset_relations(asset, assets, kms_index=kms_index): + yield asset, rel_type, target_id diff --git a/cloudg/graph/ports.py b/cloudg/graph/ports.py index 87921a9..0ba80a3 100644 --- a/cloudg/graph/ports.py +++ b/cloudg/graph/ports.py @@ -15,6 +15,11 @@ Ports are matched as numbers. A substring test on the raw string would find port 22 in ``"2200-2300"`` and miss it in ``"0-65535"``. + +The module also holds the edge classification helpers shared by the graph +builder, the reachability analysis and the ontology: which edges are filter +rules, which rules are egress, which protocols carry ports and which +sources stand for the whole internet (:func:`is_internet_source`). """ from __future__ import annotations @@ -31,24 +36,35 @@ PortRange = tuple[int, int] ALL_PORTS: tuple[PortRange, ...] = ((MIN_PORT, MAX_PORT),) +# ``port_range`` spelling of every port, as the AWS collector writes it +ALL_PORTS_RANGE = f"{MIN_PORT}-{MAX_PORT}" + # Protocol spellings that mean "every protocol": AWS ``-1`` (the collector -# writes ``ALL``), Azure ``*`` and ``Any``, GCP ``all``. -_ALL_PROTOCOLS = frozenset({"ALL", "-1", "*", "ANY"}) +# writes ``ALL``), Azure ``*`` and ``Any``, GCP ``all``. Compared upper-case. +ALL_PROTOCOLS = frozenset({"ALL", "-1", "*", "ANY"}) # Protocols that carry ports, by name and by IANA number. An empty port list # on a filter rule for one of these allows every port (GCP semantics). -_PORTED_PROTOCOLS = frozenset({"TCP", "UDP", "SCTP", "6", "17", "132"}) +PORTED_PROTOCOLS = frozenset({"TCP", "UDP", "SCTP", "6", "17", "132"}) # Port-less protocols, by name and by IANA number. AWS stores the ICMP type # and code in the port fields. -_PORTLESS_PROTOCOLS = frozenset( +PORTLESS_PROTOCOLS = frozenset( {"ICMP", "ICMPV6", "ESP", "AH", "GRE", "IPIP", "1", "58", "50", "51", "47", "4"} ) -# Edge types whose empty port list means every port. On other edges (for -# example a GCP ``INTERNET_EXPOSED`` edge for a public bucket) an empty -# ``port_range`` only means the collector had no port information. -_FILTER_RULE_EDGES = frozenset({EdgeType.SECURITY_GROUP_RULE.value, EdgeType.NACL_RULE.value}) +# Filter rule edges: an empty port list on them means every port, and they +# are followed by the exposure walk unless they are egress rules. On other +# edges (for example a GCP ``INTERNET_EXPOSED`` edge for a public bucket) an +# empty ``port_range`` only means the collector had no port information. +FILTER_RULE_EDGES = frozenset({EdgeType.SECURITY_GROUP_RULE.value, EdgeType.NACL_RULE.value}) + +# CIDRs that stand for the whole internet +INTERNET_CIDRS = frozenset({"0.0.0.0/0", "::/0"}) + +# Every source spelling that stands for the whole internet, lower-cased: the +# CIDRs above plus the Azure NSG service tags ``Internet``, ``Any`` and ``*``. +INTERNET_SOURCES = INTERNET_CIDRS | {"*", "internet", "any"} _RANGE_RE = re.compile(r"^(\d+)(?:-(\d+))?$") @@ -57,6 +73,22 @@ def _normalise_protocol(protocol: str | None) -> str: return str(protocol or "").strip().upper() +def is_internet_source(value: str | None) -> bool: + """Whether a rule source or CIDR stands for the whole internet. + + True for ``0.0.0.0/0``, ``::/0`` and the Azure NSG service tags + ``Internet``, ``Any`` and ``*``, in any letter case and with surrounding + blanks ignored. The exposure walk, the sensitive-port findings and the + ontology's ``INTERNET_REACHABLE`` relation all use this one test. + """ + return str(value or "").strip().lower() in INTERNET_SOURCES + + +def is_egress(edge_data: Mapping[str, Any]) -> bool: + """Whether a graph edge (or edge attribute mapping) is an egress rule.""" + return str(edge_data.get("direction") or "").lower() == "egress" + + def _parse_token(token: str) -> PortRange | None: """One ``"22"``, ``"1000-2000"`` or ``"*"`` token as a range, or None.""" if token in ("*", "any", "all"): @@ -101,15 +133,15 @@ def parse_port_ranges(port_range: str | None, protocol: str | None = None) -> tu () """ proto = _normalise_protocol(protocol) - if proto in _PORTLESS_PROTOCOLS: + if proto in PORTLESS_PROTOCOLS: return () tokens = [t.strip().lower() for t in str(port_range or "").split(",")] ranges = tuple(r for r in (_parse_token(t) for t in tokens if t) if r is not None) if ranges: return ranges - if proto in _ALL_PROTOCOLS: + if proto in ALL_PROTOCOLS: return ALL_PORTS - if proto in _PORTED_PROTOCOLS and not any(tokens): + if proto in PORTED_PROTOCOLS and not any(tokens): return ALL_PORTS return () @@ -122,7 +154,7 @@ def edge_port_ranges(edge_data: Mapping[str, Any]) -> tuple[PortRange, ...]: returns ``()``, since there the collector simply had no ports to record. """ port_range = str(edge_data.get("port_range") or "") - if not port_range.strip() and edge_data.get("edge_type") not in _FILTER_RULE_EDGES: + if not port_range.strip() and edge_data.get("edge_type") not in FILTER_RULE_EDGES: return () return parse_port_ranges(port_range, edge_data.get("protocol")) @@ -130,3 +162,13 @@ def edge_port_ranges(edge_data: Mapping[str, Any]) -> tuple[PortRange, ...]: def port_in_ranges(port: int, ranges: Iterable[PortRange]) -> bool: """Whether ``port`` falls inside any of the inclusive ``ranges``.""" return any(low <= port <= high for low, high in ranges) + + +def covers_all_ports(ranges: Iterable[PortRange]) -> bool: + """Whether the inclusive ``ranges`` together cover every port 0 to 65535.""" + reach = MIN_PORT - 1 + for low, high in sorted(ranges): + if low > reach + 1: + return False + reach = max(reach, high) + return reach >= MAX_PORT diff --git a/cloudg/graph/rag_export.py b/cloudg/graph/rag_export.py index 93e5b5e..353a53f 100644 --- a/cloudg/graph/rag_export.py +++ b/cloudg/graph/rag_export.py @@ -1,10 +1,15 @@ -"""RAG-ready export — generates chunked, metadata-enriched representations -of the cloud infrastructure graph for retrieval-augmented generation. +"""RAG-ready export: chunked, metadata-enriched representations of the +cloud infrastructure graph for retrieval-augmented generation. Three complementary chunking strategies: 1. Entity-centric: one chunk per cloud asset with 1-hop neighbourhood 2. Community-detection: Louvain clusters with aggregate summaries 3. Relation-group: one chunk per semantic relation group (Network, IAM, etc.) + +Relation-group chunks carry the same relations the ontology holds: the +relations inferred from each edge plus the asset-level ones inferred from +metadata (``ENCRYPTED_BY_KMS``, VPC containment from ``vpc_id``, tag +governance, rotation, EC2 security group membership). """ from __future__ import annotations @@ -21,6 +26,7 @@ RelationGroup, get_relation_group, infer_relations, + iter_asset_relations, ) from cloudg.inventory.dependencies import AssetIndex from cloudg.schema.models import ( @@ -115,8 +121,9 @@ def _index_findings( ) -> dict[str, list[Finding]]: """Group findings by the asset ID they resolve to. - ``index`` matches a finding by ID, ARN, unique name or unique ARN - tail; findings it cannot place stay keyed by their raw resource_id. + ``index`` matches a finding by ID, ARN, unique name, unique ARN tail + or unique ARN resource-part tail (:meth:`AssetIndex.resolve_finding`); + findings it cannot place stay keyed by their raw resource_id. """ findings_by_resource: dict[str, list[Finding]] = defaultdict(list) if findings: @@ -287,33 +294,53 @@ def _count_community_edges(graph: nx.DiGraph, member_set: set[str]) -> tuple[int return internal_edges, external_edges +def _label(asset_id: str, assets_by_id: dict[str, CloudAsset]) -> str: + """Display name of an asset id, or the id itself when it is not an asset.""" + asset = assets_by_id.get(asset_id) + return asset.name if asset else asset_id + + +def _triple( + subject_id: str, predicate: str, object_id: str, assets_by_id: dict[str, CloudAsset] +) -> dict[str, str]: + """A relation triple with display labels and empty edge evidence.""" + return { + "subject": _label(subject_id, assets_by_id), + "subject_id": subject_id, + "predicate": predicate, + "object": _label(object_id, assets_by_id), + "object_id": object_id, + "port_range": "", + "protocol": "", + "cidr": "", + } + + def _group_relation_triples( edges: list[NetworkEdge], assets_by_id: dict[str, CloudAsset], ) -> dict[RelationGroup, list[dict[str, str]]]: - """Group edge triples by their inferred relation group.""" + """Group relation triples by their relation group. + + The triples are the relations inferred from every edge (with the + edge's port, protocol and CIDR as evidence) followed by the asset-level + relations of every asset in ``assets_by_id``, the same two sources the + ontology is built from. + """ group_triples: dict[RelationGroup, list[dict[str, str]]] = defaultdict(list) for edge in edges: - src_name = assets_by_id.get(edge.source_id) - tgt_name = assets_by_id.get(edge.target_id) - src_label = src_name.name if src_name else edge.source_id - tgt_label = tgt_name.name if tgt_name else edge.target_id - for rel in infer_relations(edge, assets_by_id): - group = get_relation_group(rel) - group_triples[group].append( - { - "subject": src_label, - "subject_id": edge.source_id, - "predicate": rel.value, - "object": tgt_label, - "object_id": edge.target_id, - "port_range": edge.port_range or "", - "protocol": edge.protocol or "", - "cidr": edge.cidr or "", - } - ) + triple = _triple(edge.source_id, rel.value, edge.target_id, assets_by_id) + triple["port_range"] = edge.port_range or "" + triple["protocol"] = edge.protocol or "" + triple["cidr"] = edge.cidr or "" + group_triples[get_relation_group(rel)].append(triple) + + for asset, rel, object_id in iter_asset_relations(list(assets_by_id.values())): + group_triples[get_relation_group(rel)].append( + _triple(asset.id, rel.value, object_id, assets_by_id) + ) return group_triples @@ -326,6 +353,49 @@ def _triple_evidence(triple: dict[str, str]) -> str: return f" [{', '.join(parts)}]" +def _relation_group_lines( + group: RelationGroup, triples: list[dict[str, str]], type_counts: dict[str, int] +) -> list[str]: + """Content lines of a relation-group chunk: totals, distribution, triples.""" + content_lines = [ + f"Relation Group: {group.value}", + f"Total relations: {len(triples)}", + "", + "Relation type distribution:", + ] + for rt_name, count in sorted(type_counts.items(), key=lambda x: -x[1]): + content_lines.append(f" {rt_name}: {count}") + + content_lines.append("\nTriples:") + for t in triples[:50]: # Cap for chunk size + content_lines.append( + f" {t['subject']} → {t['predicate']} → {t['object']}{_triple_evidence(t)}" + ) + if len(triples) > 50: + content_lines.append(f" ... and {len(triples) - 50} more") + return content_lines + + +def _write_export(out: Path, all_chunks: list[RAGChunk], counts: dict[str, int]) -> dict[str, Path]: + """Write the JSONL chunks and the metadata index; return both paths.""" + # Write JSONL (one JSON object per line, the usual vector DB input) + chunks_path = out / "rag_chunks.jsonl" + with open(chunks_path, "w") as f: + for chunk in all_chunks: + f.write(json.dumps(chunk.to_dict(), default=str) + "\n") + + index = { + "total_chunks": len(all_chunks), + **counts, + "chunk_ids": [c.chunk_id for c in all_chunks], + "chunk_types": list({c.chunk_type for c in all_chunks}), + } + index_path = out / "rag_metadata_index.json" + with open(index_path, "w") as f: + json.dump(index, f, indent=2) + return {"chunks": chunks_path, "index": index_path} + + # --------------------------------------------------------------------------- # RAG Exporter # --------------------------------------------------------------------------- @@ -532,29 +602,10 @@ def _build_relation_group_chunk( triples: list[dict[str, str]], ) -> RAGChunk: """Assemble content and metadata for one relation-group chunk.""" - content_lines = [ - f"Relation Group: {group.value}", - f"Total relations: {len(triples)}", - "", - ] - - # Summarise relation type distribution type_counts: dict[str, int] = defaultdict(int) for t in triples: type_counts[t["predicate"]] += 1 - - content_lines.append("Relation type distribution:") - for rt_name, count in sorted(type_counts.items(), key=lambda x: -x[1]): - content_lines.append(f" {rt_name}: {count}") - - content_lines.append("\nTriples:") - for t in triples[:50]: # Cap for chunk size - content_lines.append( - f" {t['subject']} → {t['predicate']} → {t['object']}{_triple_evidence(t)}" - ) - - if len(triples) > 50: - content_lines.append(f" ... and {len(triples) - 50} more") + content_lines = _relation_group_lines(group, triples, type_counts) metadata = { "relation_group": group.value, @@ -600,38 +651,19 @@ def export_all( out.mkdir(parents=True, exist_ok=True) assets_by_id = {a.id: a for a in assets} - all_chunks: list[RAGChunk] = [] - - # Strategy 1 entity_chunks = self.export_entity_chunks(assets, edges, findings) - all_chunks.extend(entity_chunks) - - # Strategy 2 community_chunks = self.export_community_chunks(graph, assets_by_id, findings) - all_chunks.extend(community_chunks) - - # Strategy 3 relation_chunks = self.export_relation_chunks(edges, assets_by_id) - all_chunks.extend(relation_chunks) - - # Write JSONL (one JSON object per line — standard for vector DBs) - chunks_path = out / "rag_chunks.jsonl" - with open(chunks_path, "w") as f: - for chunk in all_chunks: - f.write(json.dumps(chunk.to_dict(), default=str) + "\n") - - # Write metadata index - index = { - "total_chunks": len(all_chunks), - "entity_chunks": len(entity_chunks), - "community_chunks": len(community_chunks), - "relation_group_chunks": len(relation_chunks), - "chunk_ids": [c.chunk_id for c in all_chunks], - "chunk_types": list({c.chunk_type for c in all_chunks}), - } - index_path = out / "rag_metadata_index.json" - with open(index_path, "w") as f: - json.dump(index, f, indent=2) + all_chunks = [*entity_chunks, *community_chunks, *relation_chunks] + paths = _write_export( + out, + all_chunks, + { + "entity_chunks": len(entity_chunks), + "community_chunks": len(community_chunks), + "relation_group_chunks": len(relation_chunks), + }, + ) logger.info( "RAG export complete: %d chunks (%d entity, %d community, %d relation) -> %s", @@ -641,8 +673,4 @@ def export_all( len(relation_chunks), out, ) - - return { - "chunks": chunks_path, - "index": index_path, - } + return paths diff --git a/cloudg/graph/reachability.py b/cloudg/graph/reachability.py index 0259642..8634a18 100644 --- a/cloudg/graph/reachability.py +++ b/cloudg/graph/reachability.py @@ -1,8 +1,9 @@ """BFS-based network reachability analysis. Internet exposure follows network-flow edges only. Starting from the -internet sources (``0.0.0.0/0``, ``::/0`` and edges that carry an internet -CIDR), the walk moves along an edge only when traffic can travel that way: +internet sources (``0.0.0.0/0``, ``::/0`` and edges whose CIDR stands for +the internet, see :func:`cloudg.graph.ports.is_internet_source`), the walk +moves along an edge only when traffic can travel that way: - ``INTERNET_EXPOSED``, ``LOAD_BALANCER_TARGET``, ``ROUTE`` and ``PEERING``: source to target, as the edge points. @@ -21,14 +22,29 @@ Hierarchy containment (organization, OU, account, cluster, namespace, resource group) is not traversed. +These rules are available on their own through +:meth:`ReachabilityAnalyzer.flow_successors` (the hops out of one node) and +:func:`network_flow_graph` (every hop as a directed graph), so other +consumers walk the network exactly the way the exposure analysis does. + Identity edges (``GRANTS_ACCESS``, ``ASSUMES_ROLE``, ``IAM_TRUST``, ``IAM_POLICY_ATTACHMENT``) and the other typed inventory edges (``INVOKES``, ``REFERENCES``, ``LOGS_TO`` and so on) never make a resource internet-exposed: a role granted access to a database does not open a -network path to it. Those edges still count for +network path to it. In particular ``INVOKES`` is not followed, so a Lambda +function behind a public API Gateway is not reported as internet-exposed: +the gateway is the exposed resource and the function is only invoked by +it. Those edges still count for :meth:`ReachabilityAnalyzer.compute_blast_radius`, which models what a compromised resource can reach, identity included. +Rule containers (security groups, NSGs, NACLs) and routing constructs +(target groups) are hops on the walk and are marked ``is_internet_exposed`` +like any reached node, since internet traffic does reach them, but they do +not get an exposure finding of their own: a group's open rules are +reported by the sensitive-port findings, and a target group's exposure is +reported on the targets behind it. + The walk is an over-approximation: it does not intersect ports across hops, and a rule whose source is another group or a workload is followed like any other ingress rule. @@ -43,10 +59,35 @@ import networkx as nx -from cloudg.graph.ports import edge_port_ranges, port_in_ranges -from cloudg.schema.models import AssetType, Finding, Severity +from cloudg.graph.ports import ( + FILTER_RULE_EDGES, + INTERNET_CIDRS, + edge_port_ranges, + is_egress, + is_internet_source, + port_in_ranges, +) from cloudg.schema.models import AssetType, EdgeType, Finding, Severity +__all__ = [ + "EXPECTED_EXPOSED_TYPES", + "INTERNET_CIDRS", + "NON_RESOURCE_TYPES", + "ROUTING_CONSTRUCT_TYPES", + "RULE_CONTAINER_TYPES", + "RULE_OPEN_PORT", + "RULE_SENSITIVE_EXPOSURE", + "RULE_UNEXPECTED_EXPOSURE", + "SENSITIVE_ASSET_TYPES", + "SENSITIVE_PORTS", + "FlowHop", + "ReachabilityAnalyzer", + "asset_key", + "finding_id", + "is_internet_source", + "network_flow_graph", +] + logger = logging.getLogger(__name__) # Ports that are considered sensitive when exposed to the internet @@ -81,17 +122,22 @@ AssetType.INTERNET_GATEWAY, } -# Rule containers sit on SG/NSG rule edges but are not reachable workloads; -# their open rules are reported by the sensitive-port findings instead. +# Rule containers sit on SG / NSG / NACL rule edges but are not reachable +# workloads; their open rules are reported by the sensitive-port findings. +# The same types are the traffic filters whose ATTACHED_TO edges the walk +# follows backwards. RULE_CONTAINER_TYPES = { AssetType.SECURITY_GROUP, AssetType.NSG, + AssetType.NACL, } -# Source / CIDR values that stand for the whole internet. Azure NSG rules -# use the service tags ``Internet``, ``Any`` and ``*`` instead of a CIDR. -INTERNET_CIDRS = frozenset({"0.0.0.0/0", "::/0"}) -_INTERNET_EDGE_CIDRS = frozenset({"0.0.0.0/0", "::/0", "*", "internet", "any"}) +# Load-balancer routing constructs: traffic passes through them to the +# targets, which get the exposure findings. +ROUTING_CONSTRUCT_TYPES = {AssetType.TARGET_GROUP} + +# Reached by the walk but never reported as an exposed resource +NON_RESOURCE_TYPES = RULE_CONTAINER_TYPES | ROUTING_CONSTRUCT_TYPES # Edges that carry traffic from source to target as they point _FORWARD_FLOW_EDGES = frozenset( @@ -103,13 +149,8 @@ } ) -# Filter rules: followed source to target unless they are egress rules -_RULE_EDGES = frozenset({EdgeType.SECURITY_GROUP_RULE.value, EdgeType.NACL_RULE.value}) - # ATTACHED_TO targets that filter traffic for the resource attached to them -_TRAFFIC_FILTER_TYPES = frozenset( - {AssetType.SECURITY_GROUP.value, AssetType.NSG.value, AssetType.NACL.value} -) +_TRAFFIC_FILTER_TYPES = frozenset(t.value for t in RULE_CONTAINER_TYPES) _FILTER_RELATIONSHIPS = frozenset({"PROTECTED_BY_SG", "PROTECTED_BY_NACL"}) # ATTACHED_TO sources that pass their traffic on to the resource they attach to @@ -130,19 +171,74 @@ RULE_UNEXPECTED_EXPOSURE = "internet-exposed-unexpected-asset" RULE_OPEN_PORT = "internet-open-sensitive-port" +# One network-flow hop: (next node, attributes of the graph edge used, and +# whether the edge was walked against its direction) +FlowHop = tuple[str, dict[str, Any], bool] + def finding_id(rule: str, asset: str) -> str: """Deterministic finding id: a UUID5 hash of the rule and the asset key. The same rule firing on the same asset gets the same id on every run, so reachability findings can be deduplicated and tracked across scans. + Pass :func:`asset_key` of the node, not the graph node id: collectors + give assets a fresh random id on every collection. """ return str(uuid.uuid5(_FINDING_ID_NAMESPACE, f"{rule}|{asset}")) -def _is_egress(edge_data: dict[str, Any]) -> bool: - """Whether a graph edge is an egress filter rule.""" - return str(edge_data.get("direction") or "").lower() == "egress" +def asset_key(graph: nx.DiGraph, node_id: str) -> str: + """Key of a graph node that stays the same across collections. + + The node's ARN (cloud resource id) when it has one, else the node id. + CIDR and other external placeholder nodes have no ARN; their id is + already stable. + """ + return str(graph.nodes.get(node_id, {}).get("arn") or node_id) + + +def _is_forward_hop(data: dict[str, Any], node_type: str) -> bool: + """Whether traffic at a node of ``node_type`` follows this out-edge.""" + edge_type = data.get("edge_type", "") + if edge_type in _FORWARD_FLOW_EDGES: + return True + if edge_type in FILTER_RULE_EDGES: + return not is_egress(data) + if edge_type == EdgeType.ATTACHED_TO.value: + return node_type in _INTERFACE_TYPES + if edge_type == EdgeType.CONTAINS.value: + return node_type in _NETWORK_PLACEMENT_TYPES + return False + + +def _is_reverse_hop(data: dict[str, Any], node_is_filter: bool) -> bool: + """Whether traffic at a filter node follows this in-edge backwards.""" + return data.get("edge_type") == EdgeType.ATTACHED_TO.value and ( + node_is_filter or data.get("relationship") in _FILTER_RELATIONSHIPS + ) + + +def network_flow_graph(graph: nx.DiGraph) -> nx.DiGraph: + """The network-flow hops of ``graph`` as a new directed graph. + + Every node of ``graph`` is copied with its attributes. There is an edge + ``u -> v`` exactly when :meth:`ReachabilityAnalyzer.flow_successors` + yields ``v`` for ``u``, so walking this graph forward follows the same + rules as the internet-exposure analysis (see the module docstring). + Each edge carries ``edge_type`` (the type of the underlying edge) and + ``reversed`` (True for an ``ATTACHED_TO`` edge walked from the filter + back to the resource attached to it). ``graph`` is not modified. + """ + analyzer = ReachabilityAnalyzer(graph) + flow = nx.DiGraph() + flow.add_nodes_from(graph.nodes(data=True)) + for node_id in graph.nodes: + for nxt, data, walked_back in analyzer.flow_hops(node_id): + if not flow.has_edge(node_id, nxt): + flow.add_edge( + node_id, nxt, edge_type=data.get("edge_type", ""), reversed=walked_back + ) + return flow class ReachabilityAnalyzer: @@ -165,17 +261,18 @@ def find_internet_exposed(self) -> set[str]: of every non-egress edge whose CIDR stands for the internet. The walk follows only the network-flow hops described in the module docstring, so resources reached only through IAM or other non-network edges are - not reported. + not reported. Every node returned is marked ``is_internet_exposed`` + in the graph, rule containers and target groups included. Returns: Set of node IDs that are internet-exposed. """ exposed: set[str] = set() - queue = deque(self._internet_entry_points()) + queue = deque(self.internet_entry_points()) seen = set(queue) while queue: node_id = queue.popleft() - for nxt in self._flow_successors(node_id): + for nxt in self.flow_successors(node_id): exposed.add(nxt) if nxt not in seen: seen.add(nxt) @@ -188,38 +285,55 @@ def find_internet_exposed(self) -> set[str]: logger.info("Found %d internet-exposed nodes", len(exposed)) return exposed - def _internet_entry_points(self) -> set[str]: - """Nodes that represent the internet as a traffic source.""" + def internet_entry_points(self) -> set[str]: + """Nodes that represent the internet as a traffic source. + + Nodes named ``0.0.0.0/0`` or ``::/0``, plus the source of every + non-egress edge whose ``cidr`` passes + :func:`~cloudg.graph.ports.is_internet_source` (which adds the Azure + service tags ``Internet``, ``Any`` and ``*``). + """ entry_points = { node_id for node_id, data in self._graph.nodes(data=True) if data.get("name", node_id) in INTERNET_CIDRS } for source, _target, data in self._graph.edges(data=True): - if str(data.get("cidr") or "").lower() in _INTERNET_EDGE_CIDRS and not _is_egress(data): + if is_internet_source(data.get("cidr")) and not is_egress(data): entry_points.add(source) return entry_points - def _flow_successors(self, node_id: str) -> Iterator[str]: - """Nodes that traffic arriving at ``node_id`` can travel on to.""" + def flow_hops(self, node_id: str) -> Iterator[FlowHop]: + """The network-flow hops out of ``node_id``, with the edge each uses. + + Yields ``(next_node, edge_data, reversed)``: out-edges that carry + traffic forward (``reversed`` False) and ``ATTACHED_TO`` in-edges + walked back from a traffic filter to the resource attached to it + (``reversed`` True). A node may appear more than once when several + edges lead to it. Unknown nodes yield nothing. + """ + if node_id not in self._graph: + return node_type = self._graph.nodes[node_id].get("asset_type", "") for _source, target, data in self._graph.out_edges(node_id, data=True): - edge_type = data.get("edge_type", "") - if ( - edge_type in _FORWARD_FLOW_EDGES - or (edge_type in _RULE_EDGES and not _is_egress(data)) - or (edge_type == EdgeType.ATTACHED_TO.value and node_type in _INTERFACE_TYPES) - or (edge_type == EdgeType.CONTAINS.value and node_type in _NETWORK_PLACEMENT_TYPES) - ): - yield target + if _is_forward_hop(data, node_type): + yield target, data, False # Traffic admitted by a security group / NSG / NACL reaches the # resources attached to it: walk those ATTACHED_TO edges backwards. is_filter = node_type in _TRAFFIC_FILTER_TYPES for source, _target, data in self._graph.in_edges(node_id, data=True): - if data.get("edge_type") == EdgeType.ATTACHED_TO.value and ( - is_filter or data.get("relationship") in _FILTER_RELATIONSHIPS - ): - yield source + if _is_reverse_hop(data, is_filter): + yield source, data, True + + def flow_successors(self, node_id: str) -> Iterator[str]: + """Nodes that traffic arriving at ``node_id`` can travel on to. + + This is the single-step form of the exposure walk; the rules are in + the module docstring. See :meth:`flow_hops` for the edges used and + :func:`network_flow_graph` for the whole flow graph at once. + """ + for nxt, _data, _reversed in self.flow_hops(node_id): + yield nxt def compute_blast_radius(self, node_id: str) -> dict[str, Any]: """Compute blast radius: all nodes reachable from a given node. @@ -299,7 +413,7 @@ def _exposure_finding(self, node_id: str) -> Finding | None: # Check for sensitive asset types exposed to internet if asset_type in SENSITIVE_ASSET_TYPES: return self._sensitive_exposure_finding(node_id, node_data, asset_type) - if asset_type not in EXPECTED_EXPOSED_TYPES | RULE_CONTAINER_TYPES: + if asset_type not in EXPECTED_EXPOSED_TYPES | NON_RESOURCE_TYPES: # Non-database but unexpected exposure return self._unexpected_exposure_finding(node_id, node_data, asset_type) return None @@ -310,7 +424,7 @@ def _sensitive_exposure_finding( """Critical finding for an internet-exposed sensitive data store.""" name = node_data.get("name", node_id) return Finding( - id=finding_id(RULE_SENSITIVE_EXPOSURE, node_id), + id=finding_id(RULE_SENSITIVE_EXPOSURE, asset_key(self._graph, node_id)), resource_id=node_id, resource_arn=node_data.get("arn", ""), severity=Severity.CRITICAL, @@ -336,7 +450,7 @@ def _unexpected_exposure_finding( """High finding for an unexpectedly internet-exposed resource.""" name = node_data.get("name", node_id) return Finding( - id=finding_id(RULE_UNEXPECTED_EXPOSURE, node_id), + id=finding_id(RULE_UNEXPECTED_EXPOSURE, asset_key(self._graph, node_id)), resource_id=node_id, resource_arn=node_data.get("arn", ""), severity=Severity.HIGH, @@ -356,18 +470,21 @@ def _unexpected_exposure_finding( def _sensitive_port_findings(self) -> list[Finding]: """Findings for sensitive ports open to the internet on edges. - Ports are compared as numbers against the ranges parsed from the - edge's ``port_range`` and ``protocol`` (see :mod:`cloudg.graph.ports`), - plus any ``ports`` list on the edge. Egress rules are skipped: they - point from a group to the internet and open nothing inbound. Each - edge yields at most one finding per sensitive port. + An edge qualifies when its ``cidr`` stands for the internet + (:func:`~cloudg.graph.ports.is_internet_source`: ``0.0.0.0/0``, + ``::/0`` or the Azure tags ``Internet``, ``Any`` and ``*``). Ports are + compared as numbers against the ranges parsed from the edge's + ``port_range`` and ``protocol`` (see :mod:`cloudg.graph.ports`), plus + any ``ports`` list on the edge (graphs built by hand rather than by + :class:`~cloudg.graph.builder.GraphBuilder` may carry one). Egress + rules are skipped: they point from a group to the internet and open + nothing inbound. Each edge yields at most one finding per sensitive + port. """ findings: list[Finding] = [] for source, target, data in self._graph.edges(data=True): - cidr = data.get("cidr", "") - if cidr not in ("0.0.0.0/0", "::/0"): - continue - if str(data.get("direction") or "").lower() == "egress": + cidr = str(data.get("cidr") or "") + if not is_internet_source(cidr) or is_egress(data): continue ports = data.get("ports") or [] @@ -381,18 +498,22 @@ def _sensitive_port_findings(self) -> list[Finding]: def _open_port_finding( self, source: str, target: str, cidr: str, port: int, service: str ) -> Finding: - """Critical finding for one sensitive port open from 0.0.0.0/0.""" + """Critical finding for one sensitive port open from the internet.""" target_data = self._graph.nodes.get(target, {}) + # The IPv4 and IPv6 "any" CIDRs keep the historical 0.0.0.0/0 wording; + # an Azure service tag is named as written. + origin = "0.0.0.0/0" if cidr in INTERNET_CIDRS else cidr + edge_key = f"{asset_key(self._graph, source)}->{asset_key(self._graph, target)}" return Finding( - # One finding per rule edge and port: the asset key is the edge - id=finding_id(f"{RULE_OPEN_PORT}:{port}", f"{source}->{target}"), + # One finding per rule edge and port, keyed by the stable asset keys + id=finding_id(f"{RULE_OPEN_PORT}:{port}", edge_key), resource_id=target, resource_arn=target_data.get("arn", ""), severity=Severity.CRITICAL, - title=f"Security group allows {service} (port {port}) from 0.0.0.0/0", + title=f"Security group allows {service} (port {port}) from {origin}", description=( f"A security group rule allows inbound traffic on " - f"port {port} ({service}) from 0.0.0.0/0. This is a " + f"port {port} ({service}) from {origin}. This is a " f"common attack vector." ), evidence=f"Edge from {source} to {target}, port {port}, cidr {cidr}", diff --git a/cloudg/ingest.py b/cloudg/ingest.py index 5ad868d..9796861 100644 --- a/cloudg/ingest.py +++ b/cloudg/ingest.py @@ -33,7 +33,7 @@ def parse_report(tool: str, path: str | Path) -> list[Finding]: Args: tool: One of ``prowler``, ``scoutsuite``, ``checkov``, ``trivy``. - path: The tool's native output — a report file or output directory. + path: The tool's native output: a report file or output directory. Returns: List of normalised Finding objects (empty when nothing parses). diff --git a/cloudg/inventory/__init__.py b/cloudg/inventory/__init__.py index 29540f2..bab4a56 100644 --- a/cloudg/inventory/__init__.py +++ b/cloudg/inventory/__init__.py @@ -1,7 +1,7 @@ """Scanner-independent infrastructure inventory mapping. This package maps everything deployed (or default) in a cloud account and -how it interlinks — without running any security scanner — across every +how it interlinks, without running any security scanner, across every account of an AWS Organization / Control Tower landing zone when asked. Its output can later be merged with scanner findings to build asset and compliance maps, and queried for interdependencies (DependencyGraph). diff --git a/cloudg/inventory/aws_deep.py b/cloudg/inventory/aws_deep.py index aaec7ad..e00f5e9 100644 --- a/cloudg/inventory/aws_deep.py +++ b/cloudg/inventory/aws_deep.py @@ -1,4 +1,4 @@ -"""Deep AWS inventory collector — full-account resource enumeration. +"""Deep AWS inventory collector: full-account resource enumeration. Extends the standard :class:`AsyncAWSCollector` with: @@ -14,7 +14,7 @@ Analyzer, WAF, Network Firewall, Shield, CloudTrail, Detective), and platform/data/DNS/deployment services. - A catch-all sweep over the Resource Groups Tagging API, which returns - every taggable resource in the region — so services without a dedicated + every taggable resource in the region, so services without a dedicated collector still appear on the map instead of silently missing. Global services (IAM, S3, CloudFront, Route 53, Shield, CloudFront-scope @@ -414,7 +414,7 @@ async def _collect_transit_gateways(self) -> list[CloudAsset]: ) ) - # Attachments (VPCs, VPNs, peerings — possibly in other accounts) + # Attachments (VPCs, VPNs, peerings, possibly in other accounts) try: await self._link_tgw_attachments(ec2, assets) except Exception as exc: diff --git a/cloudg/inventory/aws_services/application/_common.py b/cloudg/inventory/aws_services/application/_common.py index 06ab6c0..068d716 100644 --- a/cloudg/inventory/aws_services/application/_common.py +++ b/cloudg/inventory/aws_services/application/_common.py @@ -27,7 +27,46 @@ _MAX_PROTECTED_RESOURCES = 5000 _ARN_ACCOUNT_RE = re.compile(r"^arn:aws[a-zA-Z-]*:[a-z0-9-]+:[a-z0-9-]*:(\d{12}):") -_ALARM_RULE_RE = re.compile(r"(?:ALARM|OK|INSUFFICIENT_DATA)\(\s*\"?([^\")]+?)\"?\s*\)") +# Composite alarm rules reference child alarms as ALARM(name), OK(name) or +# INSUFFICIENT_DATA(name), optionally quoted. Parsed by a single left to right +# pass instead of a regular expression: every regex form of this backtracked +# on hostile rules ("OK(" followed by thousands of spaces, or "OK(" repeated). +_ALARM_STATES = ("INSUFFICIENT_DATA", "ALARM", "OK") + + +def _alarm_child_name(raw: str) -> str: + name = raw.strip() + if name.startswith('"'): + name = name[1:] + if name.endswith('"'): + name = name[:-1] + name = name.strip() + return "" if '"' in name else name + + +def alarm_rule_children(rule: str) -> list[str]: + """Child alarm names referenced by a composite alarm rule, in order.""" + out: list[str] = [] + i, n = 0, len(rule or "") + while i < n: + open_at = rule.find("(", i) + if open_at < 0: + break + if not any(rule.endswith(state, 0, open_at) for state in _ALARM_STATES): + i = open_at + 1 + continue + j = open_at + 1 + while j < n and rule[j] not in "()": + j += 1 + if j < n and rule[j] == ")": + name = _alarm_child_name(rule[open_at + 1 : j]) + if name: + out.append(name) + j += 1 + i = j + return out + + _ASG_POLICY_RE = re.compile(r"autoScalingGroupName/([^:]+)") _AWS_OWNED_DOC_PREFIXES = ( "AWS-", diff --git a/cloudg/inventory/aws_services/application/alarms.py b/cloudg/inventory/aws_services/application/alarms.py index cb33571..1a2ee86 100644 --- a/cloudg/inventory/aws_services/application/alarms.py +++ b/cloudg/inventory/aws_services/application/alarms.py @@ -7,9 +7,9 @@ from cloudg.inventory.aws_services._base import rel from cloudg.inventory.aws_services.application._common import ( - _ALARM_RULE_RE, _ASG_POLICY_RE, ApplicationBase, + alarm_rule_children, ) from cloudg.schema.models import AssetType, CloudAsset, EdgeType @@ -71,7 +71,7 @@ def _alarm_target_relations( "DEPENDS_ON", description="composite alarm child", ) - for child in _ALARM_RULE_RE.findall(alarm.get("AlarmRule") or "") + for child in alarm_rule_children(alarm.get("AlarmRule") or "") ] return children, [] monitored = self._alarm_dimension_targets( diff --git a/cloudg/inventory/azure_deep.py b/cloudg/inventory/azure_deep.py index abe4e25..604abd1 100644 --- a/cloudg/inventory/azure_deep.py +++ b/cloudg/inventory/azure_deep.py @@ -1,4 +1,4 @@ -"""Deep Azure inventory collector — full-subscription resource enumeration. +"""Deep Azure inventory collector: full-subscription resource enumeration. Primary path: **Azure Resource Graph** (``azure-mgmt-resourcegraph``). A few paginated KQL queries per subscription return every resource with its full diff --git a/cloudg/inventory/azure_graph/__init__.py b/cloudg/inventory/azure_graph/__init__.py index 4ab936b..1b701e4 100644 --- a/cloudg/inventory/azure_graph/__init__.py +++ b/cloudg/inventory/azure_graph/__init__.py @@ -5,9 +5,9 @@ both far faster and far richer than enumerating each service SDK. This package holds the provider-neutral half of the Azure inventory: -- :func:`run_query` (``query``) — paginated ARG queries (``skipToken``) +- :func:`run_query` (``query``): paginated ARG queries (``skipToken``) with 429 back-off, against any client exposing ``resources(request)``. -- :class:`AzureAssetBuilder` (``builder``) — turns ARM-shaped rows (from +- :class:`AzureAssetBuilder` (``builder``): turns ARM-shaped rows (from ARG *or* from SDK models serialised with :func:`to_rest`) into :class:`CloudAsset` objects carrying typed ``metadata["relations"]`` and ``metadata["aliases"]`` for the diff --git a/cloudg/inventory/azure_graph/query.py b/cloudg/inventory/azure_graph/query.py index 7b8a576..7741937 100644 --- a/cloudg/inventory/azure_graph/query.py +++ b/cloudg/inventory/azure_graph/query.py @@ -1,4 +1,13 @@ -"""Azure Resource Graph querying: paginated KQL with 429 back-off.""" +"""Azure Resource Graph querying: paginated KQL with 429 back-off. + +Resource Graph allows ~15 queries per 5-second window per user. Every page +takes a token from the shared ``azure/*/*/resourcegraph`` bucket (3 rps, +burst 15: cloudg.resilience), throttled pages wait for the server's +``Retry-After`` / ``x-ms-user-quota-resets-after`` (or exponential back-off +when there is none), and clients built by +:func:`default_graph_client_factory` also read the quota headers of every +response to pause before the quota runs out. +""" from __future__ import annotations @@ -8,13 +17,18 @@ from typing import Any, Callable from cloudg.inventory.azure_graph.helpers import _kql, to_rest +from cloudg.resilience.errors import retry_after +from cloudg.resilience.governor import get_governor +from cloudg.resilience.limiter import Scope def default_graph_client_factory(credential: Any) -> Any: """Build a real ``ResourceGraphClient`` (raises ImportError when absent).""" from azure.mgmt.resourcegraph import ResourceGraphClient # type: ignore[import-not-found] - return ResourceGraphClient(credential) + from cloudg.resilience.azure import build_azure_client + + return build_azure_client(ResourceGraphClient, credential) def _make_request( @@ -73,18 +87,46 @@ class QueryLimits: max_retries: int = 5 +_GRAPH_SCOPE = Scope("azure", None, None, "resourcegraph") + + def _fetch_page(client: Any, request: Any, max_retries: int, sleep: Callable[[float], None]) -> Any: - """One ``resources`` call, retrying throttled (429) responses.""" + """One ``resources`` call, retrying throttled (429) responses. + + Waits at least the server's Retry-After / quota-reset hint, otherwise + exponential back-off (2s, 4s, ... capped at 30s). Clients built by + :func:`default_graph_client_factory` already retry 429 in their + azure-core pipeline (and report every attempt), so a 429 that reaches + this function from them is final and is not retried again. + """ + gov = get_governor() + # Clients from default_graph_client_factory pace every HTTP attempt in + # their pipeline; others (custom factories) are paced here per page. + paced_by_client = bool(getattr(client, "_cloudg_throttle_policy", False)) attempt = 0 while True: + if not paced_by_client and gov.provider_enabled("azure"): + gov.check(_GRAPH_SCOPE) + gov.limiter.acquire_sync(_GRAPH_SCOPE, sleep=sleep) try: - return client.resources(request) + response = client.resources(request) except Exception as exc: - if _status_code(exc) == 429 and attempt < max_retries: - attempt += 1 - sleep(min(2.0**attempt, 30.0)) - continue + if not paced_by_client and _status_code(exc) == 429: + hint = retry_after(exc) + gov.on_throttle(_GRAPH_SCOPE, hint, str(exc)[:200]) + if attempt < max_retries: + attempt += 1 + gov.record_retry(_GRAPH_SCOPE) + # The limiter, paused for the hint, waits it out when enabled + delay = min(2.0**attempt, 30.0) + if not gov.provider_enabled("azure"): + delay = max(delay, min(hint or 0.0, 300.0)) + sleep(delay) + continue + gov.on_final_error(_GRAPH_SCOPE, exc, "Resource Graph query") raise + gov.on_success(_GRAPH_SCOPE) + return response def run_query( diff --git a/cloudg/inventory/dependencies.py b/cloudg/inventory/dependencies.py index 62ddd6b..ad34239 100644 --- a/cloudg/inventory/dependencies.py +++ b/cloudg/inventory/dependencies.py @@ -61,14 +61,54 @@ _STRUCTURAL_EDGES = {EdgeType.CONTAINS.value, EdgeType.PEERING.value} -class AssetIndex: - """Resolve asset references to asset IDs. +def _arn_resource_tails(arn: str) -> set[str]: + """Short forms of an ARN's resource part a reference may use. + + For ``arn:aws:rds:us-east-1:123456789012:db:orders`` these are + ``db:orders`` and ``orders``; for ``...:instance/i-0web1`` they are + ``instance/i-0web1`` and ``i-0web1``; for + ``...:function:api-handler`` they are ``function:api-handler`` and + ``api-handler``. The resource part and its last ``:`` and ``/`` + segments are returned. Strings that are not ARNs, and the trivial + segments ``""`` and ``"*"``, give nothing. + """ + parts = arn.split(":", 5) + if len(parts) < 6 or parts[0] != "arn": + return set() + resource = parts[5] + tails = {resource, resource.rsplit(":", 1)[-1], resource.rsplit("/", 1)[-1]} + return {t for t in tails if t and t != "*"} + + +def _unique(table: dict[str, set[str]], ref: str) -> str | None: + matches = table.get(ref) + return next(iter(matches)) if matches and len(matches) == 1 else None - A reference is tried, in order, as an internal ID, an ARN / resource - ID (first asset wins), a unique name and finally a unique ARN tail - (the part after the last ``/``). Names and tails shared by several - assets are ambiguous and never match. The tables are built once, so - resolving every finding of a scan stays linear. + +class AssetIndex: + """Resolve asset references to asset IDs: the one asset matcher. + + The matching tiers, tried in this order for every reference: + + 1. internal ID; + 2. ARN / resource ID, or an alias given to :meth:`add` (first asset + wins); + 3. unique name; + 4. unique ARN tail: the part after the last ``/`` (``tails``); + 5. unique ARN resource-part tail (``arn_parts``): ``db:orders``, + ``function:api-handler``, ``instance/i-0web1`` and the last ``:`` or + ``/`` segment of the resource part (``orders``, ``api-handler``, + ``i-0web1``); + 6. unique name ignoring letter case (``casefold``). + + Names and tails shared by several assets are ambiguous and never match. + Tiers 5 and 6 are off unless asked for, and they are skipped for a + reference that is the exact name of more than one asset, so an + ambiguous name stays unresolved. With the defaults, :meth:`find`, + :meth:`resolve` and :meth:`resolve_finding` behave exactly as the index + always has; turning the extra tiers on only resolves references that + were unresolved before. The tables are built once, so resolving every + finding of a scan stays linear. """ def __init__(self, assets: Iterable[CloudAsset] = ()) -> None: @@ -76,46 +116,123 @@ def __init__(self, assets: Iterable[CloudAsset] = ()) -> None: self._arns: dict[str, str] = {} self._names: dict[str, set[str]] = {} self._tails: dict[str, set[str]] = {} + self._arn_parts: dict[str, set[str]] = {} + self._folded_names: dict[str, set[str]] = {} for a in assets: self.add(a.id, a.arn, a.name) - def add(self, asset_id: str, arn: str | None = None, name: str | None = None) -> None: - """Index one asset by its ID, ARN, name and ARN tail.""" + def add( + self, + asset_id: str, + arn: str | None = None, + name: str | None = None, + *, + aliases: Iterable[str] = (), + ) -> None: + """Index one asset by its ID, ARN, name, ARN tails and ``aliases``. + + ``aliases`` are extra exact references (a KMS alias, a bare key id) + matched in the ARN tier: the first asset indexed under one wins. + """ self._ids.add(asset_id) if arn: self._arns.setdefault(arn, asset_id) self._tails.setdefault(arn.rsplit("/", 1)[-1], set()).add(asset_id) + for part in _arn_resource_tails(arn): + self._arn_parts.setdefault(part, set()).add(asset_id) + for alias in aliases: + if alias: + self._arns.setdefault(alias, asset_id) if name: self._names.setdefault(name, set()).add(asset_id) + self._folded_names.setdefault(name.casefold(), set()).add(asset_id) - def _lookup(self, refs: Iterable[str | None]) -> str | None: - candidates = [r for r in refs if r] + def _exact(self, candidates: list[str]) -> str | None: + """Tiers 1 and 2: internal ID, then ARN or alias.""" for ref in candidates: if ref in self._ids: return ref for ref in candidates: if ref in self._arns: return self._arns[ref] - for table in (self._names, self._tails): - for ref in candidates: - matches = table.get(ref) - if matches and len(matches) == 1: - return next(iter(matches)) return None + def _unique_in(self, table: dict[str, set[str]], refs: Iterable[str]) -> str | None: + for ref in refs: + found = _unique(table, ref) + if found: + return found + return None + + def _lookup( + self, + refs: Iterable[str | None], + *, + tails: bool = True, + arn_parts: bool = False, + casefold: bool = False, + ) -> str | None: + candidates = [r for r in refs if r] + found = self._exact(candidates) or self._unique_in(self._names, candidates) + if not found and tails: + found = self._unique_in(self._tails, candidates) + if found or not (arn_parts or casefold): + return found + # The opt-in tiers never pick one of several assets sharing a name + loose = [r for r in candidates if r not in self._names] + if arn_parts: + found = self._unique_in(self._arn_parts, loose) + if not found and casefold: + found = self._unique_in(self._folded_names, (r.casefold() for r in loose)) + return found + + def find( + self, + ref: str | None, + *, + tails: bool = True, + arn_parts: bool = False, + casefold: bool = False, + ) -> str | None: + """ID of the asset ``ref`` names, or ``None`` if unknown or ambiguous. + + Args: + ref: An internal ID, ARN / resource ID, alias, name or ARN tail. + tails: Try the unique ARN tail after the last ``/`` (tier 4). + arn_parts: Try the unique ARN resource-part forms (tier 5): + ``db:orders``, ``function:api-handler``, ``instance/i-0web1``, + ``orders``. + casefold: Finally try a unique name ignoring letter case (tier 6). + """ + return self._lookup((ref,), tails=tails, arn_parts=arn_parts, casefold=casefold) + def resolve(self, ref: str | None) -> str | None: - """ID of the asset ``ref`` names, or ``None`` if unknown or ambiguous.""" + """ID of the asset ``ref`` names: ID, ARN, unique name, unique ARN tail.""" return self._lookup((ref,)) - def resolve_finding(self, finding: Finding) -> str | None: - """ID of the asset a finding affects. + def resolve_finding( + self, + finding: Finding, + *, + tails: bool = True, + arn_parts: bool = False, + casefold: bool = False, + ) -> str | None: + """ID of the asset a finding affects, or ``None``. Scanner output often puts an ARN or a display name in ``resource_id``, so both ``resource_id`` and ``resource_arn`` are tried at every tier: ID first, then ARN, then unique name, then - unique ARN tail. + unique ARN tail, then (when asked for) the ARN resource-part tails + and the case-insensitive name. The ontology and the RAG export call + this with the defaults. """ - return self._lookup((finding.resource_id, finding.resource_arn)) + return self._lookup( + (finding.resource_id, finding.resource_arn), + tails=tails, + arn_parts=arn_parts, + casefold=casefold, + ) @dataclass @@ -167,7 +284,7 @@ def __init__( def find(self, ref: str) -> CloudAsset | None: """Find an asset by internal ID, ARN / resource ID, unique name or - unique ARN tail (see :class:`AssetIndex`).""" + unique ARN tail (see :meth:`AssetIndex.resolve`).""" if self._index is None: self._index = AssetIndex(self.assets.values()) asset_id = self._index.resolve(ref) @@ -319,8 +436,8 @@ def cross_account_edges(assets: list[CloudAsset], edges: list[NetworkEdge]) -> l return out -def security_coverage(assets: list[CloudAsset], edges: list[NetworkEdge]) -> dict[str, Any]: - """Which security services run where, and which workloads no scanner covers.""" +def _service_matrix(assets: list[CloudAsset]) -> dict[str, dict[str, dict[str, bool]]]: + """``account -> region -> security service -> enabled`` for the service assets.""" matrix: dict[str, dict[str, dict[str, bool]]] = {} for a in assets: svc = a.metadata.get("security_service") @@ -329,7 +446,13 @@ def security_coverage(assets: list[CloudAsset], edges: list[NetworkEdge]) -> dic acct = a.account_id or "unknown" cell = matrix.setdefault(acct, {}).setdefault(a.region, {}) cell[svc] = cell.get(svc, False) or bool(a.metadata.get("enabled")) + return matrix + +def _unscanned_workloads( + assets: list[CloudAsset], edges: list[NetworkEdge] +) -> list[dict[str, Any]]: + """Scannable workloads no enabled vulnerability scanner MONITORS.""" scanners = { a.id for a in assets @@ -339,7 +462,7 @@ def security_coverage(assets: list[CloudAsset], edges: list[NetworkEdge]) -> dic e.target_id for e in edges if e.edge_type == EdgeType.MONITORS and e.source_id in scanners } scannable = (AssetType.EC2, AssetType.CONTAINER_REGISTRY, AssetType.LAMBDA_FUNCTION) - unscanned = [ + return [ { "name": a.name, "arn": a.arn, @@ -350,13 +473,24 @@ def security_coverage(assets: list[CloudAsset], edges: list[NetworkEdge]) -> dic for a in assets if a.asset_type in scannable and a.id not in scanned ] + + +def _unprotected_entry_points( + assets: list[CloudAsset], edges: list[NetworkEdge] +) -> list[dict[str, Any]]: + """Internet-facing load balancers, APIs and CDNs no WAF PROTECTS.""" protectable = (AssetType.LOAD_BALANCER, AssetType.API_GATEWAY, AssetType.CLOUDFRONT) protected = {e.target_id for e in edges if e.edge_type == EdgeType.PROTECTS} - unprotected = [ + return [ {"name": a.name, "arn": a.arn, "type": a.asset_type.value, "account_id": a.account_id} for a in assets if a.asset_type in protectable and a.is_internet_exposed and a.id not in protected ] + + +def security_coverage(assets: list[CloudAsset], edges: list[NetworkEdge]) -> dict[str, Any]: + """Which security services run where, and which workloads no scanner covers.""" + matrix = _service_matrix(assets) gaps = sorted( f"{acct}/{region}: {svc}" for acct, regions in matrix.items() @@ -367,6 +501,6 @@ def security_coverage(assets: list[CloudAsset], edges: list[NetworkEdge]) -> dic return { "services_by_account_region": matrix, "gaps": gaps, - "workloads_without_vulnerability_scanning": unscanned, - "internet_facing_without_waf": unprotected, + "workloads_without_vulnerability_scanning": _unscanned_workloads(assets, edges), + "internet_facing_without_waf": _unprotected_entry_points(assets, edges), } diff --git a/cloudg/inventory/gcp_deep.py b/cloudg/inventory/gcp_deep.py index 6974927..f24dc74 100644 --- a/cloudg/inventory/gcp_deep.py +++ b/cloudg/inventory/gcp_deep.py @@ -26,6 +26,7 @@ from cloudg.inventory.catalogs import asset_type_map, load_catalog from cloudg.collectors.gcp import _GCP_ASSET_TYPE_MAP, GCPCollector from cloudg.coverage import ServiceStatus +from cloudg.resilience.errors import describe_error from cloudg.inventory.gcp_relations import apply_iam_policies, merge_gcp_principals from cloudg.schema.models import AssetType, CloudAsset @@ -73,7 +74,9 @@ async def _enrich(self, assets: list[CloudAsset]) -> list[CloudAsset]: if error is not None: logger.error("GCP IAM policy search failed for %s: %s", self._scope, error) status = ServiceStatus.PARTIAL if policies else ServiceStatus.FAILED - self._record("gcp_iam_policies", status, len(policies), str(error), start=start) + self._record( + "gcp_iam_policies", status, len(policies), describe_error(error), start=start + ) else: self._record("gcp_iam_policies", ServiceStatus.SUCCESS, len(policies), start=start) created = apply_iam_policies(assets, policies) diff --git a/cloudg/inventory/gcp_relations/names.py b/cloudg/inventory/gcp_relations/names.py index 73714bc..4fc8348 100644 --- a/cloudg/inventory/gcp_relations/names.py +++ b/cloudg/inventory/gcp_relations/names.py @@ -60,7 +60,7 @@ _SERVICE_ALIASES = {"sqladmin": "cloudsql"} -# (path fragment, CAI service) — first match wins for relative names +# (path fragment, CAI service); first match wins for relative names _RELATIVE_SERVICES: tuple[tuple[str, str], ...] = ( ("/cryptoKeys/", "cloudkms"), ("/keyRings/", "cloudkms"), diff --git a/cloudg/inventory/kubernetes.py b/cloudg/inventory/kubernetes.py index 8e1f1cd..f6f4ec5 100644 --- a/cloudg/inventory/kubernetes.py +++ b/cloudg/inventory/kubernetes.py @@ -169,7 +169,7 @@ def __init__( self._region = region self._account_id = account_id self.assets: list[CloudAsset] = [] - # ns, kind, name, pod labels — matched against service selectors + # ns, kind, name, pod labels, matched against service selectors self._workloads: list[tuple[str, str, str, dict[str, str]]] = [] def _add( diff --git a/cloudg/inventory/linker.py b/cloudg/inventory/linker.py index 31680c2..5c23900 100644 --- a/cloudg/inventory/linker.py +++ b/cloudg/inventory/linker.py @@ -1,8 +1,8 @@ """Cross-service relationship linker. -Derives the edges that interlink an inventory — attachment, containment, +Derives the edges that interlink an inventory (attachment, containment, routing, invocation, identity, protection, monitoring, governance and -generic references — purely from collected asset metadata. No cloud API +generic references) purely from collected asset metadata. No cloud API calls and no scanners: given any list of CloudAssets (AWS, Azure, GCP, or mixed), it produces the edges that turn a flat inventory into a connected map. @@ -46,6 +46,7 @@ IdentifierIndex, _image_registry_host, ) +from cloudg.inventory.linker_rules import RuleLinksMixin from cloudg.schema.models import AssetType, CloudAsset, EdgeType, NetworkEdge logger = logging.getLogger(__name__) @@ -87,7 +88,7 @@ _MAX_UNRESOLVED = 2000 -class RelationshipLinker(IdentifierIndex): +class RelationshipLinker(RuleLinksMixin, IdentifierIndex): """Builds interconnection edges between already-collected assets. Args: @@ -219,182 +220,10 @@ def _link_declared(self, asset: CloudAsset, edges: list[NetworkEdge]) -> None: ) # ------------------------------------------------------------------ - # Pass 2: provider-aware rules + # Pass 2: provider-aware rules (_link_rules, from RuleLinksMixin) and + # Identity Center permission sets # ------------------------------------------------------------------ - def _link_rules(self, asset: CloudAsset, edges: list[NetworkEdge]) -> None: - md = asset.metadata - raw = asset.raw_data or {} - - def res(ref: Any) -> str | None: - return self.resolve(ref, asset) - - # Security-group / NSG attachment (EC2, ELB, ENI, Lambda, RDS, Azure NIC) - for sg_ref in md.get("security_groups", []) or []: - self._add( - edges, - asset.id, - res(sg_ref), - EdgeType.ATTACHED_TO, - f"{asset.name} uses security group {sg_ref}", - "PROTECTED_BY_SG", - ) - nsg_id = md.get("nsg_id") - if nsg_id: - self._add( - edges, asset.id, res(nsg_id), EdgeType.ATTACHED_TO, relationship="PROTECTED_BY_SG" - ) - - # Subnet containment (everything with a subnet_id that isn't a subnet) - subnet_ref = md.get("subnet_id") - if subnet_ref and asset.asset_type != AssetType.SUBNET: - self._add( - edges, - res(subnet_ref), - asset.id, - EdgeType.CONTAINS, - f"Subnet {subnet_ref} contains {asset.name}", - "SUBNET_CONTAINS_INSTANCE", - ) - - # Lambda VPC config + execution role - vpc_config = md.get("vpc_config") or {} - if isinstance(vpc_config, dict): - for sn in vpc_config.get("SubnetIds", []) or []: - self._add( - edges, - res(sn), - asset.id, - EdgeType.CONTAINS, - relationship="SUBNET_CONTAINS_INSTANCE", - ) - for sg in vpc_config.get("SecurityGroupIds", []) or []: - self._add( - edges, asset.id, res(sg), EdgeType.ATTACHED_TO, relationship="PROTECTED_BY_SG" - ) - role_ref = raw.get("Role") or md.get("role_arn") - if role_ref and asset.asset_type in ( - AssetType.LAMBDA_FUNCTION, - AssetType.STATE_MACHINE, - ): - self._add( - edges, - asset.id, - res(role_ref), - EdgeType.ASSUMES_ROLE, - f"{asset.name} executes as {role_ref}", - "RUNS_ON", - ) - elif role_ref: - self._add( - edges, - asset.id, - res(role_ref), - EdgeType.REFERENCES, - f"{asset.name} references {role_ref}", - ) - - # EC2 instance profile - profile_arn = (raw.get("IamInstanceProfile") or {}).get("Arn") if raw else None - if profile_arn: - self._add( - edges, - asset.id, - res(profile_arn), - EdgeType.ASSUMES_ROLE, - "instance profile", - "RUNS_ON", - ) - - # KMS key references (secrets, volumes, tables, ...) - kms_ref = md.get("kms_key_id") - if kms_ref: - self._add( - edges, - asset.id, - res(kms_ref), - EdgeType.REFERENCES, - f"{asset.name} encrypted with {kms_ref}", - "ENCRYPTED_BY_KMS", - ) - - # CloudFront origins -> S3 buckets / load balancers, WAF -> distribution - for origin in md.get("origins", []) or []: - target = res(origin) - if target is None and isinstance(origin, str): - # bucket origins look like .s3..amazonaws.com - target = res(origin.split(".s3", 1)[0]) if ".s3" in origin else None - if target and ( - (asset.id, target) in self._pair_keys or (target, asset.id) in self._pair_keys - ): - continue # already linked by a declared (typed) relation - self._add( - edges, - asset.id, - target, - EdgeType.REFERENCES, - f"{asset.name} origin {origin}", - "SERVES_TRAFFIC_TO", - ) - web_acl = md.get("web_acl_id") - if web_acl: - self._add( - edges, res(web_acl), asset.id, EdgeType.PROTECTS, "WAF web ACL", "PROTECTED_BY_WAF" - ) - - # Route tables: routes -> gateways, associations -> subnets - if asset.asset_type == AssetType.ROUTE_TABLE: - for route in md.get("routes", []) or []: - gw = ( - route.get("GatewayId") - or route.get("NatGatewayId") - or route.get("TransitGatewayId") - or route.get("VpcPeeringConnectionId") - or route.get("NetworkInterfaceId") - ) - if gw and gw != "local": - self._add( - edges, - asset.id, - res(gw), - EdgeType.ROUTE, - f"route {route.get('DestinationCidrBlock', '')} via {gw}", - "NAT_TRANSLATED" if route.get("NatGatewayId") else "TRANSIT_ROUTED", - ) - for assoc in md.get("associations", []) or []: - self._add(edges, asset.id, res(assoc.get("SubnetId")), EdgeType.ATTACHED_TO) - - # Internet gateway attachments -> VPC - for attachment in md.get("attachments", []) or []: - if isinstance(attachment, dict) and attachment.get("VpcId"): - self._add(edges, asset.id, res(attachment.get("VpcId")), EdgeType.ATTACHED_TO) - - # ENI / EBS / EIP attachment -> instance - inst_ref = md.get("attached_instance_id") - if inst_ref: - self._add(edges, asset.id, res(inst_ref), EdgeType.ATTACHED_TO) - for inst in md.get("attached_instance_ids", []) or []: - self._add(edges, asset.id, res(inst), EdgeType.ATTACHED_TO) - eni_ref = md.get("network_interface_id") - if eni_ref and asset.asset_type != AssetType.NETWORK_INTERFACE: - self._add(edges, asset.id, res(eni_ref), EdgeType.ATTACHED_TO) - - # VPC peering - if asset.asset_type == AssetType.PEERING_CONNECTION: - req = res(md.get("requester_vpc_id")) - acc = res(md.get("accepter_vpc_id")) - self._add(edges, req, asset.id, EdgeType.PEERING, relationship="VPC_PEERED") - self._add(edges, asset.id, acc, EdgeType.PEERING, relationship="VPC_PEERED") - - # Azure VM -> NICs - for nic_ref in md.get("network_interfaces", []) or []: - self._add(edges, res(nic_ref), asset.id, EdgeType.ATTACHED_TO) - - # GCP: parent resource containment - parent = md.get("parent_full_resource_name") - if parent: - self._add(edges, res(parent), asset.id, EdgeType.CONTAINS) - def _link_sso_roles(self, edges: list[NetworkEdge]) -> None: """Identity Center permission set -> the AWSReservedSSO__ role it provisions in every assigned account.""" @@ -447,15 +276,10 @@ def _link_generic(self, asset: CloudAsset, edges: list[NetworkEdge]) -> None: target = self.resolve(value, asset) if not target or target == asset.id: continue - if (asset.id, target) in self._pair_keys or (target, asset.id) in self._pair_keys: + if self._linked_either_way(asset.id, target): continue - self._add( - edges, - asset.id, - target, - EdgeType.REFERENCES, - f"{asset.name} references {value}", - ) + description = f"{asset.name} references {value}" + self._add(edges, asset.id, target, EdgeType.REFERENCES, description) # ------------------------------------------------------------------ # Main interface diff --git a/cloudg/inventory/linker_rules.py b/cloudg/inventory/linker_rules.py new file mode 100644 index 0000000..69e4f85 --- /dev/null +++ b/cloudg/inventory/linker_rules.py @@ -0,0 +1,252 @@ +"""Provider-aware linking rules (pass 2 of :class:`~cloudg.inventory.linker.RelationshipLinker`). + +Each rule reads one family of metadata keys of an asset and emits the +typed edges it implies: security-group attachment, subnet containment, +Lambda VPC config and execution roles, instance profiles, KMS +references, CloudFront origins and WAF, route tables, gateway / ENI / +volume attachment, VPC peering, Azure NIC wiring and GCP parent +containment. The rules run in a fixed order because later ones (the +CloudFront origins) skip pairs that earlier ones already linked. +""" + +from __future__ import annotations + +from typing import Any, Callable + +from cloudg.schema.models import AssetType, CloudAsset, EdgeType, NetworkEdge + +__all__ = ["RuleLinksMixin"] + +_ROUTE_TARGET_KEYS = ( + "GatewayId", + "NatGatewayId", + "TransitGatewayId", + "VpcPeeringConnectionId", + "NetworkInterfaceId", +) +_EXECUTION_ROLE_TYPES = (AssetType.LAMBDA_FUNCTION, AssetType.STATE_MACHINE) + + +class RuleLinksMixin: + """The pass-2 rules; mixed into RelationshipLinker, which provides + ``resolve``, ``_add`` and ``_pair_keys``.""" + + _pair_keys: set[tuple[str, str]] + resolve: Callable[..., str | None] + _add: Callable[..., None] + + def _link_rules(self, asset: CloudAsset, edges: list[NetworkEdge]) -> None: + for rule in ( + self._rule_security_groups, + self._rule_subnet, + self._rule_lambda_vpc, + self._rule_roles, + self._rule_kms, + self._rule_origins, + self._rule_route_table, + self._rule_attachments, + self._rule_peering, + self._rule_parents, + ): + rule(asset, edges) + + def _res(self, ref: Any, asset: CloudAsset) -> str | None: + return self.resolve(ref, asset) + + def _linked_either_way(self, a: str, b: str) -> bool: + """True when an edge already joins ``a`` and ``b`` in either direction.""" + return (a, b) in self._pair_keys or (b, a) in self._pair_keys + + def _rule_security_groups(self, asset: CloudAsset, edges: list[NetworkEdge]) -> None: + """Security-group / NSG attachment (EC2, ELB, ENI, Lambda, RDS, Azure NIC).""" + md = asset.metadata + for sg_ref in md.get("security_groups", []) or []: + self._add( + edges, + asset.id, + self._res(sg_ref, asset), + EdgeType.ATTACHED_TO, + f"{asset.name} uses security group {sg_ref}", + "PROTECTED_BY_SG", + ) + nsg_id = md.get("nsg_id") + if nsg_id: + self._add( + edges, + asset.id, + self._res(nsg_id, asset), + EdgeType.ATTACHED_TO, + relationship="PROTECTED_BY_SG", + ) + + def _rule_subnet(self, asset: CloudAsset, edges: list[NetworkEdge]) -> None: + """Subnet containment (everything with a subnet_id that isn't a subnet).""" + subnet_ref = asset.metadata.get("subnet_id") + if subnet_ref and asset.asset_type != AssetType.SUBNET: + self._add( + edges, + self._res(subnet_ref, asset), + asset.id, + EdgeType.CONTAINS, + f"Subnet {subnet_ref} contains {asset.name}", + "SUBNET_CONTAINS_INSTANCE", + ) + + def _rule_lambda_vpc(self, asset: CloudAsset, edges: list[NetworkEdge]) -> None: + """Lambda VPC config: subnets contain the function, its groups protect it.""" + vpc_config = asset.metadata.get("vpc_config") or {} + if not isinstance(vpc_config, dict): + return + for sn in vpc_config.get("SubnetIds", []) or []: + self._add( + edges, + self._res(sn, asset), + asset.id, + EdgeType.CONTAINS, + relationship="SUBNET_CONTAINS_INSTANCE", + ) + for sg in vpc_config.get("SecurityGroupIds", []) or []: + self._add( + edges, + asset.id, + self._res(sg, asset), + EdgeType.ATTACHED_TO, + relationship="PROTECTED_BY_SG", + ) + + def _rule_roles(self, asset: CloudAsset, edges: list[NetworkEdge]) -> None: + """Execution roles (Lambda, Step Functions), other role references and + EC2 instance profiles.""" + raw = asset.raw_data or {} + role_ref = raw.get("Role") or asset.metadata.get("role_arn") + if role_ref and asset.asset_type in _EXECUTION_ROLE_TYPES: + self._add( + edges, + asset.id, + self._res(role_ref, asset), + EdgeType.ASSUMES_ROLE, + f"{asset.name} executes as {role_ref}", + "RUNS_ON", + ) + elif role_ref: + self._add( + edges, + asset.id, + self._res(role_ref, asset), + EdgeType.REFERENCES, + f"{asset.name} references {role_ref}", + ) + profile_arn = (raw.get("IamInstanceProfile") or {}).get("Arn") if raw else None + if profile_arn: + self._add( + edges, + asset.id, + self._res(profile_arn, asset), + EdgeType.ASSUMES_ROLE, + "instance profile", + "RUNS_ON", + ) + + def _rule_kms(self, asset: CloudAsset, edges: list[NetworkEdge]) -> None: + """KMS key references (secrets, volumes, tables, ...).""" + kms_ref = asset.metadata.get("kms_key_id") + if kms_ref: + self._add( + edges, + asset.id, + self._res(kms_ref, asset), + EdgeType.REFERENCES, + f"{asset.name} encrypted with {kms_ref}", + "ENCRYPTED_BY_KMS", + ) + + def _origin_target(self, origin: Any, asset: CloudAsset) -> str | None: + target = self._res(origin, asset) + if target is None and isinstance(origin, str) and ".s3" in origin: + # bucket origins look like .s3..amazonaws.com + target = self._res(origin.split(".s3", 1)[0], asset) + return target + + def _rule_origins(self, asset: CloudAsset, edges: list[NetworkEdge]) -> None: + """CloudFront origins -> S3 buckets / load balancers, WAF -> distribution.""" + md = asset.metadata + for origin in md.get("origins", []) or []: + target = self._origin_target(origin, asset) + if target and self._linked_either_way(asset.id, target): + continue # already linked by a declared (typed) relation + self._add( + edges, + asset.id, + target, + EdgeType.REFERENCES, + f"{asset.name} origin {origin}", + "SERVES_TRAFFIC_TO", + ) + web_acl = md.get("web_acl_id") + if web_acl: + self._add( + edges, + self._res(web_acl, asset), + asset.id, + EdgeType.PROTECTS, + "WAF web ACL", + "PROTECTED_BY_WAF", + ) + + def _rule_route_table(self, asset: CloudAsset, edges: list[NetworkEdge]) -> None: + """Route tables: routes -> gateways, associations -> subnets.""" + if asset.asset_type != AssetType.ROUTE_TABLE: + return + md = asset.metadata + for route in md.get("routes", []) or []: + gw = next((route.get(k) for k in _ROUTE_TARGET_KEYS if route.get(k)), None) + if gw and gw != "local": + self._add( + edges, + asset.id, + self._res(gw, asset), + EdgeType.ROUTE, + f"route {route.get('DestinationCidrBlock', '')} via {gw}", + "NAT_TRANSLATED" if route.get("NatGatewayId") else "TRANSIT_ROUTED", + ) + for assoc in md.get("associations", []) or []: + self._add( + edges, asset.id, self._res(assoc.get("SubnetId"), asset), EdgeType.ATTACHED_TO + ) + + def _rule_attachments(self, asset: CloudAsset, edges: list[NetworkEdge]) -> None: + """Internet gateway -> VPC, and ENI / EBS / EIP -> instance attachment.""" + md = asset.metadata + targets: list[Any] = [ + attachment.get("VpcId") + for attachment in md.get("attachments", []) or [] + if isinstance(attachment, dict) and attachment.get("VpcId") + ] + inst_ref = md.get("attached_instance_id") + if inst_ref: + targets.append(inst_ref) + targets.extend(md.get("attached_instance_ids", []) or []) + eni_ref = md.get("network_interface_id") + if eni_ref and asset.asset_type != AssetType.NETWORK_INTERFACE: + targets.append(eni_ref) + for ref in targets: + self._add(edges, asset.id, self._res(ref, asset), EdgeType.ATTACHED_TO) + + def _rule_peering(self, asset: CloudAsset, edges: list[NetworkEdge]) -> None: + """VPC peering: requester VPC -> peering -> accepter VPC.""" + if asset.asset_type != AssetType.PEERING_CONNECTION: + return + md = asset.metadata + req = self._res(md.get("requester_vpc_id"), asset) + acc = self._res(md.get("accepter_vpc_id"), asset) + self._add(edges, req, asset.id, EdgeType.PEERING, relationship="VPC_PEERED") + self._add(edges, asset.id, acc, EdgeType.PEERING, relationship="VPC_PEERED") + + def _rule_parents(self, asset: CloudAsset, edges: list[NetworkEdge]) -> None: + """Azure VM -> NICs, and GCP parent resource containment.""" + md = asset.metadata + for nic_ref in md.get("network_interfaces", []) or []: + self._add(edges, self._res(nic_ref, asset), asset.id, EdgeType.ATTACHED_TO) + parent = md.get("parent_full_resource_name") + if parent: + self._add(edges, self._res(parent, asset), asset.id, EdgeType.CONTAINS) diff --git a/cloudg/inventory/mapper.py b/cloudg/inventory/mapper.py index c7a00cb..f798651 100644 --- a/cloudg/inventory/mapper.py +++ b/cloudg/inventory/mapper.py @@ -1,4 +1,4 @@ -"""Inventory mapper — scanner-independent infrastructure mapping. +"""Inventory mapper: scanner-independent infrastructure mapping. Orchestrates deep collection across providers, accounts and regions, links every asset into an interconnected map, and produces exportable artifacts. @@ -24,7 +24,7 @@ Scanner findings produced elsewhere (a `cloudg run`, `cloudg ingest`, or any `CloudGEngine.scan()`) can be merged in afterwards to overlay the -inventory with risk — producing an asset map (asset → findings) and a +inventory with risk, producing an asset map (asset → findings) and a compliance map (framework → affected assets). Usage (programmatic): @@ -76,6 +76,55 @@ logger = logging.getLogger(__name__) +def _merge_copy(current: CloudAsset, duplicate: CloudAsset) -> None: + """Fold the declared relations, aliases and exposure of ``duplicate`` into ``current``.""" + for list_key in ("relations", "aliases"): + extra = duplicate.metadata.get(list_key) or [] + if extra: + merged = list(current.metadata.get(list_key) or []) + merged.extend(x for x in extra if x not in merged) + current.metadata[list_key] = merged + current.is_internet_exposed = current.is_internet_exposed or duplicate.is_internet_exposed + + +def _edge_identity(edge: NetworkEdge, source: str, target: str) -> tuple[Any, ...]: + """What makes two edges between the same endpoints the same edge. + + Besides the endpoints and the type, the rule attributes count: two + SECURITY_GROUP_RULE / NACL_RULE / INTERNET_EXPOSED edges from one CIDR + to one group for different ports or protocols are parallel rules, not + duplicates. + """ + return ( + source, + target, + edge.edge_type.value, + edge.port_range, + tuple(sorted(edge.ports)), + edge.protocol, + edge.cidr, + edge.direction, + ) + + +def _remap_edges(edges: list[NetworkEdge], remap: dict[str, str]) -> list[NetworkEdge]: + """Re-point edges at the kept copies and drop the edges that became + self-loops or exact duplicates.""" + new_edges = [] + seen: set[tuple[Any, ...]] = set() + for e in edges: + s = remap.get(e.source_id, e.source_id) + t = remap.get(e.target_id, e.target_id) + identity = _edge_identity(e, s, t) + if s == t or identity in seen: + continue + seen.add(identity) + if s != e.source_id or t != e.target_id: + e = e.model_copy(update={"source_id": s, "target_id": t}) + new_edges.append(e) + return new_edges + + def deduplicate( assets: list[CloudAsset], edges: list[NetworkEdge] ) -> tuple[list[CloudAsset], list[NetworkEdge]]: @@ -88,13 +137,11 @@ def deduplicate( out: list[CloudAsset] = [] for asset in assets: key = asset.arn - if not key: - out.append(asset) - continue - current = kept.get(key) + current = kept.get(key) if key else None if current is None: - kept[key] = asset - position[key] = len(out) + if key: + kept[key] = asset + position[key] = len(out) out.append(asset) continue # Prefer the detailed copy over sweep / placeholder discoveries @@ -102,31 +149,13 @@ def deduplicate( asset, current = current, asset kept[key] = current out[position[key]] = current - for list_key in ("relations", "aliases"): - extra = asset.metadata.get(list_key) or [] - if extra: - merged = list(current.metadata.get(list_key) or []) - merged.extend(x for x in extra if x not in merged) - current.metadata[list_key] = merged - current.is_internet_exposed = current.is_internet_exposed or asset.is_internet_exposed + _merge_copy(current, asset) remap[asset.id] = current.id if not remap: return out, edges - new_edges = [] - seen: set[tuple[str, str, str]] = set() - for e in edges: - s = remap.get(e.source_id, e.source_id) - t = remap.get(e.target_id, e.target_id) - key3 = (s, t, e.edge_type.value) - if s == t or key3 in seen: - continue - seen.add(key3) - if s != e.source_id or t != e.target_id: - e = e.model_copy(update={"source_id": s, "target_id": t}) - new_edges.append(e) logger.info("Deduplicated %d repeated assets", len(remap)) - return out, new_edges + return out, _remap_edges(edges, remap) def _account_identifier(provider: CloudProvider, account_id: str) -> str: @@ -137,22 +166,28 @@ def _account_identifier(provider: CloudProvider, account_id: str) -> str: return f"arn:aws:iam::{account_id}:root" -def add_account_hierarchy( - assets: list[CloudAsset], edges: list[NetworkEdge] -) -> tuple[list[CloudAsset], list[NetworkEdge]]: - """Ensure an account node per account and make it contain every - top-level resource (anything not already contained by something).""" +def _needs_account(asset: CloudAsset) -> bool: + return bool( + asset.account_id + and asset.account_id != "unknown" + and asset.asset_type not in _HIERARCHY_TYPES + ) + + +def _account_nodes( + assets: list[CloudAsset], +) -> tuple[dict[tuple[CloudProvider, str], CloudAsset], list[CloudAsset]]: + """The account node of every account (existing or new), and the new ones.""" by_arn = {a.arn: a for a in assets if a.arn} - accounts: dict[tuple[CloudProvider, str], CloudAsset] = {} - for a in assets: - if a.asset_type == AssetType.CLOUD_ACCOUNT and a.account_id: - accounts[(a.provider, a.account_id)] = a + accounts: dict[tuple[CloudProvider, str], CloudAsset] = { + (a.provider, a.account_id): a + for a in assets + if a.asset_type == AssetType.CLOUD_ACCOUNT and a.account_id + } new_assets: list[CloudAsset] = [] for a in assets: - if not a.account_id or a.account_id == "unknown" or a.asset_type in _HIERARCHY_TYPES: - continue key = (a.provider, a.account_id) - if key in accounts: + if not _needs_account(a) or key in accounts: continue ident = _account_identifier(a.provider, a.account_id) existing = by_arn.get(ident) @@ -170,7 +205,15 @@ def add_account_hierarchy( ) accounts[key] = node new_assets.append(node) + return accounts, new_assets + +def add_account_hierarchy( + assets: list[CloudAsset], edges: list[NetworkEdge] +) -> tuple[list[CloudAsset], list[NetworkEdge]]: + """Ensure an account node per account and make it contain every + top-level resource (anything not already contained by something).""" + accounts, new_assets = _account_nodes(assets) contained = {e.target_id for e in edges if e.edge_type == EdgeType.CONTAINS} new_edges = list(edges) for a in assets: @@ -437,7 +480,27 @@ def _link( # ------------------------------------------------------------------ async def map_inventory(self) -> InventoryResult: - """Deep-collect every provider and link the assets into a map.""" + """Deep-collect every provider and link the assets into a map. + + Cloud API throttling never fails the map: calls are rate limited and + retried (``config.ratelimit``), a service still throttled after its + retries is recorded FAILED / PARTIAL with reason "throttled", and the + run's throttling telemetry is attached as ``result.throttling``. + """ + from cloudg.resilience import configure, current_stats, get_governor, stats_scope + + configure(self._config) + # Reuse a run already bound by the caller (e.g. an MCP live tool) so + # its stats see this map's throttling too + with stats_scope(current_stats()) as stats: + result = await self._map_inventory() + if stats.eventful: + result.throttling = get_governor().summary(stats) + for line in result.throttling.get("messages", []): + logger.warning("Throttling: %s", line) + return result + + async def _map_inventory(self) -> InventoryResult: start = time.time() cfg = self._config.model_copy(deep=True) org_coverage: list[CollectionCoverage] = [] diff --git a/cloudg/inventory/mapper_result.py b/cloudg/inventory/mapper_result.py index 0fa78c2..ebab013 100644 --- a/cloudg/inventory/mapper_result.py +++ b/cloudg/inventory/mapper_result.py @@ -94,9 +94,24 @@ def _write_json(path: Path, data: Any) -> Path: return path +def _asset_flags(assets: list[CloudAsset], linked_ids: set[str]) -> dict[str, int]: + """Counts of unlinked, internet-exposed, external-account and security-gap assets.""" + flags = {"unlinked": 0, "exposed": 0, "external": 0, "gaps": 0} + for a in assets: + if a.id not in linked_ids and a.asset_type not in _HIERARCHY_TYPES: + flags["unlinked"] += 1 + if a.is_internet_exposed: + flags["exposed"] += 1 + if a.asset_type == AssetType.CLOUD_ACCOUNT and a.metadata.get("external"): + flags["external"] += 1 + if a.metadata.get("security_service") and a.metadata.get("enabled") is False: + flags["gaps"] += 1 + return flags + + @dataclass class InventoryResult: - """Complete inventory map — assets, interconnections, and summary.""" + """Complete inventory map: assets, interconnections, and summary.""" assets: list[CloudAsset] = field(default_factory=list) edges: list[NetworkEdge] = field(default_factory=list) @@ -106,22 +121,16 @@ class InventoryResult: duration_ms: int = 0 organization: dict[str, Any] | None = None unresolved_references: list[dict[str, Any]] = field(default_factory=list) + #: Throttling telemetry of the run (``cloudg.resilience`` summary: totals, + #: per-scope counters, human-readable messages, skipped scopes); None when + #: no API pushed back + throttling: dict[str, Any] | None = None @property def summary(self) -> dict[str, Any]: by_type, by_region, by_account, by_service = _asset_counts(self.assets) edge_types, relationships, linked_ids, cross_account = _edge_stats(self.assets, self.edges) - orphans = sum( - 1 - for a in self.assets - if a.id not in linked_ids and a.asset_type not in _HIERARCHY_TYPES - ) - security_gaps = sum( - 1 - for a in self.assets - if a.metadata.get("security_service") and a.metadata.get("enabled") is False - ) - + flags = _asset_flags(self.assets, linked_ids) out = { "total_assets": len(self.assets), "total_edges": len(self.edges), @@ -132,20 +141,22 @@ def summary(self) -> dict[str, Any]: "assets_by_account": by_account, "edges_by_type": edge_types, "edges_by_relationship": _by_count(relationships), - "unlinked_assets": orphans, - "internet_exposed": sum(1 for a in self.assets if a.is_internet_exposed), + "unlinked_assets": flags["unlinked"], + "internet_exposed": flags["exposed"], "accounts": len([k for k in by_account if k != "unknown"]), "cross_account_edges": cross_account, - "external_accounts": sum( - 1 - for a in self.assets - if a.asset_type == AssetType.CLOUD_ACCOUNT and a.metadata.get("external") - ), - "security_service_gaps": security_gaps, + "external_accounts": flags["external"], + "security_service_gaps": flags["gaps"], "unresolved_references": len(self.unresolved_references), } if self.organization: out["organization"] = _organization_summary(self.organization) + if self.throttling: + out["throttling"] = { + "totals": self.throttling.get("totals", {}), + "messages": self.throttling.get("messages", []), + "skipped": self.throttling.get("skipped", {}), + } return out # ------------------------------------------------------------------ @@ -200,6 +211,7 @@ def export(self, output_dir: str | Path) -> dict[str, Path]: "assets": [a.model_dump(mode="json") for a in self.assets], "edges": [e.model_dump(mode="json") for e in self.edges], "unresolved_references": self.unresolved_references, + **({"throttling": self.throttling} if self.throttling else {}), }, ) @@ -242,6 +254,7 @@ def load(cls, path: str | Path) -> "InventoryResult": regions=data.get("regions", {}), organization=org, unresolved_references=data.get("unresolved_references", []), + throttling=data.get("throttling"), ) diff --git a/cloudg/mcp/__init__.py b/cloudg/mcp/__init__.py new file mode 100644 index 0000000..3f176e0 --- /dev/null +++ b/cloudg/mcp/__init__.py @@ -0,0 +1,138 @@ +"""cloudg's MCP layer: expose cloudg's inventory, graph, findings, +compliance and ontology to AI agents over the Model Context Protocol. + +Three ways to use it: + +1. Run it as a server:: + + cloudg mcp serve # stdio, for desktop clients + cloudg mcp serve --transport http # streamable HTTP on 127.0.0.1:8765 + +2. Mount every tool, resource and prompt into an MCP server you already + have (official ``mcp`` SDK 1.x / 2.x, low-level ``Server``, or the + standalone ``fastmcp`` package):: + + from cloudg.mcp import CloudGMCPLayer + layer = CloudGMCPLayer(policy="strict", prefix="cloudg_") + layer.register_into(my_server) + +3. Call it in-process:: + + result = await layer.call_tool("find_assets", {"internet_exposed": True}) + +Every request, whichever way it arrives, passes through the same policy: +access control, rate limits, redaction, pseudonymisation, projection and +annotation. See docs/MCP.md. + +Importing this package pulls in nothing beyond cloudg's own dependencies; +the SDK adapters import ``mcp`` / ``fastmcp`` only when used. +""" + +from __future__ import annotations + +from cloudg.mcp.context import Principal, ToolContext +from cloudg.mcp.core import ( + AccessDeniedError, + Capability, + ContentAnnotations, + EmbeddedResource, + Icon, + ImageContent, + InvalidArgumentsError, + MCPLayerError, + NotFoundError, + PromptArgument, + PromptMessage, + PromptResult, + PromptSpec, + RateLimitedError, + Registry, + ResourceLink, + ResourceSpec, + ResourceTemplateSpec, + Sensitivity, + TextContent, + TextResourceContents, + ToolAnnotations, + ToolResult, + ToolSpec, +) +from cloudg.mcp.layer import CallInfo, CloudGMCPLayer +from cloudg.mcp.policy import Policy, available_profiles +from cloudg.mcp.state import Dataset, Workspace +from cloudg.mcp.transforms import Pipeline, TokenVault, TransformContext, build_transform + + +def default_registry() -> Registry: + """The full cloudg catalog (tools, resources, templates, prompts).""" + from cloudg.mcp.catalog import default_registry as _default + + return _default() + + +def register_into(layer: CloudGMCPLayer, server: object, **kwargs: object) -> object: + """Mount ``layer`` into an existing MCP server object; see + :func:`cloudg.mcp.adapters.register_into`.""" + from cloudg.mcp.adapters import register_into as _register + + return _register(layer, server, **kwargs) + + +def build_server(layer: CloudGMCPLayer, flavor: str = "auto") -> object: + """A ready-to-run server object for ``layer``; see + :func:`cloudg.mcp.adapters.build_server`.""" + from cloudg.mcp.adapters import build_server as _build + + return _build(layer, flavor=flavor) + + +def export_definitions(layer: CloudGMCPLayer, principal: Principal | None = None) -> object: + """Framework-free wire definitions plus async handlers; see + :func:`cloudg.mcp.adapters.export_definitions`.""" + from cloudg.mcp.adapters import export_definitions as _export + + return _export(layer, principal=principal) + + +__all__ = [ + "AccessDeniedError", + "CallInfo", + "Capability", + "CloudGMCPLayer", + "ContentAnnotations", + "Dataset", + "EmbeddedResource", + "Icon", + "ImageContent", + "InvalidArgumentsError", + "MCPLayerError", + "NotFoundError", + "Pipeline", + "Policy", + "Principal", + "PromptArgument", + "PromptMessage", + "PromptResult", + "PromptSpec", + "RateLimitedError", + "Registry", + "ResourceLink", + "ResourceSpec", + "ResourceTemplateSpec", + "Sensitivity", + "TextContent", + "TextResourceContents", + "TokenVault", + "ToolAnnotations", + "ToolContext", + "ToolResult", + "ToolSpec", + "TransformContext", + "Workspace", + "available_profiles", + "build_server", + "build_transform", + "default_registry", + "export_definitions", + "register_into", +] diff --git a/cloudg/mcp/__main__.py b/cloudg/mcp/__main__.py new file mode 100644 index 0000000..d3c6f0b --- /dev/null +++ b/cloudg/mcp/__main__.py @@ -0,0 +1,10 @@ +"""Entry point for ``python -m cloudg.mcp``, the cloudg MCP command line. + +The commands live in :mod:`cloudg.mcp.cli`. ``python -m cloudg.mcp serve`` +starts a stdio MCP server; ``--help`` lists the other commands. +""" + +from cloudg.mcp.cli import main + +if __name__ == "__main__": + main() diff --git a/cloudg/mcp/adapters/__init__.py b/cloudg/mcp/adapters/__init__.py new file mode 100644 index 0000000..9544ed9 --- /dev/null +++ b/cloudg/mcp/adapters/__init__.py @@ -0,0 +1,232 @@ +"""Mount the cloudg MCP layer into existing MCP servers. + +One call adds every cloudg tool, resource, resource template, prompt and +completion to a server you already run:: + + from mcp.server.mcpserver import MCPServer # or FastMCP, Server, fastmcp.FastMCP + from cloudg.mcp.layer import CloudGMCPLayer + from cloudg.mcp.adapters import register_into + + server = MCPServer("my-platform") + layer = CloudGMCPLayer(policy="standard") + register_into(layer, server, prefix="cloudg_") + +Supported servers (detected automatically): + +===================================== ========================================== +``mcp.server.mcpserver.MCPServer`` mcp 2.x (:mod:`.mcp_sdk` -> :mod:`.lowlevel`) +``mcp.server.fastmcp.FastMCP`` mcp 1.x (:mod:`.mcp_sdk` -> :mod:`.lowlevel`) +``mcp.server.lowlevel.Server`` mcp 1.x and 2.x (:mod:`.lowlevel`) +``fastmcp.FastMCP`` standalone fastmcp 2.x / 3.x / 4.x (:mod:`.fastmcp`) +===================================== ========================================== + +Guarantees, whichever server: tools keep the layer's exact ``inputSchema``, +``outputSchema``, ``annotations``, ``title`` and ``_meta``; +``structuredContent`` and ``isError`` pass through; resource templates and +prompts get cloudg's completions; progress and log notifications reach the +client; the caller's :class:`~cloudg.mcp.context.Principal` is resolved per +request (pluggable ``principal_resolver``); ``layer.notify_change`` becomes +list-changed / resource-updated notifications; and every request goes +through ``layer.call_tool`` / ``read_resource`` / ``get_prompt`` / +``complete`` so policies, transforms and middleware apply. + +Arguments of :func:`register_into`: + +* ``layer``: the cloudg layer. +* ``server``: an MCP server object from the table above. +* ``prefix``: name prefix for tools and prompts, e.g. ``"cloudg_"`` to + avoid collisions with the host's tools. It sets ``layer.prefix``: the + prefix belongs to the layer, so it also applies to every other server + this layer is (or was) mounted on. To expose the catalog under two + prefixes, mount two layers. +* ``include``: subset of ``{"tools", "resources", "templates", "prompts", + "completions", "logging", "subscriptions"}`` to register (default: all). +* ``principal_resolver``: ``fn(RequestInfo) -> Principal | None`` mapping + request facts (HTTP headers, OAuth access token, transport) to the + caller; defaults to + :func:`~cloudg.mcp.native.auth.default_principal_resolver`. +* ``list_changed``: advertise ``listChanged`` on SDK servers so clients act + on change notifications. + +Without any MCP framework, :func:`export_definitions` gives the wire +definitions plus async handlers, and :func:`build_server` creates a ready +server object (SDK, fastmcp or the dependency-free native server). +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any, Iterable + +from cloudg.mcp.adapters.export import ( + ExportedDefinitions, + export_definitions, + to_anthropic_tools, + to_langchain_tools, + to_openai_tools, +) +from cloudg.mcp.native.auth import PrincipalResolver, RequestInfo + +if TYPE_CHECKING: # pragma: no cover + from cloudg.mcp.layer import CloudGMCPLayer + +__all__ = [ + "ExportedDefinitions", + "RequestInfo", + "build_server", + "detect_server_kind", + "export_definitions", + "register_into", + "to_anthropic_tools", + "to_langchain_tools", + "to_openai_tools", +] + +logger = logging.getLogger("cloudg.mcp.adapters") + +BUILD_FLAVORS = ("auto", "sdk", "lowlevel", "mcpserver", "fastmcp", "native") + + +def detect_server_kind(server: Any) -> str: + """Classify a server object: ``fastmcp`` (standalone package), + ``sdk-highlevel`` (``MCPServer`` / SDK ``FastMCP``), ``lowlevel`` + (SDK ``Server``), ``native`` (:class:`~cloudg.mcp.native.NativeMCPServer`).""" + from cloudg.mcp.adapters.fastmcp import is_fastmcp_server + from cloudg.mcp.adapters.lowlevel import is_lowlevel_server + from cloudg.mcp.adapters.mcp_sdk import is_sdk_highlevel_server + from cloudg.mcp.native.protocol import NativeMCPServer + + if isinstance(server, NativeMCPServer): + return "native" + if is_fastmcp_server(server): + return "fastmcp" + if is_sdk_highlevel_server(server): + return "sdk-highlevel" + if is_lowlevel_server(server): + return "lowlevel" + # duck typing for subclasses / wrappers + if hasattr(server, "add_request_handler") or hasattr(server, "request_handlers"): + return "lowlevel" + if hasattr(server, "_lowlevel_server") or hasattr(server, "_mcp_server"): + return "sdk-highlevel" + raise TypeError( + f"Don't know how to register cloudg into {type(server).__module__}." + f"{type(server).__name__}; supported: mcp MCPServer / FastMCP / lowlevel Server, " + "fastmcp.FastMCP. Use export_definitions(layer) for custom servers." + ) + + +def _apply_prefix(layer: "CloudGMCPLayer", prefix: str | None) -> None: + """Set ``layer.prefix`` for :func:`register_into`, warning when that + replaces a different prefix already in use.""" + if prefix is None: + return + if layer.prefix and prefix != layer.prefix: + logger.warning( + "register_into(prefix=%r) replaces this layer's prefix %r on every server " + "it is mounted on; use one layer per prefix", + prefix, + layer.prefix, + ) + layer.set_prefix(prefix) + + +def _install(kind: str, layer: "CloudGMCPLayer", server: Any, **kwargs: Any) -> Any: + """Run the adapter for a non-native server ``kind``; ``kwargs`` are + ``include``, ``principal_resolver`` and ``list_changed``.""" + if kind == "fastmcp": + from cloudg.mcp.adapters.fastmcp import install + + kwargs.pop("list_changed", None) + return install(layer, server, **kwargs) + if kind == "sdk-highlevel": + from cloudg.mcp.adapters.mcp_sdk import install as install_sdk + + return install_sdk(layer, server, **kwargs) + from cloudg.mcp.adapters.lowlevel import install as install_lowlevel + + return install_lowlevel(layer, server, **kwargs) + + +def register_into( + layer: "CloudGMCPLayer", + server: Any, + *, + prefix: str | None = None, + include: Iterable[str] | None = None, + principal_resolver: PrincipalResolver | None = None, + list_changed: bool = True, +) -> Any: + """Register the layer's primitives into ``server`` and return the binding. + + The arguments are described under "Arguments of register_into" in the + module docstring. + """ + _apply_prefix(layer, prefix) + kind = detect_server_kind(server) + if kind == "native": + if server.layer is not layer: + raise ValueError("A NativeMCPServer serves exactly one layer (its own)") + return server + return _install( + kind, + layer, + server, + include=include, + principal_resolver=principal_resolver, + list_changed=list_changed, + ) + + +def build_server( + layer: "CloudGMCPLayer", + flavor: str = "auto", + *, + name: str | None = None, + principal_resolver: PrincipalResolver | None = None, + include: Iterable[str] | None = None, + **native_kwargs: Any, +) -> Any: + """Create a ready-to-run server object exposing ``layer``. + + Flavors: + ``auto``: SDK low-level ``Server`` if ``mcp`` is importable, else native. + ``sdk`` / ``lowlevel``: ``mcp.server.lowlevel.Server``. + ``mcpserver``: ``MCPServer`` (mcp 2.x) or ``FastMCP`` (mcp 1.x). + ``fastmcp``: standalone ``fastmcp.FastMCP``. + ``native``: :class:`~cloudg.mcp.native.NativeMCPServer` (stdlib only). + + The SDK / fastmcp objects carry the binding at ``server._cloudg_binding``. + """ + if flavor not in BUILD_FLAVORS: + raise ValueError(f"Unknown flavor {flavor!r}; choose from {BUILD_FLAVORS}") + if flavor == "auto": + from cloudg.mcp.native._shared import sdk_installed + + flavor = "lowlevel" if sdk_installed() else "native" + if flavor == "native": + from cloudg.mcp.native.protocol import NativeMCPServer + + return NativeMCPServer(layer, principal_resolver=principal_resolver, **native_kwargs) + server = _new_server(layer, flavor, name) + binding = register_into(layer, server, include=include, principal_resolver=principal_resolver) + try: + server._cloudg_binding = binding + except (AttributeError, TypeError): # pragma: no cover (slotted / frozen server classes) + logger.debug("cannot attach the cloudg binding to %s", type(server).__name__) + return server + + +def _new_server(layer: "CloudGMCPLayer", flavor: str, name: str | None) -> Any: + """A fresh, empty SDK or fastmcp server object for ``build_server``.""" + if flavor in ("sdk", "lowlevel"): + from cloudg.mcp.adapters.mcp_sdk import create_lowlevel_server + + return create_lowlevel_server(layer, name=name) + if flavor == "mcpserver": + from cloudg.mcp.adapters.mcp_sdk import create_highlevel_server + + return create_highlevel_server(layer, name=name) + import fastmcp + + return fastmcp.FastMCP(name or layer.name, instructions=layer.instructions or None) diff --git a/cloudg/mcp/adapters/_common.py b/cloudg/mcp/adapters/_common.py new file mode 100644 index 0000000..ad8f11b --- /dev/null +++ b/cloudg/mcp/adapters/_common.py @@ -0,0 +1,276 @@ +"""Helpers shared by the framework adapters. + +* ownership checks: does a tool name / prompt name / resource URI belong to + the cloudg layer, or to the host server's own primitives? +* :class:`ChangeFanout` forwards ``layer.on_change`` events to every live + client session the adapter has seen, from any thread. +* small utilities: SDK version detection, level normalisation, safe calls. +""" + +from __future__ import annotations + +import asyncio +import importlib.metadata +import logging +import weakref +from typing import TYPE_CHECKING, Any, Awaitable, Callable, Iterable + +from cloudg.mcp.context import Principal +from cloudg.mcp.core import MCPLayerError +from cloudg.mcp.native._shared import ( + expand_change, + layer_error_parts, + normalize_level, + run_on_loop, +) +from cloudg.mcp.native.auth import ( + PRINCIPAL_SCOPE_KEY, + PrincipalResolver, + RequestInfo, + default_principal_resolver, +) + +if TYPE_CHECKING: # pragma: no cover + from cloudg.mcp.layer import CloudGMCPLayer + +logger = logging.getLogger("cloudg.mcp.adapters") + +ALL_KINDS = frozenset( + {"tools", "resources", "templates", "prompts", "completions", "logging", "subscriptions"} +) + +__all__ = [ + "ALL_KINDS", + "ChangeFanout", + "call_layer", + "layer_error_parts", + "normalize_kinds", + "normalize_level", + "owns_prompt", + "owns_resource", + "owns_template", + "owns_tool", + "package_version", + "resolve_principal", + "scope_principal", + "send_change", +] + + +def package_version(name: str) -> str | None: + try: + return importlib.metadata.version(name) + except importlib.metadata.PackageNotFoundError: + return None + + +def major_version(name: str) -> int | None: + v = package_version(name) + if not v: + return None + try: + return int(v.split(".")[0]) + except ValueError: + return None + + +def normalize_kinds(include: Iterable[str] | None) -> frozenset[str]: + if include is None: + return ALL_KINDS + kinds = frozenset(include) + unknown = kinds - ALL_KINDS + if unknown: + raise ValueError(f"Unknown primitive kinds: {sorted(unknown)}; use {sorted(ALL_KINDS)}") + return kinds + + +def owns_tool(layer: "CloudGMCPLayer", name: str) -> bool: + if layer.prefix and not name.startswith(layer.prefix): + return False + return layer.internal_name(name) in layer.registry.tools + + +def owns_prompt(layer: "CloudGMCPLayer", name: str) -> bool: + if layer.prefix and not name.startswith(layer.prefix): + return False + return layer.internal_name(name) in layer.registry.prompts + + +def owns_resource(layer: "CloudGMCPLayer", uri: str) -> bool: + return uri in layer.registry.resources or layer.registry.find_template(uri) is not None + + +def owns_template(layer: "CloudGMCPLayer", uri_template: str) -> bool: + return uri_template in layer.registry.templates + + +def resolve_principal(resolver: PrincipalResolver | None, info: RequestInfo) -> Principal: + if resolver is not None: + try: + principal = resolver(info) + except Exception: + logger.exception("principal resolver failed; falling back to the default") + principal = None + if principal is not None: + return principal + return default_principal_resolver(info) + + +def lower_headers(headers: Any) -> dict[str, str]: + if not headers: + return {} + try: + items = headers.items() + except AttributeError: + return {} + return {str(k).lower(): str(v) for k, v in items} + + +def scope_principal(request: Any) -> Principal | None: + """The principal an auth middleware stored on an ASGI request + (``scope["cloudg.principal"]`` or ``scope["state"]["cloudg.principal"]``).""" + scope = getattr(request, "scope", None) + if not isinstance(scope, dict): + return None + principal = scope.get(PRINCIPAL_SCOPE_KEY) + state = scope.get("state") + if principal is None and isinstance(state, dict): + principal = state.get(PRINCIPAL_SCOPE_KEY) + return principal + + +async def call_layer( + call: Awaitable[Any], + on_layer_error: Callable[[MCPLayerError], Exception], + on_internal_error: Callable[[], Exception], +) -> Any: + """Await a layer call. Layer errors (already run through the caller's + output pipeline by the layer) map through ``on_layer_error``; anything + else is logged here and replaced by ``on_internal_error()``, so raw + exception text never reaches the client.""" + try: + return await call + except MCPLayerError as exc: + raise on_layer_error(exc) from None + except Exception: + logger.exception("cloudg MCP layer call failed") + raise on_internal_error() from None + + +async def send_change(session: Any, kind: str, uri: str | None) -> None: + """Send one change notification through an SDK ``ServerSession``.""" + if kind == "tools": + await session.send_tool_list_changed() + elif kind == "prompts": + await session.send_prompt_list_changed() + elif kind == "resources": + await session.send_resource_list_changed() + elif kind == "resource" and uri: + await session.send_resource_updated(uri) + + +Notifier = Callable[[Any, str, "str | None"], Awaitable[None]] + + +class ChangeFanout: + """Forward layer change events to tracked client sessions. + + Adapters call :meth:`track` with each session object seen while handling + a request and give a ``notifier(session, kind, uri)`` coroutine that + sends the right notification for that framework. Events raised from + worker threads are marshalled onto the event loop the adapter runs on. + Resource-updated events only go to sessions subscribed to the URI. + """ + + def __init__( + self, + layer: "CloudGMCPLayer", + notifier: Notifier, + publish: Callable[[str, "str | None"], Awaitable[None]] | None = None, + ) -> None: + self.notifier = notifier + self.publish = publish + self.sessions: "weakref.WeakSet[Any]" = weakref.WeakSet() + self.subscriptions: "weakref.WeakKeyDictionary[Any, set[str]]" = weakref.WeakKeyDictionary() + self.log_levels: "weakref.WeakKeyDictionary[Any, str]" = weakref.WeakKeyDictionary() + self.loop: asyncio.AbstractEventLoop | None = None + self._pending: set[asyncio.Future[Any]] = set() + ref = weakref.ref(self) + + def listener(kind: str, uri: str | None) -> None: + me = ref() + if me is not None: + me.on_change(kind, uri) + + layer.on_change(listener) + + def bind_loop(self) -> None: + if self.loop is None: + try: + self.loop = asyncio.get_running_loop() + except RuntimeError: + pass + + def track(self, session: Any) -> None: + self.bind_loop() + if session is None: + return + try: + self.sessions.add(session) + except TypeError: # not weak-referenceable + pass + + def subscribe(self, session: Any, uri: str) -> None: + try: + self.subscriptions.setdefault(session, set()).add(uri) + except TypeError: + pass + + def unsubscribe(self, session: Any, uri: str) -> None: + try: + self.subscriptions.get(session, set()).discard(uri) + except TypeError: + pass + + def set_level(self, session: Any, level: str) -> None: + try: + self.log_levels[session] = level + except TypeError: + pass + + def level_for(self, session: Any) -> str | None: + try: + return self.log_levels.get(session) + except TypeError: + return None + + def on_change(self, kind: str, uri: str | None) -> None: + run_on_loop(self.loop, self._schedule, kind, uri) + + def _schedule(self, kind: str, uri: str | None) -> None: + fut = asyncio.ensure_future(self._dispatch(kind, uri)) + self._pending.add(fut) + fut.add_done_callback(self._pending.discard) + + def _wants(self, session: Any, kind: str, uri: str | None) -> bool: + if kind != "resource": + return True + try: + subs = self.subscriptions.get(session) or set() + except TypeError: + subs = set() + return uri in subs + + async def _dispatch(self, kind: str, uri: str | None) -> None: + for ev_kind, ev_uri in expand_change(kind, uri): + if self.publish is not None: + try: + await self.publish(ev_kind, ev_uri) + except Exception: + logger.debug("subscription bus publish failed", exc_info=True) + for session in [s for s in list(self.sessions) if self._wants(s, ev_kind, ev_uri)]: + try: + await self.notifier(session, ev_kind, ev_uri) + except Exception: + logger.debug("dropping session after failed notification", exc_info=True) + self.sessions.discard(session) diff --git a/cloudg/mcp/adapters/export.py b/cloudg/mcp/adapters/export.py new file mode 100644 index 0000000..a1f910c --- /dev/null +++ b/cloudg/mcp/adapters/export.py @@ -0,0 +1,259 @@ +"""Framework-free export of the layer's MCP definitions. + +For custom servers, gateways and agent frameworks that do not use an MCP SDK: +:func:`export_definitions` returns the exact MCP wire definitions (as the +active policy shows them to a principal) plus async handlers for every +method. Every handler goes through the layer, so policies, transforms and +middleware still apply. + +Also small converters to function-calling formats: + +* :func:`to_openai_tools`: OpenAI ``tools=[{"type": "function", ...}]``. +* :func:`to_anthropic_tools`: Anthropic Messages API ``tools=[...]``. +* :func:`to_langchain_tools`: LangChain ``StructuredTool`` objects + (requires ``langchain-core``). +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Awaitable, Callable + +from cloudg.mcp.context import Principal +from cloudg.mcp.core import NotFoundError + +if TYPE_CHECKING: # pragma: no cover + from cloudg.mcp.layer import CloudGMCPLayer + +__all__ = [ + "ExportedDefinitions", + "export_definitions", + "to_anthropic_tools", + "to_langchain_tools", + "to_openai_tools", +] + +Handler = Callable[..., Awaitable[dict[str, Any]]] + + +@dataclass +class ExportedDefinitions: + """Wire definitions + handlers. + + Attributes: + server: ``{"name", "version", "instructions"}``. + tools / resources / resource_templates / prompts: MCP wire dicts + (camelCase, ready for ``tools/list`` etc.). + handlers: ``{method: async fn(params: dict, *, principal=None) -> result + wire dict}`` for ``tools/list``, ``tools/call``, ``resources/list``, + ``resources/templates/list``, ``resources/read``, ``prompts/list``, + ``prompts/get`` and ``completion/complete``. + tool_functions: ``{tool name: async fn(**arguments) -> result wire}``. + """ + + server: dict[str, Any] + tools: list[dict[str, Any]] + resources: list[dict[str, Any]] + resource_templates: list[dict[str, Any]] + prompts: list[dict[str, Any]] + handlers: dict[str, Handler] = field(default_factory=dict) + tool_functions: dict[str, Callable[..., Awaitable[dict[str, Any]]]] = field( + default_factory=dict + ) + + async def dispatch( + self, + method: str, + params: dict[str, Any] | None = None, + *, + principal: Principal | None = None, + ) -> dict[str, Any]: + """Run one MCP method; raises ``KeyError`` for unknown methods and + :class:`~cloudg.mcp.core.MCPLayerError` for protocol-level errors.""" + return await self.handlers[method](params or {}, principal=principal) + + def as_dict(self) -> dict[str, Any]: + return { + "server": self.server, + "tools": self.tools, + "resources": self.resources, + "resourceTemplates": self.resource_templates, + "prompts": self.prompts, + } + + +class _Handlers: + """Per-method MCP handlers bound to one layer; a handler called without + a principal acts as ``default``.""" + + def __init__(self, layer: "CloudGMCPLayer", default: Principal) -> None: + self.layer = layer + self.default = default + + def who(self, principal: Principal | None) -> Principal: + return principal or self.default + + async def tools_list( + self, params: dict[str, Any], *, principal: Principal | None = None + ) -> dict: + return {"tools": self.layer.tools_wire(self.who(principal))} + + async def tools_call( + self, params: dict[str, Any], *, principal: Principal | None = None + ) -> dict: + result = await self.layer.call_tool( + params["name"], params.get("arguments") or {}, principal=self.who(principal) + ) + return result.to_wire() + + async def resources_list( + self, params: dict[str, Any], *, principal: Principal | None = None + ) -> dict: + return {"resources": self.layer.resources_wire(self.who(principal))} + + async def templates_list( + self, params: dict[str, Any], *, principal: Principal | None = None + ) -> dict: + return {"resourceTemplates": self.layer.resource_templates_wire(self.who(principal))} + + async def resources_read( + self, params: dict[str, Any], *, principal: Principal | None = None + ) -> dict: + contents = await self.layer.read_resource(params["uri"], principal=self.who(principal)) + return {"contents": [c.to_wire() for c in contents]} + + async def prompts_list( + self, params: dict[str, Any], *, principal: Principal | None = None + ) -> dict: + return {"prompts": self.layer.prompts_wire(self.who(principal))} + + async def prompts_get( + self, params: dict[str, Any], *, principal: Principal | None = None + ) -> dict: + result = await self.layer.get_prompt( + params["name"], params.get("arguments") or {}, principal=self.who(principal) + ) + return result.to_wire() + + async def complete(self, params: dict[str, Any], *, principal: Principal | None = None) -> dict: + ctx = params.get("context") or {} + completion = await self.layer.complete( + params["ref"], + params["argument"], + principal=self.who(principal), + context_arguments=ctx.get("arguments"), + ) + return {"completion": completion} + + def by_method(self) -> dict[str, Handler]: + return { + "tools/list": self.tools_list, + "tools/call": self.tools_call, + "resources/list": self.resources_list, + "resources/templates/list": self.templates_list, + "resources/read": self.resources_read, + "prompts/list": self.prompts_list, + "prompts/get": self.prompts_get, + "completion/complete": self.complete, + } + + def tool_function(self, name: str) -> Callable[..., Awaitable[dict[str, Any]]]: + """``async fn(**arguments)`` calling tool ``name`` as ``default``.""" + layer, default = self.layer, self.default + + async def call(**arguments: Any) -> dict[str, Any]: + return (await layer.call_tool(name, arguments, principal=default)).to_wire() + + call.__name__ = name + return call + + +def export_definitions( + layer: "CloudGMCPLayer", principal: Principal | None = None +) -> ExportedDefinitions: + """Snapshot the layer's definitions for ``principal`` (default: local + user) and build per-method handlers.""" + default = principal or Principal.local() + handlers = _Handlers(layer, default) + tools = layer.tools_wire(default) + return ExportedDefinitions( + server={ + "name": layer.name, + "version": str(layer.version), + "instructions": layer.instructions, + }, + tools=tools, + resources=layer.resources_wire(default), + resource_templates=layer.resource_templates_wire(default), + prompts=layer.prompts_wire(default), + handlers=handlers.by_method(), + tool_functions={t["name"]: handlers.tool_function(t["name"]) for t in tools}, + ) + + +def to_openai_tools( + layer: "CloudGMCPLayer", principal: Principal | None = None +) -> list[dict[str, Any]]: + """OpenAI Chat Completions / Responses function-tool definitions.""" + out = [] + for t in layer.tools_wire(principal): + out.append( + { + "type": "function", + "function": { + "name": t["name"], + "description": t.get("description") or t.get("title") or t["name"], + "parameters": t["inputSchema"], + }, + } + ) + return out + + +def to_anthropic_tools( + layer: "CloudGMCPLayer", principal: Principal | None = None +) -> list[dict[str, Any]]: + """Anthropic Messages API tool definitions.""" + return [ + { + "name": t["name"], + "description": t.get("description") or t.get("title") or t["name"], + "input_schema": t["inputSchema"], + } + for t in layer.tools_wire(principal) + ] + + +def to_langchain_tools(layer: "CloudGMCPLayer", principal: Principal | None = None) -> list[Any]: + """LangChain ``StructuredTool`` objects that call the layer (async). + + The tool returns the result text; ``response_format="content_and_artifact"`` + exposes the structured content as the artifact. + """ + try: + from langchain_core.tools import StructuredTool + except ImportError as exc: # pragma: no cover (optional dependency) + raise ImportError("to_langchain_tools() needs `pip install langchain-core`") from exc + who = principal or Principal.local() + tools = [] + for t in layer.tools_wire(who): + name = t["name"] + + async def run(_name: str = name, **arguments: Any) -> tuple[str, Any]: + try: + result = await layer.call_tool(_name, arguments, principal=who) + except NotFoundError as exc: + return str(exc), None + text = "\n".join(c.text for c in result.content if hasattr(c, "text")) + return text, result.structured + + tools.append( + StructuredTool.from_function( + coroutine=run, + name=name, + description=t.get("description") or name, + args_schema=t["inputSchema"], + response_format="content_and_artifact", + ) + ) + return tools diff --git a/cloudg/mcp/adapters/fastmcp.py b/cloudg/mcp/adapters/fastmcp.py new file mode 100644 index 0000000..f44b98a --- /dev/null +++ b/cloudg/mcp/adapters/fastmcp.py @@ -0,0 +1,583 @@ +"""Adapter for the standalone ``fastmcp`` package (2.x, 3.x and 4.x). + +fastmcp manages tools, resources, templates and prompts as component +objects, so cloudg registers *component subclasses* whose definitions are +taken verbatim from the layer's wire dicts (``parameters`` = the layer's +``inputSchema``, plus ``output_schema``, ``annotations``, ``title`` and +``meta``) and whose ``run`` / ``read`` / ``render`` delegate to the layer. +fastmcp therefore never re-derives a schema from a Python signature, and +fastmcp features (its middleware, auth, ``mount``, tags, the in-memory +``fastmcp.Client``) keep working. + +On top of the components the adapter installs: + +* a fastmcp middleware that filters list results through the cloudg policy + for the calling principal (and refreshes per-principal ``outputSchema``), +* on fastmcp's inner SDK low-level server: cloudg's ``completion/complete`` + handling, ``resources/subscribe``, and (fastmcp 4 / mcp 2.x) a + ``subscriptions/listen`` handler (fastmcp itself serves none of these), +* change notifications (``layer.notify_change``) to the sessions seen, and + to listen streams on the 2026-07-28 protocol. + +Error mapping: unknown / hidden primitives raise fastmcp ``NotFoundError``; +other layer errors raise ``ToolError`` / ``ResourceError`` / ``PromptError``; +``isError`` tool results are passed through with their structured content. +""" + +from __future__ import annotations + +import base64 +import inspect +import logging +from types import SimpleNamespace +from typing import TYPE_CHECKING, Any + +from cloudg.mcp.adapters._common import ( + ChangeFanout, + call_layer, + lower_headers, + normalize_kinds, + normalize_level, + resolve_principal, + scope_principal, + send_change, +) +from cloudg.mcp.context import Principal +from cloudg.mcp.core import ( + BlobResourceContents, + MCPLayerError, + NotFoundError, + ResourceTemplateSpec, +) +from cloudg.mcp.native.auth import PrincipalResolver, RequestInfo + +if TYPE_CHECKING: # pragma: no cover + from cloudg.mcp.layer import CloudGMCPLayer + +logger = logging.getLogger("cloudg.mcp.adapters.fastmcp") + +__all__ = ["FastMCPBinding", "install", "is_fastmcp_server"] + + +def is_fastmcp_server(obj: Any) -> bool: + return any( + cls.__module__.split(".")[0] == "fastmcp" and cls.__name__ == "FastMCP" + for cls in type(obj).__mro__ + ) + + +def _fields(model: Any) -> set[str]: + return set(getattr(model, "model_fields", {}) or {}) + + +def _fastmcp_api() -> SimpleNamespace: + """The fastmcp classes the adapter needs, across fastmcp 2.x to 4.x.""" + import mcp.types as MT + from fastmcp import exceptions as fx + from fastmcp.prompts import Prompt + from fastmcp.resources import Resource, ResourceTemplate + from fastmcp.tools import Tool + + api = SimpleNamespace(MT=MT, fx=fx, Prompt=Prompt, Resource=Resource, Tool=Tool) + api.ResourceTemplate = ResourceTemplate + try: + from fastmcp.tools import ToolResult + except ImportError: # fastmcp 2.x + from fastmcp.tools.tool import ToolResult + try: + from fastmcp.prompts import PromptArgument + except ImportError: # fastmcp 2.x + from fastmcp.prompts.prompt import PromptArgument + api.ToolResult, api.PromptArgument = ToolResult, PromptArgument + try: + from fastmcp.resources import ResourceContent, ResourceResult + except ImportError: # fastmcp 2.x: read() returns str | bytes + ResourceContent = ResourceResult = None # type: ignore[assignment,misc] + try: + from fastmcp.prompts import Message, PromptResult + except ImportError: # fastmcp 2.x: render() returns list[PromptMessage] + Message = PromptResult = None # type: ignore[assignment,misc] + api.ResourceContent, api.ResourceResult = ResourceContent, ResourceResult + api.Message, api.PromptResult = Message, PromptResult + return api + + +def _component_classes(binding: "FastMCPBinding", api: SimpleNamespace) -> SimpleNamespace: + """fastmcp component subclasses whose run / read / render delegate to + ``binding``.""" + + class CloudGTool(api.Tool): # type: ignore[misc,name-defined] + async def run(self, arguments: dict[str, Any]) -> Any: + return await binding._run_tool(self.name, arguments) + + class CloudGResource(api.Resource): # type: ignore[misc,name-defined] + async def read(self) -> Any: + return await binding._read(str(self.uri)) + + class CloudGTemplate(api.ResourceTemplate): # type: ignore[misc,name-defined] + def matches(self, uri: str) -> dict[str, Any] | None: + spec = binding.layer.registry.templates.get(self.uri_template) + return spec.match(uri) if spec is not None else None + + async def create_resource(self, uri: str, params: dict[str, Any]) -> Any: + return binding._make_resource(CloudGResource, uri, self) + + async def read(self, arguments: dict[str, Any]) -> Any: + spec = binding.layer.registry.templates[self.uri_template] + return await binding._read(spec.expand(**arguments)) + + class CloudGPrompt(api.Prompt): # type: ignore[misc,name-defined] + async def render(self, arguments: dict[str, Any] | None = None) -> Any: + return await binding._render(self.name, arguments or {}) + + return SimpleNamespace( + Tool=CloudGTool, Resource=CloudGResource, Template=CloudGTemplate, Prompt=CloudGPrompt + ) + + +class FastMCPBinding: + """Registers the layer's primitives on a ``fastmcp.FastMCP`` server.""" + + def __init__( + self, + layer: "CloudGMCPLayer", + server: Any, + *, + include: Any = None, + principal_resolver: PrincipalResolver | None = None, + ) -> None: + api = _fastmcp_api() + self.layer = layer + self.server = server + self.kinds = normalize_kinds(include) + self.resolver = principal_resolver + self.MT = api.MT + self.fx = api.fx + self._ResourceContent = api.ResourceContent + self._ResourceResult = api.ResourceResult + self._FMessage = api.Message + self._FPromptResult = api.PromptResult + self._FToolResult = api.ToolResult + self._FPromptArgument = api.PromptArgument + self.fanout = ChangeFanout(layer, self._notify) + components = _component_classes(self, api) + self.CloudGTool = components.Tool + self.CloudGResource = components.Resource + self.CloudGTemplate = components.Template + self.CloudGPrompt = components.Prompt + self._result_cls = self._make_result_class(api.ToolResult) + + self.lowlevel: Any = None + self._register() + self._install_middleware() + self._install_lowlevel() + if self.lowlevel is not None: + # one fan-out for everything: the low-level binding also serves + # resources/subscribe and the 2026-07-28 subscriptions/listen bus + self.fanout = self.lowlevel.fanout + + # ------------------------------------------------------------------ + # Registration + # ------------------------------------------------------------------ + + def _register(self) -> None: + if "tools" in self.kinds: + self._register_tools() + if "resources" in self.kinds: + for spec in self.layer.registry.resources.values(): + self.server.add_resource(self._make_resource(self.CloudGResource, spec.uri, spec)) + if "templates" in self.kinds: + self._register_templates() + if "prompts" in self.kinds: + self._register_prompts() + + def _register_tools(self) -> None: + layer, tool_fields = self.layer, _fields(self.CloudGTool) + for spec in layer.registry.tools.values(): + wire = spec.to_wire(layer.exposed_name(spec.name)) + kwargs: dict[str, Any] = { + "name": wire["name"], + "description": wire.get("description") or None, + "parameters": wire["inputSchema"], + "output_schema": wire.get("outputSchema"), + "tags": {"cloudg", spec.category, *spec.tags}, + } + if wire.get("annotations"): + kwargs["annotations"] = self.MT.ToolAnnotations.model_validate(wire["annotations"]) + _add_optional(kwargs, tool_fields, wire, title=wire.get("title")) + self.server.add_tool(self.CloudGTool(**kwargs)) + + def _register_templates(self) -> None: + tmpl_fields = _fields(self.CloudGTemplate) + for tspec in self.layer.registry.templates.values(): + wire = tspec.to_wire() + kwargs: dict[str, Any] = { + "uri_template": tspec.uri_template, + "name": tspec.name, + "description": tspec.description or None, + "mime_type": tspec.mime_type, + "parameters": _template_parameters(tspec), + "tags": {"cloudg", tspec.category, *tspec.tags}, + } + _add_optional(kwargs, tmpl_fields, wire, title=tspec.title) + if wire.get("annotations") and "annotations" in tmpl_fields: + kwargs["annotations"] = self.MT.Annotations.model_validate(wire["annotations"]) + self.server.add_template(self.CloudGTemplate(**kwargs)) + + def _register_prompts(self) -> None: + layer, prompt_fields = self.layer, _fields(self.CloudGPrompt) + for pspec in layer.registry.prompts.values(): + arguments = [ + self._FPromptArgument( + name=a.name, description=a.description or None, required=a.required + ) + for a in pspec.arguments + ] + kwargs: dict[str, Any] = { + "name": layer.exposed_name(pspec.name), + "description": pspec.description or None, + "arguments": arguments, + "tags": {"cloudg", pspec.category, *pspec.tags}, + } + _add_optional(kwargs, prompt_fields, pspec.to_wire(), title=pspec.title) + self.server.add_prompt(self.CloudGPrompt(**kwargs)) + + def _make_resource(self, cls: Any, uri: str, source: Any) -> Any: + fields = _fields(cls) + kwargs: dict[str, Any] = { + "uri": uri, + "name": getattr(source, "name", "") or uri, + "description": getattr(source, "description", None) or None, + "mime_type": getattr(source, "mime_type", None) or "application/json", + } + tags = getattr(source, "tags", None) + if tags is not None: + kwargs["tags"] = {"cloudg", *tags} + title = getattr(source, "title", None) + if title and "title" in fields: + kwargs["title"] = title + to_wire = getattr(source, "to_wire", None) + if callable(to_wire) and "meta" in fields: + try: + kwargs["meta"] = to_wire().get("_meta") + except (TypeError, ValueError, AttributeError): # a fastmcp component, not a spec + logger.debug("no cloudg wire metadata for resource %s", uri) + return cls(**kwargs) + + def _install_middleware(self) -> None: + try: + from fastmcp.server.middleware import Middleware + except ImportError: # pragma: no cover (fastmcp < 2.9) + logger.warning( + "fastmcp has no middleware support; cloudg policy filtering of " + "list results is disabled (calls are still enforced)" + ) + return + binding = self + + class CloudGPolicyFilter(Middleware): # type: ignore[misc,valid-type] + async def on_list_tools(self, context: Any, call_next: Any) -> Any: + return binding._filter_tools(await call_next(context)) + + async def on_list_resources(self, context: Any, call_next: Any) -> Any: + return binding._filter_resources(await call_next(context)) + + async def on_list_resource_templates(self, context: Any, call_next: Any) -> Any: + return binding._filter_templates(await call_next(context)) + + async def on_list_prompts(self, context: Any, call_next: Any) -> Any: + return binding._filter_prompts(await call_next(context)) + + self.server.add_middleware(CloudGPolicyFilter()) + + def _install_lowlevel(self) -> None: + """Completions, resource subscriptions and (2026-07-28) the listen + bus are protocol-level features fastmcp components cannot express: + bind them on fastmcp's inner SDK low-level server.""" + kinds = self.kinds & {"completions", "subscriptions"} + inner = getattr(self.server, "_mcp_server", None) + if inner is None or not kinds: + return + from cloudg.mcp.adapters.lowlevel import install as install_lowlevel + + try: + self.lowlevel = install_lowlevel( + self.layer, + inner, + include=kinds, + principal_resolver=self.resolver, + list_changed=False, + ) + except Exception: + logger.debug("could not bind cloudg protocol features on fastmcp", exc_info=True) + + # ------------------------------------------------------------------ + # Per-request helpers + # ------------------------------------------------------------------ + + def _context(self) -> Any: + try: + from fastmcp.server.dependencies import get_context + + return get_context() + except Exception: + return None + + def request_info(self, ctx: Any = None) -> RequestInfo: + request = None + headers: dict[str, str] = {} + token = None + try: + from fastmcp.server.dependencies import get_http_request + + request = get_http_request() + except Exception: + request = None + if request is not None: + headers = lower_headers(getattr(request, "headers", None)) + try: + from fastmcp.server.dependencies import get_access_token + + token = get_access_token() + except Exception: + token = None + principal = scope_principal(request) + return RequestInfo( + transport="http" if request is not None else "stdio", + headers=headers, + session_id=headers.get("mcp-session-id"), + access_token=token, + principal=principal, + raw=ctx, + ) + + def _principal(self, ctx: Any = None) -> Principal: + return resolve_principal(self.resolver, self.request_info(ctx)) + + def _track(self, ctx: Any) -> None: + if ctx is None: + return + try: + session = ctx.session + except Exception: + return + target = getattr(session, "_connection", None) or session + self.fanout.track(target) + + def _tool_context(self, ctx: Any, principal: Principal) -> Any: + last: list[float] = [] + + async def progress(value: float, total: float | None, message: str | None) -> None: + if ctx is None or (last and value <= last[0]): + return + last[:] = [value] + await ctx.report_progress(value, total, message) + + async def log(level: Any, data: Any, logger_name: str | None) -> None: + if ctx is None: + return + text = data if isinstance(data, str) else _json(data) + await ctx.log(text, level=normalize_level(level), logger_name=logger_name) + + request_id = None + if ctx is not None: + try: + request_id = ctx.request_id + except Exception: + request_id = None + return self.layer.context( + principal, + request_id=request_id, + progress_callback=progress, + log_callback=log, + session=ctx, + ) + + # ------------------------------------------------------------------ + # Execution + # ------------------------------------------------------------------ + + def _make_result_class(self, base: Any) -> Any: + if hasattr(base, "from_mcp_result"): + return None # fastmcp 4 preserves the exact wire result itself + try: + import pydantic + + is_model = issubclass(base, pydantic.BaseModel) + except Exception: + is_model = False + if is_model: + from pydantic import PrivateAttr + + class RawToolResult(base): # type: ignore[misc,valid-type] + _cloudg_raw: Any = PrivateAttr(default=None) + + def to_mcp_result(self) -> Any: + if self._cloudg_raw is not None: + return self._cloudg_raw + return super().to_mcp_result() + + return RawToolResult + + class PlainRawToolResult(base): # type: ignore[misc,valid-type] + _cloudg_raw: Any = None + + def to_mcp_result(self) -> Any: + if self._cloudg_raw is not None: + return self._cloudg_raw + return super().to_mcp_result() + + return PlainRawToolResult + + def _tool_result(self, wire: dict[str, Any]) -> Any: + call_result = self.MT.CallToolResult.model_validate(wire) + if self._result_cls is None: + return self._FToolResult.from_mcp_result(call_result) + kwargs: dict[str, Any] = { + "content": list(call_result.content), + "structured_content": wire.get("structuredContent"), + "meta": wire.get("_meta"), + } + if "is_error" in inspect.signature(self._result_cls.__init__).parameters: + kwargs["is_error"] = bool(wire.get("isError")) + result = self._result_cls(**kwargs) + result._cloudg_raw = call_result + return result + + async def _call(self, method: str, target: str, *args: Any, error_cls: Any) -> Any: + """Run ``layer.(target, *args)`` for the current fastmcp + request; layer errors map to fastmcp errors, anything else to a + generic ``error_cls("Internal error")``.""" + ctx = self._context() + self._track(ctx) + principal = self._principal(ctx) + call = getattr(self.layer, method)( + target, *args, principal=principal, context=self._tool_context(ctx, principal) + ) + + def mapped(exc: MCPLayerError) -> Exception: + if isinstance(exc, NotFoundError): + return self.fx.NotFoundError(exc.message) + return error_cls(exc.message) + + return await call_layer(call, mapped, lambda: error_cls("Internal error")) + + async def _run_tool(self, name: str, arguments: dict[str, Any]) -> Any: + result = await self._call("call_tool", name, arguments or {}, error_cls=self.fx.ToolError) + return self._tool_result(result.to_wire()) + + async def _read(self, uri: str) -> Any: + contents = await self._call("read_resource", uri, error_cls=self.fx.ResourceError) + if self._ResourceResult is not None: + items = [] + for c in contents: + raw: Any = ( + base64.b64decode(c.blob) if isinstance(c, BlobResourceContents) else c.text + ) + items.append(self._ResourceContent(raw, mime_type=c.mime_type, meta=c.meta)) + return self._ResourceResult(items) + if not contents: + return "" + first = contents[0] + if isinstance(first, BlobResourceContents): + return base64.b64decode(first.blob) + return first.text + + async def _render(self, name: str, arguments: dict[str, Any]) -> Any: + error_cls = getattr(self.fx, "PromptError", self.fx.ToolError) + result = await self._call("get_prompt", name, arguments, error_cls=error_cls) + wire = result.to_wire() + MT = self.MT + if self._FPromptResult is None: + return [MT.PromptMessage.model_validate(m) for m in wire["messages"]] + messages = [] + for m in wire["messages"]: + content = m["content"] + ctype = content.get("type") + model = { + "text": MT.TextContent, + "image": MT.ImageContent, + "audio": getattr(MT, "AudioContent", None), + "resource": MT.EmbeddedResource, + }.get(ctype) + value = model.model_validate(content) if model is not None else _json(content) + messages.append(self._FMessage(value, role=m["role"])) + return self._FPromptResult(messages, description=wire.get("description")) + + # ------------------------------------------------------------------ + # Policy filtering of list results + # ------------------------------------------------------------------ + + def _filter_tools(self, tools: Any) -> Any: + principal = self._principal(self._context()) + wires = {w["name"]: w for w in self.layer.tools_wire(principal)} + out = [] + for tool in tools: + if isinstance(tool, self.CloudGTool): + wire = wires.get(tool.name) + if wire is None: + continue + if wire.get("outputSchema") != tool.output_schema: + tool = tool.model_copy(update={"output_schema": wire.get("outputSchema")}) + out.append(tool) + return out + + def _filter_resources(self, resources: Any) -> Any: + principal = self._principal(self._context()) + allowed = {w["uri"] for w in self.layer.resources_wire(principal)} + return [ + r for r in resources if not isinstance(r, self.CloudGResource) or str(r.uri) in allowed + ] + + def _filter_templates(self, templates: Any) -> Any: + principal = self._principal(self._context()) + allowed = {w["uriTemplate"] for w in self.layer.resource_templates_wire(principal)} + return [ + t + for t in templates + if not isinstance(t, self.CloudGTemplate) or t.uri_template in allowed + ] + + def _filter_prompts(self, prompts: Any) -> Any: + principal = self._principal(self._context()) + allowed = {w["name"] for w in self.layer.prompts_wire(principal)} + return [p for p in prompts if not isinstance(p, self.CloudGPrompt) or p.name in allowed] + + # ------------------------------------------------------------------ + # Change notifications + # ------------------------------------------------------------------ + + async def _notify(self, session: Any, kind: str, uri: str | None) -> None: + await send_change(session, kind, uri) + + +def _add_optional( + kwargs: dict[str, Any], fields: set[str], wire: dict[str, Any], *, title: str | None +) -> None: + """Set ``title`` / ``meta`` when this fastmcp version has those fields.""" + if "title" in fields and title: + kwargs["title"] = title + if "meta" in fields: + kwargs["meta"] = wire.get("_meta") + + +def _template_parameters(spec: ResourceTemplateSpec) -> dict[str, Any]: + props = {v: {"type": "string"} for v in spec.variables} + return {"type": "object", "properties": props, "required": list(props)} + + +def _json(value: Any) -> str: + import json + + return json.dumps(value, ensure_ascii=False, default=str) + + +def install( + layer: "CloudGMCPLayer", + server: Any, + *, + include: Any = None, + principal_resolver: PrincipalResolver | None = None, +) -> FastMCPBinding: + """Register the layer's primitives on a standalone ``fastmcp.FastMCP``.""" + return FastMCPBinding(layer, server, include=include, principal_resolver=principal_resolver) diff --git a/cloudg/mcp/adapters/lowlevel.py b/cloudg/mcp/adapters/lowlevel.py new file mode 100644 index 0000000..33f10f8 --- /dev/null +++ b/cloudg/mcp/adapters/lowlevel.py @@ -0,0 +1,593 @@ +"""Adapter for the official SDK's low-level ``Server`` (mcp 1.x and 2.x). + +The low-level server dispatches JSON-RPC methods to handler callables. This +adapter *wraps* the handlers already registered (if any) instead of +replacing them: list methods return the host's own primitives followed by +the cloudg ones, and ``tools/call`` / ``resources/read`` / ``prompts/get`` / +``completion/complete`` route by ownership: cloudg names / URIs go to the +:class:`~cloudg.mcp.layer.CloudGMCPLayer`, everything else to the original +handler. A server with no handlers simply becomes a cloudg server. + +Because tools are listed from the layer's wire dicts, the exact +``inputSchema`` / ``outputSchema`` / ``annotations`` / ``title`` / ``_meta`` +reach the client; the SDK never re-derives anything from Python signatures. + +Version differences handled here: + +* mcp 2.x: ``Server.add_request_handler(method, params_type, handler)`` + with ``handler(ctx: ServerRequestContext, params)``; progress via + ``ctx.session.report_progress``; log messages via + ``ctx.session.send_log_message``; legacy-era change notifications through + the per-connection ``Connection``; 2026-07-28 change notifications through + the ``subscriptions/listen`` bus (installed when the server has none). +* mcp 1.x: ``Server.request_handlers[RequestType] = handler(req)`` + returning ``ServerResult``; request context from ``server.request_context``. + +The higher-level ``MCPServer`` (2.x) and ``FastMCP`` (1.x) are adapted by +applying this module to their inner low-level server +(:mod:`cloudg.mcp.adapters.mcp_sdk`). +""" + +from __future__ import annotations + +import logging +import warnings +from typing import TYPE_CHECKING, Any, Awaitable, Callable + +from cloudg.mcp.adapters._common import ( + ChangeFanout, + call_layer, + layer_error_parts, + lower_headers, + normalize_kinds, + normalize_level, + owns_prompt, + owns_resource, + owns_template, + owns_tool, + resolve_principal, + scope_principal, + send_change, +) +from cloudg.mcp.context import Principal +from cloudg.mcp.native.auth import PrincipalResolver, RequestInfo +from cloudg.mcp.native.jsonrpc import INTERNAL_ERROR, INVALID_PARAMS, LOG_LEVELS, RESOURCE_NOT_FOUND + +if TYPE_CHECKING: # pragma: no cover + from cloudg.mcp.layer import CloudGMCPLayer + +logger = logging.getLogger("cloudg.mcp.adapters.lowlevel") + +__all__ = ["LowLevelBinding", "install", "is_lowlevel_server"] + +OrigCall = Callable[[], Awaitable[Any]] | None + +# method -> (v1 request class, v2 params class, result class, kind gate, impl name) +_METHODS: list[tuple[str, str, str, str, str]] = [ + ("tools/list", "ListToolsRequest", "PaginatedRequestParams", "tools", "_list_tools"), + ("tools/call", "CallToolRequest", "CallToolRequestParams", "tools", "_call_tool"), + ( + "resources/list", + "ListResourcesRequest", + "PaginatedRequestParams", + "resources", + "_list_resources", + ), + ( + "resources/templates/list", + "ListResourceTemplatesRequest", + "PaginatedRequestParams", + "templates", + "_list_templates", + ), + ( + "resources/read", + "ReadResourceRequest", + "ReadResourceRequestParams", + "resources", + "_read_resource", + ), + ("prompts/list", "ListPromptsRequest", "PaginatedRequestParams", "prompts", "_list_prompts"), + ("prompts/get", "GetPromptRequest", "GetPromptRequestParams", "prompts", "_get_prompt"), + ("completion/complete", "CompleteRequest", "CompleteRequestParams", "completions", "_complete"), + ( + "resources/subscribe", + "SubscribeRequest", + "SubscribeRequestParams", + "subscriptions", + "_subscribe", + ), + ( + "resources/unsubscribe", + "UnsubscribeRequest", + "UnsubscribeRequestParams", + "subscriptions", + "_unsubscribe", + ), + ("logging/setLevel", "SetLevelRequest", "SetLevelRequestParams", "logging", "_set_level"), +] + + +def is_lowlevel_server(obj: Any) -> bool: + try: + from mcp.server.lowlevel import Server + except Exception: + return False + return isinstance(obj, Server) + + +def _dump(params: Any) -> dict[str, Any]: + if params is None: + return {} + if isinstance(params, dict): + return params + try: + return params.model_dump(by_alias=True, mode="json", exclude_none=True) + except Exception: + return {} + + +class LowLevelBinding: + """The installed adapter; keeps the change fan-out alive. + + Attributes: + layer / server: What was bound. + sdk_major: 1 or 2. + fanout: Forwards ``layer.notify_change`` to connected clients. + """ + + def __init__( + self, + layer: "CloudGMCPLayer", + server: Any, + *, + include: Any = None, + principal_resolver: PrincipalResolver | None = None, + owner: Any = None, + list_changed: bool = True, + ) -> None: + import mcp.types as T + + self.layer = layer + self.server = server + self.owner = owner + self.T = T + self.kinds = normalize_kinds(include) + self.resolver = principal_resolver + self.sdk_major = 2 if hasattr(server, "add_request_handler") else 1 + self.modern_versions: tuple[str, ...] = () + if self.sdk_major == 2: + try: + from mcp_types.version import MODERN_PROTOCOL_VERSIONS + + self.modern_versions = tuple(MODERN_PROTOCOL_VERSIONS) + except Exception: # pragma: no cover (layout change) + self.modern_versions = ("2026-07-28",) + self.bus = self._find_bus() if self.sdk_major == 2 else None + self.fanout = ChangeFanout(layer, self._notify, self._publish if self.bus else None) + self._install() + if list_changed: + self._patch_initialization_options() + if self.sdk_major == 2: + self._patch_input_schema() + try: + server.__dict__.setdefault("_cloudg_bindings", []).append(self) + except (AttributeError, TypeError): # slotted server class: nothing to record on + logger.debug("cannot record the cloudg binding on %s", type(server).__name__) + + # ------------------------------------------------------------------ + # Installation + # ------------------------------------------------------------------ + + def _install(self) -> None: + T = self.T + for method, v1_req, v2_params, kind, impl_name in _METHODS: + if kind not in self.kinds: + continue + impl = getattr(self, impl_name) + if self.sdk_major == 2: + entry = self.server.get_request_handler(method) + orig = entry.handler if entry is not None else None + ptype = entry.params_type if entry is not None else getattr(T, v2_params) + self.server.add_request_handler(method, ptype, self._v2_handler(impl, orig)) + else: + req_cls = getattr(T, v1_req, None) + if req_cls is None: # very old 1.x without this method + continue + orig = self.server.request_handlers.get(req_cls) + self.server.request_handlers[req_cls] = self._v1_handler(impl, orig) + + def _v2_handler(self, impl: Any, orig: Any) -> Any: + async def handler(ctx: Any, params: Any) -> Any: + orig_call = (lambda: orig(ctx, params)) if orig is not None else None + return await impl(ctx, _dump(params), orig_call) + + return handler + + def _v1_handler(self, impl: Any, orig: Any) -> Any: + T = self.T + server = self.server + + async def handler(req: Any) -> Any: + try: + ctx = server.request_context + except LookupError: + ctx = None + + async def orig_call() -> Any: + res = await orig(req) + return getattr(res, "root", res) + + params = _dump(getattr(req, "params", None)) if req is not None else {} + result = await impl(ctx, params, orig_call if orig is not None else None) + return T.ServerResult(result) + + return handler + + def _find_bus(self) -> Any: + bus = getattr(self.owner, "_subscriptions", None) + if bus is not None: + return bus + entry = self.server.get_request_handler("subscriptions/listen") + if entry is not None: + return getattr(entry.handler, "_bus", None) + if "subscriptions" not in self.kinds: + return None + try: + from mcp.server.subscriptions import InMemorySubscriptionBus, ListenHandler + + bus = InMemorySubscriptionBus() + self.server.add_request_handler( + "subscriptions/listen", + self.T.SubscriptionsListenRequestParams, + ListenHandler(bus), + ) + return bus + except Exception: + logger.debug("could not install a subscriptions/listen handler", exc_info=True) + return None + + def _patch_initialization_options(self) -> None: + server = self.server + if getattr(server, "_cloudg_init_patched", False): + return + try: + from mcp.server.lowlevel import NotificationOptions + except Exception: # pragma: no cover + return + original = server.create_initialization_options + + def create_initialization_options( + notification_options: Any = None, *args: Any, **kwargs: Any + ) -> Any: + if notification_options is None: + notification_options = NotificationOptions( + prompts_changed=True, resources_changed=True, tools_changed=True + ) + return original(notification_options, *args, **kwargs) + + server.create_initialization_options = create_initialization_options + server._cloudg_init_patched = True + + def _patch_input_schema(self) -> None: + server = self.server + previous = getattr(server, "get_tool_input_schema", None) + layer = self.layer + + def get_tool_input_schema(name: str) -> Any: + if owns_tool(layer, name): + spec = layer.registry.tools.get(layer.internal_name(name)) + return spec.input_schema if spec is not None else None + return previous(name) if previous is not None else None + + if "tools" in self.kinds and previous is not None: + server.get_tool_input_schema = get_tool_input_schema + + # ------------------------------------------------------------------ + # Request context helpers + # ------------------------------------------------------------------ + + def _is_modern(self, ctx: Any) -> bool: + return bool(self.modern_versions) and getattr(ctx, "protocol_version", None) in ( + self.modern_versions + ) + + def _session_key(self, ctx: Any) -> Any: + if ctx is None: + return None + session = getattr(ctx, "session", None) + if self.sdk_major == 2: + if self._is_modern(ctx): + return None + return getattr(session, "_connection", None) + return session + + def request_info(self, ctx: Any) -> RequestInfo: + request = getattr(ctx, "request", None) if ctx is not None else None + headers = lower_headers(getattr(request, "headers", None)) if request is not None else {} + principal = scope_principal(request) + token = None + try: + from mcp.server.auth.middleware.auth_context import get_access_token + + token = get_access_token() + except Exception: + token = None + client_info = None + try: + params = ctx.session.client_params # type: ignore[union-attr] + info = getattr(params, "client_info", None) or getattr(params, "clientInfo", None) + if info is not None: + client_info = info.model_dump(mode="json") + except Exception: + client_info = None + return RequestInfo( + transport="http" if request is not None else "stdio", + headers=headers, + session_id=headers.get("mcp-session-id"), + client_info=client_info, + access_token=token, + principal=principal, + raw=ctx, + ) + + def _principal(self, ctx: Any) -> Principal: + return resolve_principal(self.resolver, self.request_info(ctx)) + + def _track(self, ctx: Any) -> None: + self.fanout.track(self._session_key(ctx)) + + def _error(self, code: int, message: str, data: Any = None) -> Exception: + if self.sdk_major == 2: + from mcp.shared.exceptions import MCPError + + return MCPError(code, message, data) + from mcp.shared.exceptions import McpError + from mcp.types import ErrorData + + return McpError(ErrorData(code=code, message=message, data=data)) + + async def _through_layer(self, call: Any, **error_kw: Any) -> Any: + """Await a layer call; layer errors become SDK errors with the same + JSON-RPC code, anything else a generic internal error.""" + return await call_layer( + call, + lambda exc: self._error(*layer_error_parts(exc, **error_kw)), + lambda: self._error(INTERNAL_ERROR, "Internal error"), + ) + + def _tool_context(self, ctx: Any, principal: Principal) -> Any: + session = getattr(ctx, "session", None) if ctx is not None else None + request_id = getattr(ctx, "request_id", None) if ctx is not None else None + modern = self._is_modern(ctx) + key = self._session_key(ctx) + fanout = self.fanout + last: list[float] = [] + + async def progress(value: float, total: float | None, message: str | None) -> None: + if session is None: + return + if last and value <= last[0]: + return + last[:] = [value] + if self.sdk_major == 2: + await session.report_progress(value, total, message) + return + meta = getattr(ctx, "meta", None) + token = getattr(meta, "progressToken", None) if meta is not None else None + if token is None: + return + try: + await session.send_progress_notification( + token, value, total, message=message, related_request_id=request_id + ) + except TypeError: # older 1.x without ``message`` + await session.send_progress_notification(token, value, total) + + async def log(level: Any, data: Any, logger_name: str | None) -> None: + if session is None: + return + lvl = normalize_level(level) + if not modern: + threshold = fanout.level_for(key) or "info" + if LOG_LEVELS.index(lvl) < LOG_LEVELS.index(threshold): + return + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + await session.send_log_message( + level=lvl, data=data, logger=logger_name, related_request_id=request_id + ) + + return self.layer.context( + principal, + request_id=request_id, + progress_callback=progress, + log_callback=log, + session=session, + ) + + async def _merge( + self, orig_call: OrigCall, ours: list[dict[str, Any]], field: str, key: str, model: Any + ) -> Any: + if orig_call is None: + return model.model_validate({field: ours}) + base = await orig_call() + data = ( + base + if isinstance(base, dict) + else base.model_dump(by_alias=True, mode="json", exclude_none=True) + ) + if data.get("nextCursor"): + return base # cloudg entries are appended to the host's last page + existing = list(data.get(field) or []) + seen = {item.get(key) for item in existing} + clashes = [o[key] for o in ours if o.get(key) in seen] + if clashes: + logger.warning("host server already defines %s %s; keeping the host's", field, clashes) + data = dict(data) + data[field] = existing + [o for o in ours if o.get(key) not in seen] + return model.model_validate(data) + + # ------------------------------------------------------------------ + # Method implementations (params are wire dicts) + # ------------------------------------------------------------------ + + async def _list_tools(self, ctx: Any, p: dict[str, Any], orig_call: OrigCall) -> Any: + self._track(ctx) + ours = self.layer.tools_wire(self._principal(ctx)) + return await self._merge(orig_call, ours, "tools", "name", self.T.ListToolsResult) + + async def _call_tool(self, ctx: Any, p: dict[str, Any], orig_call: OrigCall) -> Any: + name = str(p.get("name", "")) + if not owns_tool(self.layer, name): + if orig_call is not None: + return await orig_call() + raise self._error(INVALID_PARAMS, f"Unknown tool: {name}") + self._track(ctx) + principal = self._principal(ctx) + tctx = self._tool_context(ctx, principal) + call = self.layer.call_tool( + name, p.get("arguments") or {}, principal=principal, context=tctx + ) + result = await self._through_layer(call, data={"name": name}) + return self.T.CallToolResult.model_validate(result.to_wire()) + + async def _list_resources(self, ctx: Any, p: dict[str, Any], orig_call: OrigCall) -> Any: + self._track(ctx) + ours = self.layer.resources_wire(self._principal(ctx)) + return await self._merge(orig_call, ours, "resources", "uri", self.T.ListResourcesResult) + + async def _list_templates(self, ctx: Any, p: dict[str, Any], orig_call: OrigCall) -> Any: + self._track(ctx) + ours = self.layer.resource_templates_wire(self._principal(ctx)) + return await self._merge( + orig_call, ours, "resourceTemplates", "uriTemplate", self.T.ListResourceTemplatesResult + ) + + async def _read_resource(self, ctx: Any, p: dict[str, Any], orig_call: OrigCall) -> Any: + uri = str(p.get("uri", "")) + if not owns_resource(self.layer, uri): + if orig_call is not None: + return await orig_call() + code = INVALID_PARAMS if self._is_modern(ctx) else RESOURCE_NOT_FOUND + raise self._error(code, f"Unknown resource: {uri}", {"uri": uri}) + self._track(ctx) + principal = self._principal(ctx) + tctx = self._tool_context(ctx, principal) + code = INVALID_PARAMS if self._is_modern(ctx) else RESOURCE_NOT_FOUND + call = self.layer.read_resource(uri, principal=principal, context=tctx) + contents = await self._through_layer(call, not_found_code=code, data={"uri": uri}) + return self.T.ReadResourceResult.model_validate( + {"contents": [c.to_wire() for c in contents]} + ) + + async def _list_prompts(self, ctx: Any, p: dict[str, Any], orig_call: OrigCall) -> Any: + self._track(ctx) + ours = self.layer.prompts_wire(self._principal(ctx)) + return await self._merge(orig_call, ours, "prompts", "name", self.T.ListPromptsResult) + + async def _get_prompt(self, ctx: Any, p: dict[str, Any], orig_call: OrigCall) -> Any: + name = str(p.get("name", "")) + if not owns_prompt(self.layer, name): + if orig_call is not None: + return await orig_call() + raise self._error(INVALID_PARAMS, f"Unknown prompt: {name}") + self._track(ctx) + principal = self._principal(ctx) + tctx = self._tool_context(ctx, principal) + call = self.layer.get_prompt( + name, p.get("arguments") or {}, principal=principal, context=tctx + ) + result = await self._through_layer(call) + return self.T.GetPromptResult.model_validate(result.to_wire()) + + async def _complete(self, ctx: Any, p: dict[str, Any], orig_call: OrigCall) -> Any: + ref = p.get("ref") or {} + argument = p.get("argument") or {} + ours = ref.get("type") == "ref/prompt" and owns_prompt(self.layer, str(ref.get("name"))) + ours = ours or ( + ref.get("type") == "ref/resource" and owns_template(self.layer, str(ref.get("uri"))) + ) + if not ours: + if orig_call is not None: + return await orig_call() + return self.T.CompleteResult.model_validate({"completion": {"values": []}}) + context = p.get("context") or {} + call = self.layer.complete( + ref, + argument, + principal=self._principal(ctx), + context_arguments=context.get("arguments"), + ) + completion = await self._through_layer(call) + return self.T.CompleteResult.model_validate({"completion": completion}) + + async def _subscribe(self, ctx: Any, p: dict[str, Any], orig_call: OrigCall) -> Any: + key = self._session_key(ctx) + self.fanout.track(key) + self.fanout.subscribe(key, str(p.get("uri", ""))) + if orig_call is not None: + return await orig_call() + return self.T.EmptyResult() + + async def _unsubscribe(self, ctx: Any, p: dict[str, Any], orig_call: OrigCall) -> Any: + self.fanout.unsubscribe(self._session_key(ctx), str(p.get("uri", ""))) + if orig_call is not None: + return await orig_call() + return self.T.EmptyResult() + + async def _set_level(self, ctx: Any, p: dict[str, Any], orig_call: OrigCall) -> Any: + level = p.get("level") + if level in LOG_LEVELS: + self.fanout.set_level(self._session_key(ctx), str(level)) + if orig_call is not None: + return await orig_call() + return self.T.EmptyResult() + + # ------------------------------------------------------------------ + # Change notifications + # ------------------------------------------------------------------ + + async def _notify(self, session: Any, kind: str, uri: str | None) -> None: + await send_change(session, kind, uri) + + async def _publish(self, kind: str, uri: str | None) -> None: + from mcp.shared.subscriptions import ( + PromptsListChanged, + ResourcesListChanged, + ResourceUpdated, + ToolsListChanged, + ) + + event: Any + if kind == "tools": + event = ToolsListChanged() + elif kind == "prompts": + event = PromptsListChanged() + elif kind == "resources": + event = ResourcesListChanged() + elif kind == "resource" and uri: + event = ResourceUpdated(uri=uri) + else: + return + await self.bus.publish(event) + + +def install( + layer: "CloudGMCPLayer", + server: Any, + *, + include: Any = None, + principal_resolver: PrincipalResolver | None = None, + owner: Any = None, + list_changed: bool = True, +) -> LowLevelBinding: + """Bind ``layer`` into an mcp low-level ``Server`` (1.x or 2.x).""" + return LowLevelBinding( + layer, + server, + include=include, + principal_resolver=principal_resolver, + owner=owner, + list_changed=list_changed, + ) diff --git a/cloudg/mcp/adapters/mcp_sdk.py b/cloudg/mcp/adapters/mcp_sdk.py new file mode 100644 index 0000000..a2a5607 --- /dev/null +++ b/cloudg/mcp/adapters/mcp_sdk.py @@ -0,0 +1,108 @@ +"""Adapter for the official ``mcp`` SDK's high-level servers. + +* mcp **2.x** ``MCPServer`` (``from mcp.server.mcpserver import MCPServer``, + formerly FastMCP). +* mcp **1.x** ``FastMCP`` (``from mcp.server.fastmcp import FastMCP``). + +Neither class has a public API for registering a tool with an explicit JSON +input schema; their decorators always derive schemas from function +signatures. Both are thin layers over a low-level ``Server`` +(``MCPServer._lowlevel_server`` / ``FastMCP._mcp_server``), so cloudg binds +there with :mod:`cloudg.mcp.adapters.lowlevel`: the host's own decorated +tools keep working and cloudg's primitives are served with their exact wire +definitions. The inner attribute is private, so this path is checked against +mcp 1.x (1.30) and 2.x (2.3); :func:`install` raises a clear error if a +future SDK renames it. + +Also provides constructors for fresh SDK servers used by +:func:`cloudg.mcp.adapters.build_server` and :mod:`cloudg.mcp.server`. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from cloudg.mcp.adapters._common import major_version +from cloudg.mcp.adapters.lowlevel import LowLevelBinding +from cloudg.mcp.adapters.lowlevel import install as install_lowlevel +from cloudg.mcp.native.auth import PrincipalResolver + +if TYPE_CHECKING: # pragma: no cover + from cloudg.mcp.layer import CloudGMCPLayer + +__all__ = [ + "create_highlevel_server", + "create_lowlevel_server", + "inner_lowlevel_server", + "install", + "is_sdk_highlevel_server", + "sdk_major", +] + + +def sdk_major() -> int | None: + """Major version of the installed ``mcp`` package (``None`` if absent).""" + return major_version("mcp") + + +def is_sdk_highlevel_server(obj: Any) -> bool: + module = type(obj).__module__ + return module.startswith("mcp.server.mcpserver") or module.startswith("mcp.server.fastmcp") + + +def inner_lowlevel_server(server: Any) -> Any: + for attr in ("_lowlevel_server", "_mcp_server"): + inner = getattr(server, attr, None) + if inner is not None: + return inner + raise TypeError( + f"Cannot find the low-level server inside {type(server).__name__}; this mcp SDK " + "version is not supported by cloudg's adapter (tested with mcp 1.30 and 2.3). " + "Build a low-level Server with cloudg.mcp.adapters.build_server(layer) instead." + ) + + +def install( + layer: "CloudGMCPLayer", + server: Any, + *, + include: Any = None, + principal_resolver: PrincipalResolver | None = None, + list_changed: bool = True, +) -> LowLevelBinding: + """Bind ``layer`` into an SDK ``MCPServer`` (2.x) or ``FastMCP`` (1.x).""" + return install_lowlevel( + layer, + inner_lowlevel_server(server), + include=include, + principal_resolver=principal_resolver, + owner=server, + list_changed=list_changed, + ) + + +def create_lowlevel_server(layer: "CloudGMCPLayer", *, name: str | None = None) -> Any: + """A bare SDK low-level ``Server`` carrying the layer's identity.""" + from mcp.server.lowlevel import Server + + kwargs: dict[str, Any] = {"version": str(layer.version)} + if layer.instructions: + kwargs["instructions"] = layer.instructions + try: + return Server(name or layer.name, **kwargs) + except TypeError: # very old 1.x without instructions= + kwargs.pop("instructions", None) + return Server(name or layer.name, **kwargs) + + +def create_highlevel_server(layer: "CloudGMCPLayer", *, name: str | None = None) -> Any: + """A fresh ``MCPServer`` (mcp 2.x) or ``FastMCP`` (mcp 1.x).""" + if (sdk_major() or 0) >= 2: + from mcp.server.mcpserver import MCPServer + + return MCPServer( + name or layer.name, instructions=layer.instructions or None, version=str(layer.version) + ) + from mcp.server.fastmcp import FastMCP + + return FastMCP(name or layer.name, instructions=layer.instructions or None) diff --git a/cloudg/mcp/catalog/__init__.py b/cloudg/mcp/catalog/__init__.py new file mode 100644 index 0000000..ef2baca --- /dev/null +++ b/cloudg/mcp/catalog/__init__.py @@ -0,0 +1,89 @@ +"""The cloudg MCP catalog: every tool, resource, resource template and +prompt the layer exposes, grouped by category. + +Each module has a ``register(registry)`` function. :func:`default_registry` +registers them all; build your own :class:`~cloudg.mcp.core.Registry` from +a subset to expose less:: + + from cloudg.mcp.core import Registry + from cloudg.mcp.catalog import inventory, graph + + reg = Registry() + inventory.register(reg) + graph.register(reg) + layer = CloudGMCPLayer(registry=reg) + +(or filter with ``CloudGMCPLayer(include_categories=[...])``). +""" + +from __future__ import annotations + +import logging +from types import ModuleType + +from cloudg.mcp.core import Registry + +logger = logging.getLogger("cloudg.mcp") + +CATEGORIES: dict[str, str] = { + "workspace": "Load, select, snapshot, diff and unload datasets", + "inventory": "Search, inspect and aggregate assets, accounts, regions, tags, coverage, org", + "graph": "Relationships: neighbours, paths, exposure, attack / lateral paths, " + "dependencies, blast radius, centrality, sub-graphs", + "findings": "Browse, summarise, prioritise, suppress and ingest security findings", + "compliance": "Framework posture, controls and compliance gaps", + "ontology": "RDF ontology, read-only SPARQL, semantic neighbourhoods, RAG chunks", + "export": "Terraform recreation and report files", + "live": "Collect from cloud APIs and run scanners (credentials / binaries; open world)", + "meta": "Server capabilities and cloudg's vocabulary", + "privacy": "Privacy policy, pseudonym vault and data-handling controls", + "prompts": "Packaged analysis workflows", +} + + +def _modules() -> list[ModuleType]: + from cloudg.mcp.catalog import ( + compliance, + export, + findings, + graph, + inventory, + live, + meta, + ontology, + prompts, + resources, + workspace, + ) + + return [ + workspace, + inventory, + graph, + findings, + compliance, + ontology, + export, + live, + meta, + resources, + prompts, + ] + + +def default_registry() -> Registry: + """A registry with the full cloudg catalog (plus the privacy module + when it is available).""" + reg = Registry() + for mod in _modules(): + mod.register(reg) + try: + from cloudg.mcp.catalog import privacy + except ImportError: + logger.debug("cloudg.mcp.catalog.privacy not available; skipping privacy tools") + else: + privacy.register(reg) + return reg + + +__all__ = ["CATEGORIES", "default_registry"] diff --git a/cloudg/mcp/catalog/_common.py b/cloudg/mcp/catalog/_common.py new file mode 100644 index 0000000..7827403 --- /dev/null +++ b/cloudg/mcp/catalog/_common.py @@ -0,0 +1,545 @@ +"""Shared helpers for the catalog modules. + +Conventions every catalog tool follows so an agent can rely on them: + +- ``dataset`` argument: empty means the active dataset. +- Asset references (``ref``) accept the internal id, the ARN / cloud + resource id, a unique name, or a unique ARN tail; misses come back as an + error listing close matches. +- List tools are bounded: ``limit`` (default 50) plus an opaque ``cursor``; + results carry ``total``, ``returned``, ``next_cursor`` and ``truncated``. + A cursor is tied to the query and dataset version that produced it, so a + stale cursor is rejected instead of silently returning the wrong page. +- Assets are summarised as compact "briefs" with a ``uri`` + (``cloudg://assets/{id}``) the client can read for full detail. +""" + +from __future__ import annotations + +import dataclasses +import difflib +import hashlib +import json +import re +from enum import Enum +from typing import Annotated, Any, Callable, Iterable, Sequence, TypeVar +from urllib.parse import quote + +from pydantic import BaseModel, ConfigDict, Field + +from cloudg.mcp.core import InvalidArgumentsError +from cloudg.mcp.state import SEVERITY_RANK, Dataset +from cloudg.schema.models import CloudAsset, Finding, NetworkEdge, Severity + +T = TypeVar("T") +E = TypeVar("E", bound=Enum) +F = TypeVar("F", bound=Callable[..., Any]) + +#: Dataset summary keys every headline view (live run results, prompt +#: overviews) shows, in display order. +SUMMARY_CORE_KEYS = ( + "open_findings", + "severity_breakdown", + "accounts", + "regions", + "internet_exposed", + "cross_account_edges", +) +DEFAULT_LIMIT = 50 +MAX_LIMIT = 500 + + +class Catalog: + """Records tool / resource / prompt declarations made at import time + and replays them onto a :class:`~cloudg.mcp.core.Registry`. + + Catalog modules declare their handlers as module-level functions with + ``@CATALOG.tool(...)`` (same arguments as :meth:`Registry.tool`) and + their ``register(reg)`` is just ``CATALOG.register(reg)``. Declaration + order is registration order. + """ + + def __init__(self) -> None: + self._specs: list[tuple[str, tuple[Any, ...], dict[str, Any], Callable[..., Any]]] = [] + + def _record(self, kind: str, *args: Any, **kwargs: Any) -> Callable[[F], F]: + def deco(fn: F) -> F: + self._specs.append((kind, args, kwargs, fn)) + return fn + + return deco + + def tool(self, *args: Any, **kwargs: Any) -> Callable[[F], F]: + return self._record("tool", *args, **kwargs) + + def resource(self, *args: Any, **kwargs: Any) -> Callable[[F], F]: + return self._record("resource", *args, **kwargs) + + def resource_template(self, *args: Any, **kwargs: Any) -> Callable[[F], F]: + return self._record("resource_template", *args, **kwargs) + + def prompt(self, *args: Any, **kwargs: Any) -> Callable[[F], F]: + return self._record("prompt", *args, **kwargs) + + def register(self, reg: Any) -> None: + for kind, args, kwargs, fn in self._specs: + getattr(reg, kind)(*args, **kwargs)(fn) + + +DatasetArg = Annotated[ + str, Field(description="Dataset name (see list_datasets). Empty = the active dataset.") +] +RefArg = Annotated[ + str, + Field( + min_length=1, + description="Asset reference: internal id, ARN / cloud resource id, unique name, " + "or unique ARN tail (e.g. 'function:api' or 'i-0abc').", + ), +] +Limit = Annotated[int, Field(ge=1, le=MAX_LIMIT, description="Maximum items to return.")] +Cursor = Annotated[ + str, Field(description="Opaque cursor from a previous call's next_cursor; empty = first page.") +] + +ASSET_FIELDS = ( + "id", + "name", + "type", + "provider", + "region", + "account_id", + "arn", + "internet_exposed", + "open_findings", + "max_severity", + "tags", + "uri", +) +FINDING_FIELDS = ( + "id", + "title", + "severity", + "risk_score", + "source_tool", + "resource_id", + "resource_arn", + "asset_id", + "asset_name", + "compliance_frameworks", + "is_suppressed", + "cvss_score", + "detected_at", + "uri", +) + + +def ws_dataset(ctx: Any, dataset: str = "") -> Dataset: + return ctx.workspace.get(dataset or None) + + +def asset_uri(asset_id: str) -> str: + return "cloudg://assets/" + quote(asset_id, safe="") + + +def finding_uri(finding_id: str) -> str: + return "cloudg://findings/" + quote(finding_id, safe="") + + +def max_severity(findings: Iterable[Finding]) -> str | None: + best = None + for f in findings: + if best is None or SEVERITY_RANK[f.severity.value] > SEVERITY_RANK[best]: + best = f.severity.value + return best + + +def asset_brief(ds: Dataset, a: CloudAsset, *, tags: bool = False) -> dict[str, Any]: + open_f = ds.open_findings(a.id) + out: dict[str, Any] = { + "id": a.id, + "name": a.name, + "type": a.asset_type.value, + "provider": a.provider.value, + "region": a.region, + "account_id": a.account_id, + "arn": a.arn, + "internet_exposed": a.is_internet_exposed, + "open_findings": len(open_f), + "max_severity": max_severity(open_f), + "uri": asset_uri(a.id), + } + if tags: + out["tags"] = dict(a.tags) + return out + + +def node_brief(ds: Dataset, node_id: str) -> dict[str, Any]: + """Brief for a graph node: an asset, or a placeholder (CIDR / external).""" + a = ds.by_id.get(node_id) + if a is not None: + return asset_brief(ds, a) + data = ds.graph.nodes.get(node_id, {}) if node_id in ds.graph else {} + return { + "id": node_id, + "name": data.get("name", node_id), + "type": data.get("asset_type", "EXTERNAL"), + "external": True, + } + + +def edge_brief(ds: Dataset, e: NetworkEdge) -> dict[str, Any]: + s, t = ds.by_id.get(e.source_id), ds.by_id.get(e.target_id) + out: dict[str, Any] = { + "id": e.id, + "source": e.source_id, + "source_name": s.name if s else e.source_id, + "target": e.target_id, + "target_name": t.name if t else e.target_id, + "edge_type": e.edge_type.value, + "relationship": e.relationship, + } + for k in ("port_range", "protocol", "cidr", "description"): + v = getattr(e, k) + if v: + out[k] = v + if e.ports: + out["ports"] = e.ports[:20] + if e.edge_type.value in ("SECURITY_GROUP_RULE", "NACL_RULE"): + out["direction"] = e.direction + if s and t and s.account_id and t.account_id and s.account_id != t.account_id: + out["cross_account"] = True + return out + + +def finding_brief(ds: Dataset, f: Finding) -> dict[str, Any]: + aid = ds.finding_asset_id(f) + a = ds.by_id.get(aid) if aid else None + return { + "id": f.id, + "title": f.title, + "severity": f.severity.value, + "risk_score": f.risk_score, + "source_tool": f.source_tool, + "resource_id": f.resource_id, + "resource_arn": f.resource_arn, + "asset_id": a.id if a else None, + "asset_name": a.name if a else None, + "compliance_frameworks": list(f.compliance_frameworks), + "is_suppressed": f.is_suppressed, + "cvss_score": f.cvss_score, + "detected_at": f.detected_at.isoformat() if f.detected_at else None, + "uri": finding_uri(f.id), + } + + +# --------------------------------------------------------------------------- +# Argument parsing +# --------------------------------------------------------------------------- + + +def parse_enum(value: str, enum: type[E], label: str) -> E: + """Case-insensitive enum lookup with a helpful error.""" + v = (value or "").strip().upper().replace("-", "_").replace(" ", "_") + for attr in ("value", "name"): + for m in enum: + if getattr(m, attr) == v: + return m + names = [m.value for m in enum] + close = difflib.get_close_matches(v, names, n=3, cutoff=0.5) + hint = f" Did you mean {', '.join(close)}?" if close else "" + shown = ", ".join(names) if len(names) <= 30 else ", ".join(names[:30]) + ", ..." + # The value itself goes in data, not in the text (see cloudg.mcp.state.base) + raise InvalidArgumentsError( + f"Unknown {label}.{hint} Valid values: {shown}", + data={"value": value, "valid": names, "close_matches": close}, + ) + + +def parse_enums(values: Sequence[str] | None, enum: type[E], label: str) -> set[E]: + return {parse_enum(v, enum, label) for v in values or []} + + +def parse_severity(value: str) -> Severity: + return parse_enum(value, Severity, "severity") + + +class ArgsFilter: + """Base for the frozen dataclass filters the list tools build from + their arguments (FindingFilter, AssetFilter).""" + + @classmethod + def from_args(cls: type[T], args: dict[str, Any]) -> T: + """Build from a tool's arguments (``locals()``): the matching keys.""" + names = (f.name for f in dataclasses.fields(cls)) # type: ignore[arg-type] + return cls(**{name: args[name] for name in names if name in args}) + + def key(self) -> dict[str, Any]: + """The filter as a dict, for cursor fingerprints.""" + return dataclasses.asdict(self) # type: ignore[call-overload] + + +def check_fields(fields: Sequence[str] | None, allowed: Sequence[str]) -> list[str] | None: + if not fields: + return None + bad = [f for f in fields if f not in allowed] + if bad: + raise InvalidArgumentsError( + f"Unknown field(s) {', '.join(bad)}. Valid fields: {', '.join(allowed)}" + ) + out = list(dict.fromkeys(["id", *fields])) + return out + + +def project(item: dict[str, Any], fields: list[str] | None) -> dict[str, Any]: + if not fields: + return item + return {k: item.get(k) for k in fields} + + +# --------------------------------------------------------------------------- +# Pagination +# --------------------------------------------------------------------------- + + +def fingerprint(ds: Dataset | None, **params: Any) -> str: + """Hash of the query and the dataset state a cursor belongs to: name, + version and the per-load instance id, so a dataset replaced under the + same name (version back to 0) rejects the old cursors.""" + raw = json.dumps( + { + "ds": ds.name if ds else None, + "v": ds.version if ds else None, + "i": getattr(ds, "instance_id", None) if ds else None, + **params, + }, + sort_keys=True, + default=str, + ) + return hashlib.sha256(raw.encode()).hexdigest()[:12] + + +def encode_cursor(offset: int, fp: str) -> str: + # Short and punctuated on purpose: privacy policies redact long opaque + # base64-looking tokens (and anything starting "eyJ") as secrets. + return f"c{offset}.{fp}" + + +_CURSOR_RE = re.compile(r"^c(\d{1,9})\.([0-9a-f]{12})$") + + +def decode_cursor(cursor: str, fp: str) -> int: + if not cursor: + return 0 + m = _CURSOR_RE.match(cursor.strip()) + if not m: + raise InvalidArgumentsError( + "Invalid cursor. Pass the next_cursor value from the previous result unchanged, " + "or omit cursor to start from the first page." + ) + if m.group(2) != fp: + raise InvalidArgumentsError( + "Stale cursor: the query arguments or the dataset changed since it was issued. " + "Repeat the call without cursor to start again." + ) + return int(m.group(1)) + + +def paginate( + items: Sequence[T], limit: int, cursor: str, fp: str +) -> tuple[list[T], dict[str, Any]]: + """Slice ``items``; returns (page, envelope fields).""" + start = decode_cursor(cursor, fp) + page = list(items[start : start + limit]) + end = start + len(page) + more = end < len(items) + return page, { + "total": len(items), + "offset": start, + "returned": len(page), + "next_cursor": encode_cursor(end, fp) if more else None, + "truncated": more, + } + + +def bounded(items: Sequence[T], limit: int) -> tuple[list[T], dict[str, Any]]: + """First ``limit`` items plus total/truncated (no cursor).""" + return list(items[:limit]), { + "total": len(items), + "returned": min(limit, len(items)), + "truncated": len(items) > limit, + } + + +# --------------------------------------------------------------------------- +# Output schemas (documented shapes; every field optional so that policy +# transforms such as projection / redaction never make results invalid) +# --------------------------------------------------------------------------- + + +class _Out(BaseModel): + model_config = ConfigDict(extra="allow") + + +class AssetBriefOut(_Out): + id: str | None = None + name: str | None = None + type: str | None = None + provider: str | None = None + region: str | None = None + account_id: str | None = None + arn: str | None = None + internet_exposed: bool | None = None + open_findings: int | None = None + max_severity: str | None = None + uri: str | None = None + + +class PageOut(_Out): + dataset: str | None = None + total: int | None = None + offset: int | None = None + returned: int | None = None + next_cursor: str | None = None + truncated: bool | None = None + + +class AssetPageOut(PageOut): + items: list[AssetBriefOut] = Field(default_factory=list) + + +class EdgeOut(_Out): + id: str | None = None + source: str | None = None + source_name: str | None = None + target: str | None = None + target_name: str | None = None + edge_type: str | None = None + relationship: str | None = None + + +class EdgePageOut(PageOut): + items: list[EdgeOut] = Field(default_factory=list) + + +class FindingBriefOut(_Out): + id: str | None = None + title: str | None = None + severity: str | None = None + risk_score: float | None = None + source_tool: str | None = None + asset_id: str | None = None + asset_name: str | None = None + compliance_frameworks: list[str] | None = None + is_suppressed: bool | None = None + uri: str | None = None + + +class FindingPageOut(PageOut): + items: list[FindingBriefOut] = Field(default_factory=list) + + +class AssetDetailOut(AssetBriefOut): + tags: dict[str, str] | None = None + relations: dict[str, Any] | None = None + findings: list[FindingBriefOut] | None = None + metadata_keys: list[str] | None = None + + +class NeighborsOut(_Out): + dataset: str | None = None + asset: AssetBriefOut | None = None + nodes: list[AssetBriefOut] = Field(default_factory=list) + edges: list[EdgeOut] = Field(default_factory=list) + truncated: bool | None = None + + +class PathOut(_Out): + length: int | None = None + nodes: list[AssetBriefOut] = Field(default_factory=list) + edge_types: list[str] = Field(default_factory=list) + + +class PathsOut(_Out): + dataset: str | None = None + paths: list[PathOut] = Field(default_factory=list) + total_found: int | None = None + truncated: bool | None = None + + +class DependencyOut(_Out): + dataset: str | None = None + asset: AssetBriefOut | None = None + total: int | None = None + items: list[dict[str, Any]] = Field(default_factory=list) + truncated: bool | None = None + + +class SummaryOut(_Out): + dataset: str | None = None + total_assets: int | None = None + total_edges: int | None = None + total_findings: int | None = None + open_findings: int | None = None + severity_breakdown: dict[str, int] | None = None + + +class CountOut(_Out): + dataset: str | None = None + group_by: str | None = None + total: int | None = None + groups: dict[str, int] = Field(default_factory=dict) + + +class RiskOut(_Out): + asset: AssetBriefOut | None = None + score: float | None = None + components: dict[str, Any] | None = None + top_findings: list[FindingBriefOut] = Field(default_factory=list) + + +class TopRisksOut(_Out): + dataset: str | None = None + items: list[RiskOut] = Field(default_factory=list) + total_candidates: int | None = None + + +class WorkspaceOut(_Out): + active_dataset: str | None = None + datasets: list[dict[str, Any]] = Field(default_factory=list) + allowed_roots: list[str] | None = None + output_dir: str | None = None + + +class DiffSectionOut(_Out): + count: int | None = None + items: list[dict[str, Any]] = Field(default_factory=list) + truncated: bool | None = None + + +class DiffOut(_Out): + base: str | None = None + target: str | None = None + assets: dict[str, Any] | None = None + edges: dict[str, Any] | None = None + findings: dict[str, Any] | None = None + + +class SparqlOut(_Out): + dataset: str | None = None + query_type: str | None = None + rows: list[dict[str, Any]] = Field(default_factory=list) + returned: int | None = None + truncated: bool | None = None + answer: bool | None = None + + +class ComplianceSummaryOut(_Out): + dataset: str | None = None + frameworks: list[dict[str, Any]] = Field(default_factory=list) + + +def schema(model: type[BaseModel]) -> dict[str, Any]: + from cloudg.mcp.schema import output_schema_for + + return output_schema_for(model) diff --git a/cloudg/mcp/catalog/_compliance.py b/cloudg/mcp/catalog/_compliance.py new file mode 100644 index 0000000..5e9d1b2 --- /dev/null +++ b/cloudg/mcp/catalog/_compliance.py @@ -0,0 +1,330 @@ +"""Helpers behind the compliance tools: the shipped ruleset catalog, the +per-framework rollup, control rows, one control's status and the +per-asset gap rows.""" + +from __future__ import annotations + +import functools +import logging +from pathlib import Path +from typing import Any + +from cloudg.mcp.catalog._common import asset_brief, finding_brief, max_severity +from cloudg.mcp.state import SEVERITY_RANK, Dataset, ReferenceNotFoundError +from cloudg.schema.models import ComplianceResult, Finding + +logger = logging.getLogger("cloudg.mcp") + +_STATUS_KEY = {"FAIL": "controls_failing", "PASS": "controls_passing"} # nosec B105 - status labels, not credentials + + +def _load_yaml(path: Path) -> dict[str, Any]: + import yaml + + data = yaml.safe_load(path.read_text()) + return data if isinstance(data, dict) else {} + + +def _add_ruleset(out: dict[str, dict[str, Any]], path: Path, data: dict[str, Any]) -> None: + fw = str(data.get("framework") or path.stem) + entry = out.setdefault( + fw, + { + "framework": fw, + "versions": set(), + "providers": set(), + "sources": set(), + "files": [], + "controls": {}, + }, + ) + if data.get("version"): + entry["versions"].add(str(data["version"])) + if data.get("provider"): + entry["providers"].add(str(data["provider"])) + if data.get("source"): + entry["sources"].add(str(data["source"]).split(",")[0]) + entry["files"].append(path.name) + for c in data.get("controls", []) or []: + cid = str(c.get("id", "")) + if cid: + entry["controls"].setdefault( + cid, + { + "title": c.get("title"), + "severity": c.get("severity"), + "checks": len(c.get("checks", []) or []), + }, + ) + + +@functools.lru_cache(maxsize=4) +def ruleset_catalog(rules_dir: str) -> dict[str, dict[str, Any]]: + """framework -> {versions, providers, files, controls: {id: title}}.""" + out: dict[str, dict[str, Any]] = {} + for path in sorted(Path(rules_dir).rglob("*.yaml")): + if path.name == "check_equivalence.yaml": + continue + try: + data = _load_yaml(path) + except Exception as exc: # a broken ruleset file must not break the catalog + logger.debug("Skipping ruleset %s: %s", path, exc) + continue + _add_ruleset(out, path, data) + return out + + +def rules_dir(ctx: Any) -> str: + rd = getattr(getattr(ctx.config, "rulesets", None), "rules_dir", None) + if rd: + return str(rd) + from cloudg.config import _default_rules_dir + + return _default_rules_dir() + + +def fw_match(name: str, wanted: str) -> bool: + return not wanted or wanted.lower() == name.lower() or wanted.lower() in name.lower() + + +# --------------------------------------------------------------------------- +# Framework rollup +# --------------------------------------------------------------------------- + + +def _rollup_entry(fws: dict[str, dict[str, Any]], fw: str) -> dict[str, Any]: + return fws.setdefault( + fw, + { + "framework": fw, + "controls_total": 0, + "controls_failing": 0, + "controls_passing": 0, + "controls_other": 0, + "findings": set(), + "assets": set(), + "severity": {}, + }, + ) + + +def _collect_rollup(ds: Dataset, framework: str) -> dict[str, dict[str, Any]]: + fws: dict[str, dict[str, Any]] = {} + for c in ds.compliance: + if fw_match(c.framework, framework): + e = _rollup_entry(fws, c.framework) + e["controls_total"] += 1 + e[_STATUS_KEY.get(c.status.value, "controls_other")] += 1 + e["findings"].update(c.finding_ids) + for f in ds.findings: + if f.is_suppressed: + continue + for fw in f.compliance_frameworks: + if fw_match(fw, framework): + _rollup_entry(fws, fw)["findings"].add(f.id) + return fws + + +def _settle_findings(ds: Dataset, e: dict[str, Any]) -> None: + """Drop missing / suppressed findings and count severities and assets.""" + for fid in list(e["findings"]): + f = ds.findings_by_id.get(fid) + if f is None or f.is_suppressed: + e["findings"].discard(fid) + continue + e["severity"][f.severity.value] = e["severity"].get(f.severity.value, 0) + 1 + aid = ds.finding_asset_id(f) + if aid: + e["assets"].add(aid) + + +def _rollup_row(e: dict[str, Any]) -> dict[str, Any]: + total = e["controls_total"] + return { + "framework": e["framework"], + "controls_evaluated": total, + "controls_failing": e["controls_failing"], + "controls_passing": e["controls_passing"], + "controls_other": e["controls_other"], + "pass_rate": round(e["controls_passing"] / total * 100, 1) if total else None, + "open_findings": len(e["findings"]), + "severity_breakdown": dict( + sorted(e["severity"].items(), key=lambda kv: -SEVERITY_RANK[kv[0]]) + ), + "affected_assets": len(e["assets"]), + "uri": f"cloudg://compliance/{e['framework']}", + } + + +def framework_rollup(ds: Dataset, framework: str = "") -> list[dict[str, Any]]: + """Per-framework posture from compliance results + finding tags.""" + fws = _collect_rollup(ds, framework) + for e in fws.values(): + _settle_findings(ds, e) + return [_rollup_row(e) for e in sorted(fws.values(), key=lambda x: -len(x["findings"]))] + + +# --------------------------------------------------------------------------- +# Controls +# --------------------------------------------------------------------------- + + +def open_control_findings(ds: Dataset, c: ComplianceResult) -> list[Finding]: + return [ + ds.findings_by_id[i] + for i in c.finding_ids + if i in ds.findings_by_id and not ds.findings_by_id[i].is_suppressed + ] + + +def ruleset_control_rows(cat: dict[str, Any], framework: str) -> tuple[list[str], list[dict]]: + matches = [fw for fw in cat if fw_match(fw, framework)] + if not matches: + raise ReferenceNotFoundError( + "No ruleset framework matches the name. Known: " + ", ".join(sorted(cat)), + data={"value": framework, "known": sorted(cat)}, + ) + rows = [ + {"framework": fw, "control_id": cid, **meta} + for fw in matches + for cid, meta in cat[fw]["controls"].items() + ] + return matches, rows + + +def dataset_control_rows(ds: Dataset, framework: str, status: str) -> list[dict[str, Any]]: + rows = [] + for c in ds.compliance: + if not fw_match(c.framework, framework) or (status and c.status.value != status): + continue + open_f = open_control_findings(ds, c) + rows.append( + { + "framework": c.framework, + "control_id": c.control_id, + "control_title": c.control_title, + "status": c.status.value, + "findings": len(c.finding_ids), + "open_findings": len(open_f), + "max_severity": max_severity(open_f), + } + ) + rows.sort( + key=lambda r: ( + -SEVERITY_RANK.get(r["max_severity"] or "", -1), + -r["open_findings"], + r["control_id"], + ) + ) + return rows + + +def control_definition(cat: dict[str, Any], framework: str, control_id: str) -> dict | None: + for fw, e in cat.items(): + if fw_match(fw, framework) and control_id in e["controls"]: + return {"framework": fw, "control_id": control_id, **e["controls"][control_id]} + return None + + +def control_results(ds: Dataset, framework: str, control_id: str) -> list[ComplianceResult]: + return [ + c + for c in ds.compliance + if fw_match(c.framework, framework) + and (c.control_id == control_id or c.control_id.endswith(control_id)) + ] + + +def missing_control_error(ds: Dataset, framework: str, control_id: str) -> Exception: + known = sorted({c.control_id for c in ds.compliance if fw_match(c.framework, framework)}) + hint = ( + "Controls in this dataset: " + ", ".join(known[:10]) + if known + else "Use list_controls(framework=..., source='ruleset') to browse." + ) + return ReferenceNotFoundError( + "No such control for that framework. " + hint, + data={"value": control_id, "framework": framework, "known": known[:10]}, + ) + + +def control_state(results: list[ComplianceResult], open_f: list[Finding]) -> str: + if not results: + return "NOT_EVALUATED" + failing = any(c.status.value == "FAIL" for c in results) + status = "FAIL" if failing and open_f else results[0].status.value + if status == "FAIL" and not open_f: + return "PASS (all findings suppressed)" + return status + + +def control_findings(ds: Dataset, results: list[ComplianceResult]) -> list[Finding]: + fids = list(dict.fromkeys(i for c in results for i in c.finding_ids)) + return [ + ds.findings_by_id[i] + for i in fids + if i in ds.findings_by_id and not ds.findings_by_id[i].is_suppressed + ] + + +# --------------------------------------------------------------------------- +# Gaps +# --------------------------------------------------------------------------- + + +def _failing_controls(ds: Dataset, framework: str) -> dict[str, list[str]]: + ctrl_by_finding: dict[str, list[str]] = {} + for c in ds.compliance: + if c.status.value == "FAIL" and fw_match(c.framework, framework): + for fid in c.finding_ids: + ctrl_by_finding.setdefault(fid, []).append(f"{c.framework}:{c.control_id}") + return ctrl_by_finding + + +def _gap_groups(ds: Dataset, framework: str, min_rank: int) -> dict[str, dict[str, Any]]: + ctrl_by_finding = _failing_controls(ds, framework) + per_asset: dict[str, dict[str, Any]] = {} + for f in ds.findings: + if f.is_suppressed or SEVERITY_RANK[f.severity.value] < min_rank: + continue + fws = [fw for fw in f.compliance_frameworks if fw_match(fw, framework)] + ctrls = ctrl_by_finding.get(f.id, []) + if not fws and not ctrls: + continue + aid = ds.finding_asset_id(f) or f"unmapped:{f.resource_arn or f.resource_id}" + e = per_asset.setdefault(aid, {"frameworks": set(), "controls": set(), "findings": []}) + e["frameworks"].update(fws or [c.split(":")[0] for c in ctrls]) + e["controls"].update(ctrls) + e["findings"].append(f) + return per_asset + + +def gap_rows(ds: Dataset, framework: str, min_rank: int) -> list[dict[str, Any]]: + rows = [] + for aid, e in _gap_groups(ds, framework, min_rank).items(): + a = ds.by_id.get(aid) + base = ( + asset_brief(ds, a) if a else {"id": aid, "name": aid.split(":", 1)[1], "unmapped": True} + ) + rows.append( + { + **base, + "frameworks": sorted(e["frameworks"]), + "failing_controls": sorted(e["controls"])[:20], + "gap_findings": len(e["findings"]), + "gap_max_severity": max_severity(e["findings"]), + } + ) + rows.sort( + key=lambda r: ( + -SEVERITY_RANK.get(r["gap_max_severity"] or "", -1), + not r.get("internet_exposed"), + -r["gap_findings"], + r["name"], + ) + ) + return rows + + +def briefs(ds: Dataset, findings: list[Finding]) -> list[dict[str, Any]]: + return [finding_brief(ds, f) for f in findings] diff --git a/cloudg/mcp/catalog/_findings.py b/cloudg/mcp/catalog/_findings.py new file mode 100644 index 0000000..d4f6e0e --- /dev/null +++ b/cloudg/mcp/catalog/_findings.py @@ -0,0 +1,236 @@ +"""Helpers behind the findings tools: filtering, sorting, grouping, risk +ranking, compliance merging and report ingestion.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Any, Callable, Sequence + +from cloudg.mcp.catalog._common import ( + ArgsFilter, + asset_brief, + finding_brief, + parse_enums, + parse_severity, +) +from cloudg.mcp.core import InvalidArgumentsError +from cloudg.mcp.state import SEVERITY_RANK, Dataset, check_prowler_input +from cloudg.schema.models import ComplianceResult, ComplianceStatus, Finding, Severity + +FindingPredicate = Callable[[Finding], bool] + + +@dataclass(frozen=True) +class FindingFilter(ArgsFilter): + """The filters list_findings accepts (empty = no filter).""" + + severities: Sequence[str] | None = None + min_severity: str = "" + source_tool: str = "" + framework: str = "" + resource: str = "" + query: str = "" + include_suppressed: bool = False + + +def _predicates(flt: FindingFilter) -> list[FindingPredicate]: + preds: list[FindingPredicate] = [] + if not flt.include_suppressed: + preds.append(lambda f: not f.is_suppressed) + sevs = {s.value for s in parse_enums(flt.severities, Severity, "severity")} + if sevs: + preds.append(lambda f: f.severity.value in sevs) + if flt.min_severity: + min_rank = SEVERITY_RANK[parse_severity(flt.min_severity).value] + preds.append(lambda f: SEVERITY_RANK[f.severity.value] >= min_rank) + tool = flt.source_tool.lower() + if tool: + preds.append(lambda f: tool in f.source_tool.lower()) + fw = flt.framework.lower() + if fw: + preds.append(lambda f: any(fw in x.lower() for x in f.compliance_frameworks)) + q = flt.query.lower().strip() + if q: + preds.append(lambda f: q in f.title.lower() or q in (f.description or "").lower()) + return preds + + +def filter_findings(ds: Dataset, flt: FindingFilter | None = None) -> list[Finding]: + """Findings of ``ds`` matching every filter of ``flt``.""" + flt = flt or FindingFilter() + preds = _predicates(flt) + asset_id = ds.resolve_asset(flt.resource).id if flt.resource else None + pool = ds.findings_by_asset.get(asset_id, []) if asset_id else ds.findings + return [f for f in pool if all(p(f) for p in preds)] + + +_SORT_KEYS: dict[str, Callable[[Finding], Any]] = { + "risk": lambda f: (-f.risk_score, -SEVERITY_RANK[f.severity.value], f.id), + "severity": lambda f: (-SEVERITY_RANK[f.severity.value], -f.risk_score, f.id), + "source_tool": lambda f: (f.source_tool, -f.risk_score, f.id), + "title": lambda f: (f.title.lower(), f.id), +} + + +def sort_findings(items: list[Finding], sort_by: str) -> list[Finding]: + if sort_by == "detected_at": + return sorted(items, key=lambda f: (f.detected_at, f.id), reverse=True) + return sorted(items, key=_SORT_KEYS.get(sort_by, _SORT_KEYS["title"])) + + +def finding_group_keys(ds: Dataset, f: Finding, group_by: str) -> list[str]: + """The findings_summary groups a finding counts towards.""" + if group_by == "severity": + return [f.severity.value] + if group_by == "source_tool": + return [f.source_tool] + if group_by == "framework": + return list(f.compliance_frameworks or ["(none)"]) + aid = ds.finding_asset_id(f) + a = ds.by_id.get(aid) if aid else None + if a is None: + return ["unmapped"] + if group_by == "asset_type": + return [a.asset_type.value] + if group_by == "account": + return [a.account_id or "unknown"] + return [a.name] + + +def summarise_findings(ds: Dataset, group_by: str, include_suppressed: bool) -> dict[str, Any]: + """Group counts with a severity breakdown per group.""" + groups: dict[str, dict[str, int]] = {} + for f in ds.findings: + if f.is_suppressed and not include_suppressed: + continue + for key in finding_group_keys(ds, f, group_by): + g = groups.setdefault(key, {"total": 0}) + g["total"] += 1 + g[f.severity.value] = g.get(f.severity.value, 0) + 1 + return groups + + +def _risk_row(ds: Dataset, aid: str, open_f: list[Finding], reachable: set[str]) -> dict: + a = ds.by_id[aid] + base = max(f.risk_score for f in open_f) + exposed = a.is_internet_exposed or aid in reachable + exposure = 1.5 if exposed else 1.0 + blast = len(ds.dependency_graph().dependents(aid, max_depth=6)) + blast_f = 1.0 + min(1.0, math.log10(1 + blast) / 2) + volume = 1.0 + min(0.5, 0.05 * (len(open_f) - 1)) + return { + "asset": asset_brief(ds, a), + "score": round(base * exposure * blast_f * volume, 2), + "components": { + "max_finding_risk": base, + "exposure_factor": exposure, + "internet_reachable": exposed, + "transitive_dependents": blast, + "blast_factor": round(blast_f, 3), + "open_findings": len(open_f), + "volume_factor": round(volume, 3), + }, + "top_findings": [ + finding_brief(ds, f) for f in sorted(open_f, key=lambda f: -f.risk_score)[:3] + ], + } + + +def risk_ranking(ds: Dataset, min_severity: str = "") -> list[dict[str, Any]]: + """Assets ranked by worst finding x exposure x blast radius x volume.""" + min_rank = SEVERITY_RANK[parse_severity(min_severity).value] if min_severity else -1 + reachable = ds.internet_reachable() + rows = [] + for aid, fs in ds.findings_by_asset.items(): + open_f = [ + f for f in fs if not f.is_suppressed and SEVERITY_RANK[f.severity.value] >= min_rank + ] + if open_f and aid in ds.by_id: + rows.append(_risk_row(ds, aid, open_f, reachable)) + rows.sort(key=lambda r: (-r["score"], r["asset"]["name"])) + return rows + + +def _merge_one(n: ComplianceResult, c: ComplianceResult, alive: list[str]) -> ComplianceResult: + ids = list(dict.fromkeys([*n.finding_ids, *alive])) + title = n.control_title + if c.control_title and (not title or title.startswith(f"{c.framework} ")): + title = c.control_title + n = n.model_copy(update={"finding_ids": ids, "control_title": title}) + if ids: + n = n.model_copy(update={"status": ComplianceStatus.FAIL}) + return n + + +def merge_compliance( + old: list[ComplianceResult], new: list[ComplianceResult], finding_ids: set[str] +) -> list[ComplianceResult]: + """Combine loaded compliance results with a fresh normalisation. + + Controls the normaliser evaluated take its result, plus any surviving + finding ids the loaded result already mapped to them; controls it did + not evaluate (scanner-native mappings, PASS results, other tools) are + kept with their finding ids filtered to findings that still exist. A + kept result whose findings all disappeared is dropped.""" + fresh = {(c.framework, c.control_id): c for c in new} + out: list[ComplianceResult] = [] + for c in old: + key = (c.framework, c.control_id) + alive = [i for i in c.finding_ids if i in finding_ids] + if key in fresh: + fresh[key] = _merge_one(fresh[key], c, alive) + elif alive or not c.finding_ids: + out.append(c.model_copy(update={"finding_ids": alive})) + return out + list(fresh.values()) + + +def renormalise(ws: Any, ds: Dataset, *extra: list[Finding]) -> None: + """Run the normaliser over the dataset's findings (plus ``extra``), + keeping suppression flags, and store the result. Existing compliance + results are merged with the new ones (see :func:`merge_compliance`).""" + from cloudg.normaliser import FindingsNormaliser + + rules = getattr(getattr(ws.config, "rulesets", None), "rules_dir", None) + suppressed = {f.id for f in ds.findings if f.is_suppressed} + sr = FindingsNormaliser(rules_dir=rules).normalise(ds.findings, *extra, assets=ds.assets) + for f in sr.findings: + if f.id in suppressed: + f.is_suppressed = True + merged = merge_compliance(ds.compliance, sr.compliance, {f.id for f in sr.findings}) + ds.replace_findings(sr.findings, merged) + + +def check_report_tools(reports: dict[str, list[str]]) -> None: + """Refuse tools cloudg has no report parser for.""" + from cloudg.ingest import SUPPORTED_TOOLS + + bad = [t for t in reports if t.lower() not in SUPPORTED_TOOLS] + if bad: + raise InvalidArgumentsError( + f"Unsupported tool(s) {', '.join(bad)}. Supported: {', '.join(SUPPORTED_TOOLS)}" + ) + + +def parse_reports(ws: Any, reports: dict[str, list[str]]) -> tuple[list, list, list[Finding]]: + """Parse every (tool, path) of ``reports``: (per_path, errors, findings).""" + from cloudg.ingest import parse_report + + check_report_tools(reports) + per_path: list[dict[str, Any]] = [] + errors: list[dict[str, Any]] = [] + new: list[Finding] = [] + for tool, paths in reports.items(): + for p in paths: + safe = ws.check_path(p, must_exist=False) + try: + if tool.lower() == "prowler" and safe.exists(): + check_prowler_input(safe) + fs = parse_report(tool, safe) + except (ValueError, FileNotFoundError, InvalidArgumentsError) as exc: + msg = exc.message if isinstance(exc, InvalidArgumentsError) else str(exc) + errors.append({"tool": tool, "path": p, "error": msg}) + continue + per_path.append({"tool": tool, "path": p, "findings": len(fs)}) + new.extend(fs) + return per_path, errors, new diff --git a/cloudg/mcp/catalog/_graph.py b/cloudg/mcp/catalog/_graph.py new file mode 100644 index 0000000..13001cb --- /dev/null +++ b/cloudg/mcp/catalog/_graph.py @@ -0,0 +1,480 @@ +"""Helpers behind the graph tools: node resolution (including the virtual +internet source), path description, exposure evidence, edge filters, the +breadth-first neighbourhood walk and the attack / lateral path searches. + +Port and internet-source parsing is delegated to :mod:`cloudg.graph.ports` +so the MCP tools read rules exactly like the reachability analysis: Azure +``*`` ports and ``Internet`` / ``Any`` / ``*`` sources, GCP rules with no +ports (every port of the protocol), AWS ICMP type / code pairs (no ports) +and egress rules (never an internet entry). +""" + +from __future__ import annotations + +from collections import deque +from dataclasses import dataclass +from typing import Any, Iterable + +import networkx as nx + +from cloudg.graph.ports import edge_port_ranges, is_egress, is_internet_source, port_in_ranges +from cloudg.mcp.catalog._common import asset_brief, node_brief, parse_enums +from cloudg.mcp.core import InvalidArgumentsError +from cloudg.mcp.state import Dataset +from cloudg.schema.models import AssetType, EdgeType, NetworkEdge + +CROWN_JEWELS = { + AssetType.RDS_INSTANCE: 10, + AssetType.AURORA_CLUSTER: 10, + AssetType.AZURE_SQL: 10, + AssetType.CLOUD_SQL: 10, + AssetType.DYNAMODB_TABLE: 9, + AssetType.DATA_WAREHOUSE: 9, + AssetType.SECRET: 9, + AssetType.KEY_VAULT: 9, + AssetType.KMS_KEY: 8, + AssetType.S3_BUCKET: 8, + AssetType.BLOB_STORAGE: 8, + AssetType.GCS_BUCKET: 8, + AssetType.FILE_SYSTEM: 7, + AssetType.CACHE_CLUSTER: 6, + AssetType.SEARCH_DOMAIN: 7, + AssetType.IAM_ROLE: 5, + AssetType.IAM_USER: 5, + AssetType.ACCESS_KEY: 6, +} +IDENTITY_EDGES = { + EdgeType.IAM_TRUST.value, + EdgeType.ASSUMES_ROLE.value, + EdgeType.GRANTS_ACCESS.value, +} +SENSITIVE_PORTS = {22, 3389, 3306, 5432, 1433, 27017, 6379, 9200, 5601, 8080, 8443} + +CATEGORY = "graph" +#: Tool options shared by every read-only graph tool. +GRAPH_RO = dict(read_only=True, idempotent=True, open_world=False, category=CATEGORY) + +# Virtual node standing for "the internet" when a tool is asked for +# source='internet': linked to every internet entry point of the dataset. +INTERNET = "__internet__" +_INTERNET_ALIASES = ("internet", "0.0.0.0/0", "::/0", "public") +_COMPACT_KEYS = ( + "id", + "name", + "type", + "account_id", + "internet_exposed", + "open_findings", + "max_severity", + "external", +) + + +# --------------------------------------------------------------------------- +# Edges +# --------------------------------------------------------------------------- + + +def edge_type_values(values: list[str] | None) -> set[str]: + """The EdgeType values named in a tool's ``edge_types`` argument.""" + return {t.value for t in parse_enums(values, EdgeType, "edge type")} + + +def edge_data(e: NetworkEdge) -> dict[str, Any]: + """The attributes :mod:`cloudg.graph.ports` reads from a graph edge.""" + return { + "edge_type": e.edge_type.value, + "port_range": e.port_range, + "protocol": e.protocol, + "cidr": e.cidr, + "direction": e.direction, + } + + +def allows_port(e: NetworkEdge, port: int) -> bool: + """Whether a rule edge opens ``port`` (explicit port list or port range).""" + if port in e.ports: + return True + return port_in_ranges(port, edge_port_ranges(edge_data(e))) + + +def open_ports(e: NetworkEdge, wanted: Iterable[int]) -> set[int]: + """The ports of ``wanted`` a rule edge opens.""" + ranges = edge_port_ranges(edge_data(e)) + return {p for p in wanted if p in e.ports or port_in_ranges(p, ranges)} + + +def from_internet(e: NetworkEdge) -> bool: + """An ingress edge whose source or CIDR stands for the whole internet.""" + if is_egress(edge_data(e)): + return False + return is_internet_source(e.source_id) or is_internet_source(e.cidr) + + +# --------------------------------------------------------------------------- +# Nodes +# --------------------------------------------------------------------------- + + +def internet_alias(ref: str) -> bool: + return ref.strip().lower() in _INTERNET_ALIASES + + +def internet_nodes(ds: Dataset) -> list[str]: + """Internet entry points of the dataset (same rule as the reachability + analysis), sorted with the 0.0.0.0/0 and ::/0 nodes first.""" + + def build() -> list[str]: + from cloudg.graph.reachability import ReachabilityAnalyzer + + found = ReachabilityAnalyzer(ds.graph).internet_entry_points() + return sorted(found, key=lambda n: (n not in ("0.0.0.0/0", "::/0"), n)) + + return ds.cached("internet_nodes", build) + + +def resolve_node(ds: Dataset, ref: str, *, virtual_internet: bool = False) -> str: + """Graph node for a reference: an asset, a placeholder node, or for + 'internet' the first internet entry point (or :data:`INTERNET` when + ``virtual_internet``).""" + if internet_alias(ref): + nodes = internet_nodes(ds) + if not nodes: + raise InvalidArgumentsError( + "This dataset has no internet node (no internet-sourced ingress rules or " + "INTERNET_EXPOSED edges). Use internet_exposure to see assets flagged as exposed." + ) + return INTERNET if virtual_internet else nodes[0] + if ref in ds.graph and ref not in ds.by_id: + return ref # placeholder node (CIDR / external id) + return ds.resolve_asset(ref).id + + +def brief(ds: Dataset, node: str) -> dict[str, Any]: + if node == INTERNET: + return {"id": INTERNET, "name": "internet", "type": "EXTERNAL", "external": True} + return node_brief(ds, node) + + +def compact(b: dict[str, Any]) -> dict[str, Any]: + """Path hops repeat a lot: keep only what explains the hop.""" + return {k: b[k] for k in _COMPACT_KEYS if b.get(k) is not None} + + +# --------------------------------------------------------------------------- +# Paths +# --------------------------------------------------------------------------- + + +def path_graph(ds: Dataset, mode: str) -> nx.DiGraph | nx.Graph: + if mode == "flow": + return ds.flow_graph + if mode == "undirected": + return ds.cached("undirected", lambda: ds.graph.to_undirected(as_view=False)) + return ds.graph + + +def with_internet(ds: Dataset, g: Any) -> Any: + """Copy of ``g`` with :data:`INTERNET` linked to every entry point.""" + g = g.copy() + g.add_node(INTERNET) + for n in internet_nodes(ds): + if n in g: + g.add_edge(INTERNET, n, edge_type="INTERNET") + return g + + +def hop_type(data: dict[str, Any]) -> str: + if data.get("derived") == "sg_admits" or data.get("reversed"): + return "SG_ADMITS" + return data.get("edge_type", "?") + + +def describe_path(ds: Dataset, g: Any, path: list[str]) -> dict[str, Any]: + edge_types = [hop_type(g.get_edge_data(u, v) or {}) for u, v in zip(path, path[1:])] + briefs = [brief(ds, n) for n in path] + return { + "length": len(path) - 1, + "nodes": [compact(b) for b in briefs], + "edge_types": edge_types, + "summary": " -> ".join(b["name"] for b in briefs), + } + + +def simple_paths(g: Any, s: str, t: str, max_depth: int, max_paths: int) -> list[list[str]]: + """Shortest simple paths from ``s`` to ``t`` (one beyond ``max_paths`` so + the caller can tell more exist). Paths from :data:`INTERNET` are + returned without the virtual first node.""" + skip = 1 if s == INTERNET else 0 + paths: list[list[str]] = [] + try: + for p in nx.shortest_simple_paths(g, s, t): + if len(p) - 1 - skip > max_depth or len(paths) > max_paths: + break + paths.append(p[skip:]) + except (nx.NetworkXNoPath, nx.NodeNotFound): + return paths # no (further) path: what was found so far is the answer + return paths + + +@dataclass +class AttackSearch: + """Inputs of the attack-path search.""" + + ds: Dataset + targets: dict[str, int] + max_depth: int + + +def attack_graph(ds: Dataset) -> nx.DiGraph: + """The flow graph plus a virtual entry linked to every internet entry + point and every asset flagged internet-exposed.""" + g = with_internet(ds, ds.flow_graph) + for a in ds.assets: + if a.is_internet_exposed: + g.add_edge(INTERNET, a.id, edge_type="INTERNET_EXPOSED") + return g + + +def attack_paths_found(search: AttackSearch, g: nx.DiGraph) -> list[dict[str, Any]]: + ds = search.ds + _, paths = nx.single_source_dijkstra(g, INTERNET, cutoff=search.max_depth + 1) + found = [] + for tid, weight in search.targets.items(): + p = paths.get(tid, [INTERNET])[1:] # drop the virtual entry + if not p: + continue + on_path = sum(len(ds.open_findings(n)) for n in p if n in ds.by_id) + first = g.get_edge_data(INTERNET, p[0]) or {} + desc = describe_path(ds, g, p) + desc["entry"] = ( + "internet-exposed asset" + if first.get("edge_type") == "INTERNET_EXPOSED" + else "internet-sourced rule" + ) + desc["target_weight"] = weight + desc["findings_on_path"] = on_path + desc["score"] = round(weight * 1.0 + min(on_path, 10) * 0.5 - (len(p) - 1) * 0.3, 2) + found.append(desc) + found.sort(key=lambda d: (-d["score"], d["length"])) + return found + + +def identity_graph(ds: Dataset) -> nx.DiGraph: + def build() -> nx.DiGraph: + g = nx.DiGraph() + for e in ds.edges: + if e.edge_type.value in IDENTITY_EDGES: + g.add_edge( + e.source_id, + e.target_id, + edge_type=e.edge_type.value, + relationship=e.relationship, + ) + return g + + return ds.cached("identity_graph", build) + + +def lateral_starts(ds: Dataset) -> list[str]: + starts = [a.id for a in ds.assets if a.is_internet_exposed] + starts += [ + a.id + for a in ds.assets + if a.asset_type == AssetType.CLOUD_ACCOUNT and a.metadata.get("external") + ] + return starts + + +def lateral_chains(ds: Dataset, g: nx.DiGraph, starts: list[str], max_depth: int) -> list: + chains: list[list[str]] = [] + for s in dict.fromkeys(starts): + if s not in g: + continue + for tgt, p in nx.single_source_shortest_path(g, s, cutoff=max_depth).items(): + if tgt != s and g.out_degree(tgt) == 0: + chains.append(p) + + def weight(p: list[str]) -> int: + return max( + (CROWN_JEWELS.get(ds.by_id[n].asset_type, 0) for n in p if n in ds.by_id), default=0 + ) + + chains.sort(key=lambda p: (-weight(p), len(p))) + return chains + + +# --------------------------------------------------------------------------- +# Exposure +# --------------------------------------------------------------------------- + + +def _filter_nodes(ds: Dataset, asset_id: str) -> tuple[list[tuple[str, str | None]], bool]: + """The asset plus the security groups / NSGs attached to it, and + whether a PROTECTS edge (WAF) points at it.""" + via_nodes: list[tuple[str, str | None]] = [(asset_id, None)] + protected = False + for e, other, d in ds.adjacency.get(asset_id, []): + if d == "in" and e.edge_type == EdgeType.PROTECTS: + protected = True + a = ds.by_id.get(other) + if ( + d == "out" + and e.edge_type == EdgeType.ATTACHED_TO + and a is not None + and a.asset_type in (AssetType.SECURITY_GROUP, AssetType.NSG) + ): + via_nodes.append((other, a.name)) + return via_nodes, protected + + +def exposure_evidence(ds: Dataset, asset_id: str) -> dict[str, Any]: + """Internet-sourced ingress rules reaching an asset directly or through + a security group / NSG attached to it, sensitive ports among them, WAF + coverage.""" + via_nodes, protected = _filter_nodes(ds, asset_id) + rules = [] + ports: set[int] = set() + for node, via in via_nodes: + for e, other, d in ds.adjacency.get(node, []): + if d != "in" or is_egress(edge_data(e)): + continue + if not (is_internet_source(e.cidr) or is_internet_source(other)): + continue + rules.append( + { + "edge_type": e.edge_type.value, + "ports": e.port_range or ",".join(str(p) for p in e.ports[:10]) or "all", + "protocol": e.protocol, + "via": via or e.description, + } + ) + ports |= open_ports(e, SENSITIVE_PORTS) + return { + "internet_rules": rules[:10], + "sensitive_ports_open": sorted(ports), + "protected_by_waf": protected, + } + + +# --------------------------------------------------------------------------- +# Neighbourhood walk and dependency helpers +# --------------------------------------------------------------------------- + + +@dataclass +class Walk: + """Options of the breadth-first neighbourhood walk.""" + + direction: str = "both" + edge_types: set[str] | None = None + depth: int = 1 + cap: int = 100 + + +def _hops(ds: Dataset, node: str, walk: Walk) -> Iterable[tuple[NetworkEdge, str]]: + for e, other, dirn in ds.adjacency.get(node, []): + if walk.direction != "both" and dirn != walk.direction: + continue + if walk.edge_types and e.edge_type.value not in walk.edge_types: + continue + yield e, other + + +def bfs(ds: Dataset, start: str, walk: Walk) -> tuple[list[str], list[Any], bool]: + """Nodes and edges within ``walk.depth`` hops of ``start``; at most + ``cap`` nodes and ``2 * cap`` edges.""" + seen = {start} + nodes: list[str] = [] + edges: list[Any] = [] + seen_edges: set[str] = set() + queue: deque[tuple[str, int]] = deque([(start, 0)]) + truncated = False + while queue: + node, d = queue.popleft() + if d >= walk.depth: + continue + for e, other in _hops(ds, node, walk): + if e.id not in seen_edges: + if len(edges) >= walk.cap * 2: + truncated = True + continue + seen_edges.add(e.id) + edges.append(e) + if other in seen: + continue + if len(nodes) >= walk.cap: + truncated = True + continue + seen.add(other) + nodes.append(other) + queue.append((other, d + 1)) + return nodes, edges, truncated + + +def dep_items(ds: Dataset, links: list[Any]) -> list[dict[str, Any]]: + out = [] + for link in links: + a = ds.by_id.get(link.asset_id) + if a is None: + continue + out.append( + { + **asset_brief(ds, a), + "via": link.via_edge, + "relationship": link.relationship, + "depth": link.depth, + "parent_id": link.parent_id, + } + ) + return out + + +def prune_tree(nodes: list[dict[str, Any]], budget: list[int]) -> list[dict[str, Any]]: + """Keep at most ``budget[0]`` nodes; ``budget[1]`` becomes 1 when a node + had to be dropped.""" + out = [] + for n in nodes: + if budget[0] <= 0: + budget[1] = 1 + break + budget[0] -= 1 + children = n.get("children") or [] + out.append({**n, "children": prune_tree(children, budget)}) + return out + + +# --------------------------------------------------------------------------- +# Edge search +# --------------------------------------------------------------------------- + + +@dataclass +class EdgeQuery: + """get_edges filters, with endpoints already resolved to node ids.""" + + source: str | None = None + target: str | None = None + edge_types: set[str] | None = None + relationship: str = "" + internet_only: bool = False + port: int | None = None + + def _endpoint_ok(self, e: NetworkEdge) -> bool: + if self.source == INTERNET: + if not from_internet(e): + return False + elif self.source and e.source_id != self.source: + return False + return not self.target or e.target_id == self.target + + def matches(self, e: NetworkEdge) -> bool: + if not self._endpoint_ok(e): + return False + if self.edge_types and e.edge_type.value not in self.edge_types: + return False + if self.relationship and (e.relationship or "") != self.relationship: + return False + if self.internet_only and not from_internet(e): + return False + return self.port is None or allows_port(e, self.port) diff --git a/cloudg/mcp/catalog/_graph_impact.py b/cloudg/mcp/catalog/_graph_impact.py new file mode 100644 index 0000000..9d14d2a --- /dev/null +++ b/cloudg/mcp/catalog/_graph_impact.py @@ -0,0 +1,400 @@ +"""Graph tools about impact and structure: blast radius, dependencies, +shared dependencies, cross-account edges, centrality, sub-graph export, +security-service coverage and graph statistics. Registered by +:func:`cloudg.mcp.catalog.graph.register` after the path and exposure +tools, so the tool order is unchanged.""" + +from __future__ import annotations + +from typing import Annotated, Any, Callable, Literal + +import networkx as nx +from pydantic import Field + +from cloudg.mcp.catalog._common import ( + DEFAULT_LIMIT, + Catalog, + Cursor, + DatasetArg, + DependencyOut, + Limit, + RefArg, + asset_brief, + asset_uri, + bounded, + fingerprint, + node_brief, + paginate, + parse_enums, + schema, + ws_dataset, +) +from cloudg.mcp.catalog._graph import ( + GRAPH_RO, + edge_type_values, + CROWN_JEWELS, + Walk, + bfs, + dep_items, + prune_tree, + resolve_node, +) +from cloudg.mcp.core import Sensitivity +from cloudg.mcp.state import Dataset +from cloudg.schema.models import AssetType + + +GraphFormat = Literal["d3", "cytoscape", "graphml"] +Metric = Literal["degree", "in_degree", "out_degree", "betweenness"] +DepthArg = Annotated[int, Field(ge=1, le=15)] +DepLimit = Annotated[int, Field(ge=1, le=500)] + +CATALOG = Catalog() + + +@CATALOG.tool(title="Blast radius", sensitivity=Sensitivity.CONFIDENTIAL, **GRAPH_RO) +def blast_radius( + ctx: Any, + ref: RefArg, + max_depth: DepthArg = 10, + limit: DepLimit = 50, + dataset: DatasetArg = "", +) -> dict: + """What is affected if this asset breaks, is deleted, changed or + compromised: (1) dependency blast radius, meaning everything that + transitively depends on it, with accounts and exposed assets + affected; (2) network / graph reach: everything reachable from it, + with the sensitive data stores among them and a 0-10 risk score.""" + ds = ws_dataset(ctx, dataset) + a = ds.resolve_asset(ref) + items = dep_items(ds, ds.dependency_graph().dependents(a.id, max_depth)) + net = ds.reachability.compute_blast_radius(a.id) + reach = sorted(n for n in net["reachable_nodes"] if n in ds.by_id) + sensitive = [n for n in reach if ds.by_id[n].asset_type in CROWN_JEWELS] + by_depth: dict[str, int] = {} + for i in items: + by_depth[str(i["depth"])] = by_depth.get(str(i["depth"]), 0) + 1 + ctx.link(asset_uri(a.id), a.name) + return { + "dataset": ds.name, + "asset": asset_brief(ds, a), + "dependency": { + "transitive_dependents": len(items), + "by_depth": by_depth, + "accounts_affected": sorted({i["account_id"] for i in items if i["account_id"]}), + "internet_exposed_dependents": sum(1 for i in items if i["internet_exposed"]), + "items": items[:limit], + "truncated": len(items) > limit, + }, + "network": { + "reachable_assets": len(reach), + "max_depth": net["depth"], + "risk_score": net["risk_score"], + "sensitive_reachable": [node_brief(ds, n) for n in sensitive[:limit]], + }, + } + + +def _dependency_result( + ctx: Any, ref: str, dataset: str, limit: int, walk: Callable[[Dataset, str], list] +) -> dict: + """Shared body of depends_on / dependents.""" + ds = ws_dataset(ctx, dataset) + a = ds.resolve_asset(ref) + page, env = bounded(dep_items(ds, walk(ds, a.id)), limit) + return {"dataset": ds.name, "asset": asset_brief(ds, a), **env, "items": page} + + +@CATALOG.tool( + title="Depends on", + sensitivity=Sensitivity.CONFIDENTIAL, + output_schema=schema(DependencyOut), + **GRAPH_RO, +) +def depends_on( + ctx: Any, + ref: RefArg, + max_depth: DepthArg = 3, + limit: DepLimit = 100, + dataset: DatasetArg = "", +) -> dict: + """Everything this asset needs (upstream), breadth-first: its role, + keys, image, subnet, log group, the queue that triggers it... Each + item says which edge it came through and at what depth.""" + return _dependency_result( + ctx, ref, dataset, limit, lambda ds, aid: ds.dependency_graph().depends_on(aid, max_depth) + ) + + +@CATALOG.tool( + title="Dependents", + sensitivity=Sensitivity.CONFIDENTIAL, + output_schema=schema(DependencyOut), + **GRAPH_RO, +) +def dependents( + ctx: Any, + ref: RefArg, + max_depth: DepthArg = 3, + limit: DepLimit = 100, + dataset: DatasetArg = "", +) -> dict: + """Everything that needs this asset (downstream), i.e. what breaks if it + is deleted or changed. Breadth-first, with the edge and depth of each.""" + return _dependency_result( + ctx, ref, dataset, limit, lambda ds, aid: ds.dependency_graph().dependents(aid, max_depth) + ) + + +@CATALOG.tool(title="Dependency tree", sensitivity=Sensitivity.CONFIDENTIAL, **GRAPH_RO) +def dependency_tree( + ctx: Any, + ref: RefArg, + direction: Annotated[ + Literal["up", "down", "both"], + Field(description="up = what it depends on, down = what depends on it."), + ] = "both", + max_depth: Annotated[int, Field(ge=1, le=6)] = 3, + include_hierarchy: Annotated[ + bool, Field(description="Include org / account containment edges.") + ] = False, + max_nodes: Annotated[int, Field(ge=1, le=1000)] = 200, + dataset: DatasetArg = "", +) -> dict: + """Nested upstream / downstream dependency tree of one asset (the + same view as `cloudg deps`), pruned to max_nodes.""" + ds = ws_dataset(ctx, dataset) + a = ds.resolve_asset(ref) + tree = ds.dependency_graph(include_hierarchy).tree(a.id, direction, max_depth) + budget = [max_nodes, 0] + out: dict[str, Any] = {"dataset": ds.name, "asset": tree["asset"]} + for key in ("depends_on", "dependents"): + if key in tree: + out[key] = prune_tree(tree[key], budget) + out["truncated"] = bool(budget[1]) + return out + + +@CATALOG.tool(title="Shared dependencies", sensitivity=Sensitivity.CONFIDENTIAL, **GRAPH_RO) +def shared_dependencies( + ctx: Any, top: Annotated[int, Field(ge=1, le=200)] = 25, dataset: DatasetArg = "" +) -> dict: + """Single points of failure: assets the most other assets directly + depend on (one KMS key behind forty resources, one role behind every + function), excluding pure containment.""" + ds = ws_dataset(ctx, dataset) + items = ds.dependency_graph().shared_dependencies(top) + return {"dataset": ds.name, "items": items, "returned": len(items)} + + +@CATALOG.tool(title="Largest blast radius", sensitivity=Sensitivity.CONFIDENTIAL, **GRAPH_RO) +def largest_blast_radius( + ctx: Any, top: Annotated[int, Field(ge=1, le=100)] = 15, dataset: DatasetArg = "" +) -> dict: + """Assets whose failure or change affects the most others + (transitive dependents), with accounts and internet-exposed assets + affected. Use blast_radius(ref) for the detail of one.""" + ds = ws_dataset(ctx, dataset) + items = ds.dependency_graph().blast_radius(top=top) + return {"dataset": ds.name, "items": items, "returned": len(items)} + + +@CATALOG.tool(title="Cross-account edges", sensitivity=Sensitivity.CONFIDENTIAL, **GRAPH_RO) +def cross_account_edges( + ctx: Any, + account_id: Annotated[str, Field(description="Only edges touching this account.")] = "", + external_only: Annotated[ + bool, Field(description="Only edges to accounts outside the mapped estate.") + ] = False, + edge_types: list[str] | None = None, + limit: Limit = DEFAULT_LIMIT, + cursor: Cursor = "", + dataset: DatasetArg = "", +) -> dict: + """Relationships that cross account / subscription / project + boundaries (trust, grants, peering, shared resources), flagging + external accounts. Summarised by account pair.""" + from cloudg.inventory.dependencies import cross_account_edges as _xa + + ds = ws_dataset(ctx, dataset) + types = edge_type_values(edge_types) + rows = [ + r + for r in ds.cached("cross_account", lambda: _xa(ds.assets, ds.edges)) + if (not account_id or account_id in (r["source_account"], r["target_account"])) + and (not external_only or r["external"]) + and (not types or r["edge_type"] in types) + ] + pairs: dict[str, int] = {} + for r in rows: + k = f"{r['source_account']} -> {r['target_account']}" + pairs[k] = pairs.get(k, 0) + 1 + page, env = paginate( + rows, limit, cursor, fingerprint(ds, a=account_id, x=external_only, t=sorted(types)) + ) + return {"dataset": ds.name, **env, "by_account_pair": pairs, "items": page} + + +def _centrality(g: nx.DiGraph, metric: str) -> dict[str, float]: + if metric == "degree": + return nx.degree_centrality(g) + if metric == "in_degree": + return nx.in_degree_centrality(g) + if metric == "out_degree": + return nx.out_degree_centrality(g) + n = g.number_of_nodes() + k = None if n <= 1500 else min(n, 300) + return nx.betweenness_centrality(g, k=k, seed=7) + + +@CATALOG.tool(title="Most central assets", sensitivity=Sensitivity.CONFIDENTIAL, **GRAPH_RO) +def centrality_top( + ctx: Any, + metric: Metric = "degree", + top: Annotated[int, Field(ge=1, le=200)] = 20, + asset_types: list[str] | None = None, + dataset: DatasetArg = "", +) -> dict: + """Most connected / most "in the middle" assets by graph centrality. + High betweenness marks chokepoints that many paths cross. They are good + places for controls and bad places for compromise. Betweenness is + sampled on large graphs.""" + ds = ws_dataset(ctx, dataset) + types = parse_enums(asset_types, AssetType, "asset type") + g = ds.graph + scores = ds.cached(f"centrality:{metric}", lambda: _centrality(g, metric)) + ranked = [ + (n, s) + for n, s in sorted(scores.items(), key=lambda kv: -kv[1]) + if n in ds.by_id and (not types or ds.by_id[n].asset_type in types) + ][:top] + return { + "dataset": ds.name, + "metric": metric, + "sampled": metric == "betweenness" and g.number_of_nodes() > 1500, + "items": [ + {**asset_brief(ds, ds.by_id[n]), "score": round(s, 5), "degree": g.degree(n)} + for n, s in ranked + ], + } + + +def _subgraph_nodes(ds: Dataset, seeds: list[str], walk: Walk) -> tuple[list[str], bool]: + keep: list[str] = [] + truncated = False + for ref in seeds: + sid = resolve_node(ds, ref) + nodes, _, tr = bfs(ds, sid, walk) + truncated = truncated or tr + for n in [sid, *nodes]: + if n not in keep: + keep.append(n) + if len(keep) > walk.cap: + return keep[: walk.cap], True + return keep, truncated + + +def _render_subgraph(ds: Dataset, keep: list[str], types: set[str] | None, fmt: str) -> Any: + from cloudg.graph.builder import GraphBuilder + + sub = GraphBuilder() + keep_set = set(keep) + sub.build( + [ds.by_id[n] for n in keep if n in ds.by_id], + [ + e + for e in ds.edges + if e.source_id in keep_set + and e.target_id in keep_set + and (types is None or e.edge_type.value in types) + ], + ) + if fmt == "d3": + return sub, sub.to_d3_json() + if fmt == "cytoscape": + return sub, sub.to_cytoscape_json() + return sub, "\n".join(nx.generate_graphml(sub.graph)) + + +# RESTRICTED: the GraphML form is one opaque string the privacy layer can +# only scan as free text, so strict ceilings hide the whole tool. +@CATALOG.tool(title="Export sub-graph", sensitivity=Sensitivity.RESTRICTED, **GRAPH_RO) +def subgraph_export( + ctx: Any, + seeds: Annotated[ + list[str], Field(min_length=1, max_length=50, description="Asset refs to start from.") + ], + depth: Annotated[int, Field(ge=0, le=4)] = 1, + format: GraphFormat = "d3", # pylint: disable=redefined-builtin # MCP argument name + max_nodes: Annotated[int, Field(ge=1, le=2000)] = 300, + edge_types: list[str] | None = None, + dataset: DatasetArg = "", +) -> dict: + """The neighbourhood of a few assets as a self-contained graph in + D3 ({nodes, links}), Cytoscape ({elements}) or GraphML (string), for + visualisation or handing to another tool. Use the + cloudg://graph/{format} resources for the whole graph.""" + ds = ws_dataset(ctx, dataset) + types = edge_type_values(edge_types) or None + keep, truncated = _subgraph_nodes(ds, seeds, Walk("both", types, depth, max_nodes)) + sub, data = _render_subgraph(ds, keep, types, format) + return { + "dataset": ds.name, + "format": format, + "nodes": sub.graph.number_of_nodes(), + "edges": sub.graph.number_of_edges(), + "truncated": truncated, + "graph": data, + } + + +@CATALOG.tool(title="Security service coverage", sensitivity=Sensitivity.CONFIDENTIAL, **GRAPH_RO) +def security_coverage( + ctx: Any, limit: Annotated[int, Field(ge=1, le=500)] = 50, dataset: DatasetArg = "" +) -> dict: + """Which security services (GuardDuty, Security Hub, Inspector, + Config, Access Analyzer, Macie...) are enabled per account / region, + the gaps, workloads without vulnerability scanning, and + internet-facing entry points without a WAF.""" + from cloudg.inventory.dependencies import security_coverage as _cov + + ds = ws_dataset(ctx, dataset) + cov = ds.cached("security_coverage", lambda: _cov(ds.assets, ds.edges)) + out: dict[str, Any] = { + "dataset": ds.name, + "services_by_account_region": cov["services_by_account_region"], + } + for key in ("gaps", "workloads_without_vulnerability_scanning", "internet_facing_without_waf"): + items, env = bounded(cov[key], limit) + out[key] = {**env, "items": items} + return out + + +@CATALOG.tool(title="Graph statistics", sensitivity=Sensitivity.INTERNAL, **GRAPH_RO) +def graph_stats(ctx: Any, dataset: DatasetArg = "") -> dict: + """Shape of the relationship graph: nodes (assets vs placeholders), + edges by type and relationship, connected components, the largest + component, isolated assets and density.""" + ds = ws_dataset(ctx, dataset) + g = ds.graph + n_nodes = g.number_of_nodes() + comps = sorted((len(c) for c in nx.weakly_connected_components(g)), reverse=True) + rels: dict[str, int] = {} + for e in ds.edges: + if e.relationship: + rels[e.relationship] = rels.get(e.relationship, 0) + 1 + return { + "dataset": ds.name, + "nodes": n_nodes, + "asset_nodes": len(ds.by_id), + "placeholder_nodes": n_nodes - sum(1 for n in g if n in ds.by_id), + "edges": len(ds.edges), + "graph_edges": g.number_of_edges(), + "edges_by_type": ds.summary()["edges_by_type"], + "edges_by_relationship": dict(sorted(rels.items(), key=lambda kv: -kv[1])[:30]), + "components": len(comps), + "largest_component": comps[0] if comps else 0, + "isolated_assets": sum(1 for n in ds.by_id if n in g and g.degree(n) == 0), + "density": round(nx.density(g), 6) if n_nodes > 1 else 0.0, + "mean_degree": round(sum(d for _, d in g.degree()) / n_nodes, 3) if n_nodes else 0.0, + } diff --git a/cloudg/mcp/catalog/_inventory.py b/cloudg/mcp/catalog/_inventory.py new file mode 100644 index 0000000..1adb934 --- /dev/null +++ b/cloudg/mcp/catalog/_inventory.py @@ -0,0 +1,292 @@ +"""Helpers behind the inventory tools: the asset filter, sort and group +keys, metadata truncation and the organization topology views.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from typing import Any, Callable, Sequence + +from cloudg.inventory.mapper_result import _service_of +from cloudg.mcp.catalog._common import ArgsFilter, parse_enum, parse_enums, parse_severity +from cloudg.mcp.state import SEVERITY_RANK, Dataset +from cloudg.schema.models import AssetType, CloudAsset, CloudProvider + +_HIERARCHY = {AssetType.ORGANIZATION, AssetType.ORG_UNIT, AssetType.CLOUD_ACCOUNT} + +Predicate = Callable[[CloudAsset], bool] + + +@dataclass(frozen=True) +class AssetFilter(ArgsFilter): + """The filters find_assets / count_assets accept (empty = no filter).""" + + query: str = "" + provider: str = "" + asset_types: Sequence[str] | None = None + region: str = "" + account_id: str = "" + tag: str = "" + internet_exposed: bool | None = None + has_findings: bool | None = None + min_severity: str = "" + + +def _tag_predicate(tag: str) -> Predicate: + key, has_val, val = tag.partition("=") + if has_val: + return lambda a: key in a.tags and a.tags[key] == val + return lambda a: key in a.tags + + +def _query_predicate(query: str) -> Predicate: + q = query.lower().strip() + + def match(a: CloudAsset) -> bool: + return ( + q in a.name.lower() + or q in (a.arn or "").lower() + or q in a.id.lower() + or any(q in str(v).lower() for v in a.tags.values()) + ) + + return match + + +def _finding_predicate(ds: Dataset, flt: AssetFilter) -> Predicate | None: + min_rank = SEVERITY_RANK[parse_severity(flt.min_severity).value] if flt.min_severity else None + if flt.has_findings is None and min_rank is None: + return None + + def match(a: CloudAsset) -> bool: + open_f = ds.open_findings(a.id) + if flt.has_findings is not None and bool(open_f) != flt.has_findings: + return False + return min_rank is None or any(SEVERITY_RANK[f.severity.value] >= min_rank for f in open_f) + + return match + + +def _predicates(ds: Dataset, flt: AssetFilter) -> list[Predicate]: + preds: list[Predicate] = [] + if flt.provider: + prov = parse_enum(flt.provider, CloudProvider, "provider") + preds.append(lambda a: a.provider == prov) + types = parse_enums(flt.asset_types, AssetType, "asset type") + if types: + preds.append(lambda a: a.asset_type in types) + if flt.region: + region = flt.region.lower() + preds.append(lambda a: a.region.lower() == region) + if flt.account_id: + preds.append(lambda a: (a.account_id or "") == flt.account_id) + if flt.tag: + preds.append(_tag_predicate(flt.tag)) + if flt.internet_exposed is not None: + preds.append(lambda a: a.is_internet_exposed == flt.internet_exposed) + by_findings = _finding_predicate(ds, flt) + if by_findings is not None: + preds.append(by_findings) + if flt.query.strip(): + preds.append(_query_predicate(flt.query)) + return preds + + +def filter_assets(ds: Dataset, flt: AssetFilter | None = None) -> list[CloudAsset]: + """Assets of ``ds`` matching every filter of ``flt``.""" + preds = _predicates(ds, flt or AssetFilter()) + return [a for a in ds.assets if all(p(a) for p in preds)] + + +def risk_key(ds: Dataset, a: CloudAsset) -> tuple: + open_f = ds.open_findings(a.id) + best = max((f.risk_score for f in open_f), default=0.0) + return (best + (2.0 if a.is_internet_exposed else 0.0), len(open_f)) + + +def sort_assets(ds: Dataset, assets: list[CloudAsset], sort_by: str, descending: bool) -> None: + """Sort in place; 'risk' and 'findings' put the highest first unless + ``descending`` flips them.""" + keys: dict[str, Callable[[CloudAsset], Any]] = { + "name": lambda a: (a.name.lower(), a.id), + "type": lambda a: (a.asset_type.value, a.name.lower()), + "region": lambda a: (a.region, a.name.lower()), + "account": lambda a: (a.account_id or "", a.name.lower()), + "risk": lambda a: risk_key(ds, a), + "findings": lambda a: (len(ds.open_findings(a.id)), a.name.lower()), + } + reverse = descending if sort_by not in ("risk", "findings") else not descending + assets.sort(key=keys[sort_by], reverse=reverse) + + +def _worst_severity(ds: Dataset, a: CloudAsset) -> str: + fs = ds.open_findings(a.id) + return max((f.severity.value for f in fs), key=lambda s: SEVERITY_RANK[s], default="NONE") + + +def group_key(ds: Dataset, a: CloudAsset, group_by: str) -> str: + """The count_assets group an asset falls in.""" + keys: dict[str, Callable[[CloudAsset], str]] = { + "type": lambda x: x.asset_type.value, + "provider": lambda x: x.provider.value, + "region": lambda x: x.region, + "account": lambda x: x.account_id or "unknown", + "service": _service_of, + "exposure": lambda x: "internet_exposed" if x.is_internet_exposed else "internal", + } + fn = keys.get(group_by) + return fn(a) if fn else _worst_severity(ds, a) + + +def relation_view(ds: Dataset, a: CloudAsset, max_relations: int) -> dict[str, Any]: + """get_asset's relation counts by edge type plus the first edges.""" + from cloudg.mcp.catalog._common import edge_brief + + rel: dict[str, Any] = {"outgoing": {}, "incoming": {}} + out_edges = [e for e in ds.edges if e.source_id == a.id] + in_edges = [e for e in ds.edges if e.target_id == a.id] + for side, edges in (("outgoing", out_edges), ("incoming", in_edges)): + for e in edges: + rel[side][e.edge_type.value] = rel[side].get(e.edge_type.value, 0) + 1 + rel["sample"] = [edge_brief(ds, e) for e in (out_edges + in_edges)[:max_relations]] + rel["total"] = len(out_edges) + len(in_edges) + rel["truncated"] = rel["total"] > max_relations + return rel + + +def truncate_metadata(meta: dict[str, Any], max_chars: int) -> dict[str, Any]: + """Fit ``meta`` into ``max_chars`` of JSON, largest keys dropped first.""" + if len(json.dumps(meta, default=str)) <= max_chars: + return {"truncated": False, "metadata": meta} + sizes = {k: len(json.dumps(v, default=str)) for k, v in meta.items()} + kept, total = {}, 0 + for k, v in meta.items(): + if total + sizes[k] > max_chars: + continue + kept[k] = v + total += sizes[k] + return { + "truncated": True, + "key_sizes": dict(sorted(sizes.items(), key=lambda kv: -kv[1])[:50]), + "hint": "Metadata exceeds max_chars; request specific keys.", + "metadata": kept, + } + + +def account_rows(ds: Dataset) -> list[dict[str, Any]]: + """One row per account / subscription / project, most assets first.""" + accts: dict[str, dict[str, Any]] = {} + names: dict[str, tuple[str, bool]] = {} + for a in ds.assets: + if a.asset_type == AssetType.CLOUD_ACCOUNT and a.account_id: + names[a.account_id] = (a.name, bool(a.metadata.get("external"))) + acct = a.account_id or "unknown" + e = accts.setdefault( + acct, + { + "account_id": acct, + "provider": a.provider.value, + "assets": 0, + "regions": set(), + "internet_exposed": 0, + "open_findings": 0, + }, + ) + e["assets"] += 1 + e["regions"].add(a.region) + e["internet_exposed"] += int(a.is_internet_exposed) + e["open_findings"] += len(ds.open_findings(a.id)) + items = [] + for acct, e in sorted(accts.items(), key=lambda kv: -kv[1]["assets"]): + name, external = names.get(acct, (None, False)) + items.append({**e, "regions": sorted(e["regions"]), "name": name, "external": external}) + return items + + +# --------------------------------------------------------------------------- +# Organization topology +# --------------------------------------------------------------------------- + + +def _pick(d: dict[str, Any], keys: Sequence[str]) -> dict[str, Any]: + return {k: d.get(k) for k in keys} + + +def org_from_discovery( + ds: Dataset, org: dict[str, Any], include_policies: bool, max_accounts: int +) -> dict[str, Any]: + """Topology from organization discovery (``Dataset.organization``).""" + accounts = list((org.get("accounts") or {}).values()) + ous = list((org.get("ous") or {}).values()) + out: dict[str, Any] = { + "dataset": ds.name, + "source": "organization_discovery", + "organization_id": org.get("organization_id"), + "management_account_id": org.get("management_account_id"), + "feature_set": org.get("feature_set"), + "control_tower_enabled": bool(org.get("control_tower_enabled", org.get("landing_zone"))), + "governed_regions": org.get("governed_regions", []), + "shared_accounts": org.get("shared_accounts", {}), + "ous": [_pick(o, ("id", "name", "path", "parent_id")) for o in ous], + "accounts": [ + _pick(a, ("id", "name", "status", "ou_path", "parent_id")) + for a in accounts[:max_accounts] + ], + "total_accounts": len(accounts), + "truncated": len(accounts) > max_accounts, + } + if include_policies: + out["policies"] = [ + {**_pick(p, ("id", "name", "type", "aws_managed")), "targets": p.get("targets", [])} + for p in org.get("policies", []) + ] + return out + + +def _org_policies(ds: Dataset) -> list[dict[str, Any]]: + pols = [a for a in ds.assets if a.asset_type == AssetType.ORG_POLICY] + pol_ids = {p.id for p in pols} + targets: dict[str, list[str]] = {} + for e in ds.edges: + if e.edge_type.value == "GOVERNS" and e.source_id in pol_ids: + targets.setdefault(e.source_id, []).append(e.target_id) + return [{"id": p.id, "name": p.name, "targets": targets.get(p.id, [])} for p in pols] + + +def org_from_assets(ds: Dataset, include_policies: bool, max_accounts: int) -> dict[str, Any]: + """Topology rebuilt from the hierarchy assets on the map.""" + hier = [a for a in ds.assets if a.asset_type in _HIERARCHY] + if not hier: + return { + "dataset": ds.name, + "source": None, + "note": "No organization data in this dataset. Map with organization discovery " + "enabled to get it.", + } + ids = {a.id for a in hier} + parents = { + e.target_id: e.source_id + for e in ds.edges + if e.edge_type.value == "CONTAINS" and e.source_id in ids and e.target_id in ids + } + nodes = [ + { + "id": a.id, + "name": a.name, + "type": a.asset_type.value, + "account_id": a.account_id, + "parent_id": parents.get(a.id), + "external": bool(a.metadata.get("external")), + } + for a in hier + ] + out = { + "dataset": ds.name, + "source": "map_assets", + "nodes": nodes[:max_accounts], + "total": len(nodes), + "truncated": len(nodes) > max_accounts, + } + if include_policies: + out["policies"] = _org_policies(ds) + return out diff --git a/cloudg/mcp/catalog/_live.py b/cloudg/mcp/catalog/_live.py new file mode 100644 index 0000000..c137546 --- /dev/null +++ b/cloudg/mcp/catalog/_live.py @@ -0,0 +1,445 @@ +"""Helpers behind the live tools: engine and config set-up, credential +checks, progress wiring, landing results, the live-operation guard and +the scanner job.""" + +from __future__ import annotations + +import asyncio +import time +from dataclasses import dataclass +from typing import Any, Awaitable, Callable + +from cloudg.mcp.catalog._common import SUMMARY_CORE_KEYS +from cloudg.mcp.core import AccessDeniedError, InvalidArgumentsError, RateLimitedError +from cloudg.mcp.state import Dataset +from cloudg.resilience import stats_scope + +PROVIDERS = ("aws", "azure", "gcp") +SCANNERS = ("prowler", "scoutsuite", "checkov", "trivy", "iam") + +PHASES = { + "inventory_mapping": 1, + "collection": 1, + "scanning": 2, + "ingest": 2, + "analysis": 3, + "normalisation": 4, + "reporting": 5, +} + + +def _engine(cfg: Any) -> Any: + from cloudg import api + + return api.CloudGEngine(cfg) + + +def _config(ctx: Any, providers: list[str] | None, regions: list[str] | None, **inv: Any) -> Any: + cfg = ctx.workspace.config.model_copy(deep=True) + if providers: + bad = [p for p in providers if p.lower() not in PROVIDERS] + if bad: + raise InvalidArgumentsError( + f"Unknown provider(s) {', '.join(bad)}. Use aws, azure, gcp." + ) + cfg.providers = [p.lower() for p in providers] + if regions: + for p in cfg.providers: + sub = getattr(cfg, p, None) + if sub is not None and hasattr(sub, "regions"): + sub.regions = list(regions) + for k, v in inv.items(): + if v is not None: + setattr(cfg.inventory, k, v) + return cfg + + +def _check_aws(cfg: Any) -> str | None: + import os + + try: + import boto3 + from botocore.exceptions import BotoCoreError + except ImportError: + return "boto3 is not installed (pip install 'cloudg[aws]')" + if cfg.access_key_id and cfg.secret_access_key: + return None + token_file = cfg.web_identity_token_file or os.environ.get("AWS_WEB_IDENTITY_TOKEN_FILE") + if cfg.role_arn and token_file: + return ( + None + if os.path.isfile(token_file) + else (f"web identity token file {token_file} does not exist") + ) + try: + session = boto3.Session(profile_name=cfg.profile) if cfg.profile else boto3.Session() + creds = session.get_credentials() + except BotoCoreError as exc: + return str(exc) + if creds is None: + return ( + "no credentials found (set aws.profile / access keys in the cloudg config, " + "AWS_PROFILE / AWS_ACCESS_KEY_ID in the server's environment, or run on a " + "host with an instance role)" + ) + return None + + +def _check_azure(cfg: Any) -> str | None: + try: + from cloudg.credentials import build_azure_credential + + build_azure_credential(cfg).get_token("https://management.azure.com/.default") + except ImportError: + return "the Azure SDK is not installed (pip install 'cloudg[azure]')" + except Exception as exc: # azure.identity raises CredentialUnavailableError & co + return f"{type(exc).__name__}: {str(exc).splitlines()[0][:300]}" + return None + + +def _check_gcp(cfg: Any) -> str | None: + try: + from cloudg.credentials import build_gcp_credentials + + build_gcp_credentials(cfg) + except ImportError: + return "google-auth is not installed (pip install 'cloudg[gcp]')" + except Exception as exc: # google.auth.exceptions.DefaultCredentialsError + return f"{type(exc).__name__}: {str(exc).splitlines()[0][:300]}" + return None + + +_CHECKS = {"aws": _check_aws, "azure": _check_azure, "gcp": _check_gcp} + + +def _hook_progress(ctx: Any, engine: Any, total: float) -> list[str]: + """Wire the engine's phase hook to MCP progress notifications.""" + loop = asyncio.get_running_loop() + phases: list[str] = [] + + def on_phase(phase: str) -> None: + phases.append(phase) + step = min(float(PHASES.get(phase, len(phases))), total - 0.5) + asyncio.run_coroutine_threadsafe(ctx.report_progress(step, total, f"phase: {phase}"), loop) + + errors: list[str] = [] + + def on_error(phase: str, exc: Exception) -> None: + errors.append(f"{phase}: {exc}") + + engine.on_phase_start = on_phase + engine.on_error = on_error + engine._mcp_errors = errors # surfaced in the tool result + return phases + + +async def _done(ctx: Any, total: float) -> None: + """Let queued phase notifications go out, then report completion + (MCP progress must increase, so 'done' has to be last).""" + for _ in range(3): + await asyncio.sleep(0) + await ctx.report_progress(total, total, "done") + + +def _land( + ctx: Any, ds: Dataset, tool: str, started: float, extra: dict[str, Any], replace: bool +) -> dict: + """Add a live result to the workspace and build the tool result. Runs + inside the guarded operation, so callers joined by single-flight get + this same result (and dataset) instead of a duplicate.""" + ctx.workspace.add(ds, activate=True, replace=replace) + summary = ds.summary() + out = { + "dataset": ds.name, + "active": True, + "tool": tool, + "duration_s": round(time.monotonic() - started, 1), + "summary": { + k: summary[k] + for k in ( + "total_assets", + "total_edges", + "total_findings", + *SUMMARY_CORE_KEYS, + "assets_by_type", + "providers", + ) + }, + **extra, + "next_steps": [ + "dataset_summary", + "find_assets", + "top_risks", + "attack_paths", + "export_report(format='inventory') to persist it", + ], + } + throttling = ds.metadata.get("throttling") + if throttling: + out["throttling"] = _throttling_summary(throttling) + return out + + +def _run_throttling(result: Any, stats: Any) -> dict[str, Any] | None: + """Throttling of one live run: the result's own report when it has one + (InventoryResult.throttling), else the run's bound resilience stats.""" + throttling = getattr(result, "throttling", None) + if throttling: + return throttling + if stats is not None and stats.eventful: + from cloudg.resilience import get_governor + + return get_governor().summary(stats) + return None + + +def _throttling_summary(throttling: dict[str, Any]) -> dict[str, Any]: + return { + "totals": throttling.get("totals", {}), + "messages": list(throttling.get("messages", []))[:20], + "skipped": throttling.get("skipped", {}), + } + + +def _links(ctx: Any, out: dict[str, Any]) -> dict[str, Any]: + ctx.link( + f"cloudg://datasets/{out['dataset']}/summary", + f"{out['dataset']} summary", + title="Dataset summary", + priority=1.0, + ) + ctx.link("cloudg://workspace", "workspace") + return out + + +def _name(ctx: Any, name: str, prefix: str, replace: bool) -> str: + """Validate (or generate) the result dataset's name before any cloud call.""" + if name: + return ctx.workspace.check_name(name, replace=replace) + return ctx.workspace.unique_name(f"{prefix}-{time.strftime('%Y%m%d-%H%M%S')}") + + +# --------------------------------------------------------------------------- +# Live operation guard (cloudg.resilience): single-flight, cooldowns, caps +# --------------------------------------------------------------------------- + +FORCE_ROLES = {"admin", "operator"} + + +class LiveOperationRefused(RateLimitedError): + """The live operation guard refused the call (cooldown, busy or caller + quota). ``data`` carries the guard's reason, retry_after_seconds and + scopes, plus the dataset to use instead.""" + + +def _aws_scope_ids(sub: Any) -> list[str]: + return list(sub.accounts or []) or ([f"profile:{sub.profile}"] if sub.profile else []) + + +def _azure_scope_ids(sub: Any) -> list[str]: + return list(sub.subscription_ids or []) + + +def _gcp_scope_ids(sub: Any) -> list[str]: + return list(sub.project_ids or []) or ( + [f"org:{sub.organization_id}"] if sub.organization_id else [] + ) + + +_SCOPE_IDS = {"aws": _aws_scope_ids, "azure": _azure_scope_ids, "gcp": _gcp_scope_ids} + + +def live_scopes(cfg: Any) -> list[str]: + """Guard scopes for the providers of ``cfg``: provider plus account / + profile / subscription / project when the config names them, otherwise + the bare provider.""" + scopes: list[str] = [] + for provider in cfg.providers: + sub = getattr(cfg, provider, None) + ids_of = _SCOPE_IDS.get(provider) + ids = ids_of(sub) if ids_of is not None and sub is not None else [] + scopes += [f"{provider}/{i}" for i in ids] or [provider] + return sorted(set(scopes)) + + +def _latest_live_dataset(ctx: Any, providers: Any = None) -> str | None: + """Newest live dataset that covers every provider in ``providers`` (all + live datasets when none are given); None rather than a dataset of the + wrong providers.""" + wanted = {str(p).lower() for p in providers or []} + live = [ + d + for d in ctx.workspace.datasets() + if (d.source == "live" or d.kind == "live") + and wanted <= {str(p).lower() for p in d.providers or []} + ] + live.sort(key=lambda d: d.loaded_at) + return live[-1].name if live else None + + +def _bypass(ctx: Any, force: bool) -> bool: + if not force: + return False + if not FORCE_ROLES & set(ctx.principal.roles): + raise AccessDeniedError( + "force=true (skip the live cooldown) needs the admin or operator role. Use the " + "dataset from the last collection, or wait for the cooldown to end." + ) + return True + + +async def _guarded( + ctx: Any, op: str, cfg: Any, parts: tuple, scopes: list[str], force: bool, factory: Any +) -> dict[str, Any]: + """Run ``factory`` under the workspace's live guard and map a refusal to + a readable tool error. The returned dict says whether this call joined + an identical operation that was already running.""" + from cloudg.resilience import LiveOperationRejected, operation_key + + guard = ctx.workspace.live_guard + token = object() + key = operation_key(op, sorted(cfg.providers), scopes, *parts) + + async def run() -> tuple[object, dict[str, Any]]: + return token, await factory() + + try: + owner, out = await guard.run( + key, run, scopes=scopes, caller=ctx.principal.id, bypass_cooldown=_bypass(ctx, force) + ) + except LiveOperationRejected as exc: + data = exc.to_dict() + latest = _latest_live_dataset(ctx, getattr(cfg, "providers", None)) + msg = exc.message.rstrip(".") + "." + if latest: + data["use_dataset"] = latest + msg += ( + f" The last live result is already loaded as dataset {latest!r}; query it " + "instead of collecting again." + ) + else: + msg += " Query an already loaded dataset (list_datasets) in the meantime." + if data.get("reason") == "cooldown": + msg += " An admin or operator can pass force=true to skip the cooldown." + raise LiveOperationRefused(msg, data=data) from None + out = dict(out) + if owner is not token: + out["joined"] = True + out["note"] = ( + "An identical live operation was already running; this call shares its " + "result instead of collecting again." + ) + return _links(ctx, out) + + +@dataclass +class ScanJob: + """Everything run_scanners hands to the guarded scan operation.""" + + ds: Dataset + cfg: Any + iac_dir: Any + images: list[str] | None + out_dir: Any + normalise: bool + started: float + + +def _scanner_config(ctx: Any, scanners: list[str] | None) -> Any: + if scanners: + bad = [s for s in scanners if s.lower() not in SCANNERS] + if bad: + raise InvalidArgumentsError( + f"Unknown scanner(s) {', '.join(bad)}. Use: {', '.join(SCANNERS)}" + ) + cfg = _config(ctx, None, None) + if scanners: + cfg.scanners.enabled = [s.lower() for s in scanners] + return cfg + + +def _check_images(images: list[str] | None) -> None: + """Refuse image references that the scanner would read as options.""" + bad = [i for i in images or [] if not i.strip() or i.lstrip().startswith("-")] + if bad: + raise InvalidArgumentsError( + "Container image references must be non-empty and must not start with '-' " + "(they are passed to trivy on its command line).", + data={"value": bad}, + ) + + +async def _run_scan(ctx: Any, job: ScanJob) -> dict[str, Any]: + ds, cfg = job.ds, job.cfg + engine = _engine(cfg) + _hook_progress(ctx, engine, 3.0) + with stats_scope() as stats: + findings = await engine.scan( + ds.assets, + ds.edges, + iac_dir=str(job.iac_dir) if job.iac_dir else None, + images=job.images, + output_dir=job.out_dir, + ) + before = len(ds.findings) + if job.normalise: + from cloudg.mcp.catalog.findings import renormalise + + await asyncio.to_thread(renormalise, ctx.workspace, ds, findings) + else: + ds.add_findings(findings) + ctx.workspace.mutated(ds) + await _done(ctx, 3) + sev: dict[str, int] = {} + for f in findings: + sev[f.severity.value] = sev.get(f.severity.value, 0) + 1 + throttling = _run_throttling(None, stats) + extra = {"throttling": _throttling_summary(throttling)} if throttling else {} + return { + **extra, + "dataset": ds.name, + "scanners": cfg.scanners.enabled, + "new_findings": len(findings), + "severity_breakdown": sev, + "findings_before": before, + "findings_after": len(ds.findings), + "output_dir": str(job.out_dir), + "duration_s": round(time.monotonic() - job.started, 1), + "errors": getattr(engine, "_mcp_errors", []), + "next_steps": ["top_risks", "list_findings(min_severity='HIGH')", "compliance_summary"], + } + + +@dataclass +class LiveRun: + """One guarded live collection: how to call the engine, turn its + result into a dataset, and what to add to the tool result.""" + + tool: str + total: float + started: float + replace: bool + call: Callable[[Any], Awaitable[Any]] + build: Callable[[Any], Dataset] + extra_fields: Callable[[Any, Any, Dataset], dict[str, Any]] + note: str = "" + + +async def run_live(ctx: Any, cfg: Any, run: LiveRun) -> dict[str, Any]: + """Run the engine for ``run`` and land the result as a dataset.""" + engine = _engine(cfg) + _hook_progress(ctx, engine, run.total) + if run.note: + await ctx.report_progress(0.5, run.total, run.note) + with stats_scope() as stats: + result = await run.call(engine) + ds = run.build(result) + throttling = _run_throttling(result, stats) + if throttling: + ds.metadata["throttling"] = throttling + await _done(ctx, run.total) + return _land(ctx, ds, run.tool, run.started, run.extra_fields(result, engine, ds), run.replace) + + +def engine_errors(engine: Any) -> list[str]: + return getattr(engine, "_mcp_errors", []) diff --git a/cloudg/mcp/catalog/_resources.py b/cloudg/mcp/catalog/_resources.py new file mode 100644 index 0000000..28aab22 --- /dev/null +++ b/cloudg/mcp/catalog/_resources.py @@ -0,0 +1,161 @@ +"""Completions and content builders shared by the resources and prompts.""" + +from __future__ import annotations + +from typing import Any + +from cloudg.mcp.catalog._common import asset_brief, finding_brief, max_severity +from cloudg.mcp.state import SEVERITY_RANK, Dataset, ReferenceNotFoundError +from cloudg.schema.models import AssetType, Severity + +GRAPH_FORMATS = ("d3", "cytoscape", "graphml") +ONTOLOGY_FORMATS = { + "turtle": "text/turtle", + "json-ld": "application/ld+json", + "xml": "application/rdf+xml", + "nt": "application/n-triples", +} + + +# --------------------------------------------------------------------------- +# Completions (fn(ctx, partial) -> list[str]) +# --------------------------------------------------------------------------- + + +def _ctx_dataset(ctx: Any) -> Dataset | None: + name = (getattr(ctx, "meta", {}) or {}).get("arguments", {}).get("dataset") or None + try: + return ctx.workspace.get(name) + except Exception: + return None + + +def complete_dataset(ctx: Any, partial: str = "") -> list[str]: + return [n for n in ctx.workspace.names() if n.lower().startswith(partial.lower())] + + +def complete_asset(ctx: Any, partial: str = "") -> list[str]: + ds = _ctx_dataset(ctx) + if ds is None: + return [] + p = partial.lower() + names: dict[str, int] = {} + for a in ds.assets: + names[a.name] = names.get(a.name, 0) + 1 + starts, contains = [], [] + for a in ds.assets: + # Unique names complete to the name; shared names to the ARN / id + label = a.name if names[a.name] == 1 else (a.arn or a.id) + keys = [k.lower() for k in (a.name, a.arn or "", a.id)] + if not p or any(k.startswith(p) for k in keys): + starts.append(label) + elif any(p in k for k in keys): + contains.append(label) + # The full list: the layer caps what it sends and computes hasMore + return list(dict.fromkeys(sorted(starts) + sorted(contains))) + + +def complete_finding(ctx: Any, partial: str = "") -> list[str]: + ds = _ctx_dataset(ctx) + if ds is None: + return [] + ranked = sorted(ds.findings, key=lambda f: -f.risk_score) + return [f.id for f in ranked if f.id.startswith(partial)] + + +def complete_severity(ctx: Any, partial: str = "") -> list[str]: + return [s.value for s in Severity if s.value.startswith(partial.upper())] + + +def complete_framework(ctx: Any, partial: str = "") -> list[str]: + names: set[str] = set() + ds = _ctx_dataset(ctx) + if ds is not None: + names |= {c.framework for c in ds.compliance} + names |= {fw for f in ds.findings for fw in f.compliance_frameworks} + if not names: + from cloudg.mcp.catalog.compliance import _rules_dir, ruleset_catalog + + names = set(ruleset_catalog(_rules_dir(ctx))) + return sorted(n for n in names if partial.lower() in n.lower()) + + +def complete_asset_type(ctx: Any, partial: str = "") -> list[str]: + return [t.value for t in AssetType if t.value.startswith(partial.upper())] + + +def complete_graph_format(ctx: Any, partial: str = "") -> list[str]: + return [f for f in GRAPH_FORMATS if f.startswith(partial.lower())] + + +def complete_ontology_format(ctx: Any, partial: str = "") -> list[str]: + return [f for f in ONTOLOGY_FORMATS if f.startswith(partial.lower())] + + +# --------------------------------------------------------------------------- +# Content builders shared with prompts +# --------------------------------------------------------------------------- + + +def asset_detail(ds: Dataset, ref: str) -> dict[str, Any]: + if ref in ds.graph and ref not in ds.by_id: + return { + "id": ref, + "external": True, + "name": ds.graph.nodes[ref].get("name", ref), + "degree": ds.graph.degree(ref), + } + a = ds.resolve_asset(ref) + out = asset_brief(ds, a, tags=True) + rel: dict[str, dict[str, int]] = {"outgoing": {}, "incoming": {}} + for e, _, d in ds.adjacency.get(a.id, []): + side = rel["outgoing" if d == "out" else "incoming"] + side[e.edge_type.value] = side.get(e.edge_type.value, 0) + 1 + out["relations"] = rel + out["findings"] = [ + finding_brief(ds, f) + for f in sorted(ds.open_findings(a.id), key=lambda f: -f.risk_score)[:10] + ] + out["metadata_keys"] = sorted(a.metadata)[:100] + needs, needed_by = ds.dependency_graph().direct_counts(a.id) + out["dependencies"] = {"direct_depends_on": needs, "direct_dependents": needed_by} + return out + + +def finding_detail(ds: Dataset, f: Any) -> dict[str, Any]: + """Content of ``cloudg://findings/{finding_id}``.""" + return { + **finding_brief(ds, f), + "description": f.description, + "evidence": f.evidence, + "remediation": f.remediation, + } + + +def framework_detail(ds: Dataset, framework: str) -> dict[str, Any]: + """Content of ``cloudg://compliance/{framework}``.""" + from cloudg.mcp.catalog._compliance import framework_rollup, open_control_findings + + rows = [ + r for r in framework_rollup(ds, framework) if r["framework"].lower() == framework.lower() + ] or framework_rollup(ds, framework) + if not rows: + raise ReferenceNotFoundError( + "No compliance data for that framework.", data={"value": framework} + ) + fw_names = {r["framework"] for r in rows} + failing = [] + for c in ds.compliance: + if c.framework in fw_names and c.status.value == "FAIL": + open_f = open_control_findings(ds, c) + failing.append( + { + "framework": c.framework, + "control_id": c.control_id, + "control_title": c.control_title, + "open_findings": len(open_f), + "max_severity": max_severity(open_f), + } + ) + failing.sort(key=lambda r: (-SEVERITY_RANK.get(r["max_severity"] or "", -1), r["control_id"])) + return {"dataset": ds.name, "posture": rows, "failing_controls": failing[:300]} diff --git a/cloudg/mcp/catalog/_sparql.py b/cloudg/mcp/catalog/_sparql.py new file mode 100644 index 0000000..2c09486 --- /dev/null +++ b/cloudg/mcp/catalog/_sparql.py @@ -0,0 +1,247 @@ +"""Read-only SPARQL for the ontology tools: query vetting and bounded +result collection. + +Vetting refuses SPARQL UPDATE, FROM / FROM NAMED and SERVICE before the +query reaches rdflib. The UPDATE check is a linear scan of the query's +prologue (no backtracking regular expression), so a hostile query string +cannot stall it. + +Result collection is bounded: rows are read one at a time and the read +stops at ``limit`` rows or when the deadline passes. rdflib evaluates +SELECT patterns lazily, so a query whose pattern multiplies out (a +cartesian product of unrelated triple patterns) stops after the first +``limit`` rows instead of enumerating every combination. Clauses that +need the whole solution set first (ORDER BY, GROUP BY / aggregates, +DISTINCT over a large set) and CONSTRUCT / DESCRIBE, which rdflib builds +eagerly, are only bounded by the tool's ``timeout_seconds``. +""" + +from __future__ import annotations + +import re +import time +from typing import Any, Callable, Iterable + +from cloudg.mcp.core import InvalidArgumentsError + +QUERY_DEADLINE_S = 60.0 + +_FORBIDDEN_NODES = {"ServiceGraphPattern"} +_COMMENT = re.compile(r"#[^\n]*") +_UPDATE_KEYWORDS = frozenset( + {"INSERT", "DELETE", "LOAD", "CLEAR", "DROP", "CREATE", "ADD", "MOVE", "COPY", "WITH"} +) +_PROLOGUE_KEYWORDS = frozenset({"PREFIX", "BASE"}) + + +def namespaces() -> dict[str, Any]: + from rdflib.namespace import OWL, RDF, RDFS, XSD + + from cloudg.graph.ontology import CM, CMP, CMR + + return {"cm": CM, "cmr": CMR, "cmp": CMP, "rdf": RDF, "rdfs": RDFS, "owl": OWL, "xsd": XSD} + + +def compact(term: Any) -> str: + s = str(term) + for prefix, ns in namespaces().items(): + if s.startswith(str(ns)): + return f"{prefix}:{s[len(str(ns)) :]}" + return s + + +def _skip_space(text: str, i: int) -> int: + while i < len(text) and text[i].isspace(): + i += 1 + return i + + +def _word_end(text: str, i: int) -> int: + while i < len(text) and text[i].isalpha(): + i += 1 + return i + + +def first_keyword(text: str) -> str: + """The first keyword after the PREFIX / BASE prologue, upper-cased. + + Each PREFIX / BASE declaration is skipped up to the ``>`` that closes + its IRI. The scan only moves forward, so it is linear in the length of + ``text``.""" + i = 0 + while True: + i = _skip_space(text, i) + end = _word_end(text, i) + word = text[i:end].upper() + if word not in _PROLOGUE_KEYWORDS: + return word + close = text.find(">", end) + if close < 0: + return "" + i = close + 1 + + +def _walk(node: Any) -> Any: + from rdflib.plugins.sparql.parserutils import CompValue + + if isinstance(node, CompValue): + yield node.name + for v in node.values(): + yield from _walk(v) + elif isinstance(node, (list, tuple)): + for v in node: + yield from _walk(v) + + +def _literal_end(query: str, start: int, unterminated: set[str]) -> int: + """End of the string literal opening at ``start``, or -1 when none + does. ``unterminated`` remembers the quote forms already known to have + no closing quote further on, so no stretch of the query is scanned + twice for the same form.""" + quote = query[start] + triple = quote * 3 + if query.startswith(triple, start) and triple not in unterminated: + close = query.find(triple, start + 3) + if close >= 0: + return close + 3 + unterminated.add(triple) + if quote in unterminated: + return -1 + pos = start + 1 + while pos < len(query): + char = query[pos] + if char == "\\": + pos += 2 + elif char == quote: + return pos + 1 + else: + pos += 1 + unterminated.add(quote) + return -1 + + +def _blank_strings(query: str) -> str: + """``query`` with every SPARQL string literal (``"..."``, ``'...'`` and + their triple-quoted forms) replaced by ``""``, in linear time. A regular + expression doing the same backtracks quadratically on a query full of + unterminated quotes.""" + out: list[str] = [] + pos = 0 + unterminated: set[str] = set() + while pos < len(query): + end = _literal_end(query, pos, unterminated) if query[pos] in "\"'" else -1 + if end < 0: + out.append(query[pos]) + pos += 1 + else: + out.append('""') + pos = end + return "".join(out) + + +def prepare_readonly(query: str) -> Any: + """Parse a SPARQL query and refuse anything that is not a local, + read-only SELECT / ASK / CONSTRUCT / DESCRIBE.""" + from rdflib.plugins.sparql import prepareQuery + + stripped = _COMMENT.sub(" ", _blank_strings(query)) + if first_keyword(stripped) in _UPDATE_KEYWORDS: + raise InvalidArgumentsError( + "Only read-only queries are allowed (SELECT, ASK, " + "CONSTRUCT, DESCRIBE); SPARQL UPDATE is rejected." + ) + try: + prepared = prepareQuery(query, initNs=namespaces()) + except Exception as exc: # rdflib raises pyparsing / ValueError / custom errors + msg = str(exc).splitlines()[0][:300] + raise InvalidArgumentsError( + f"SPARQL parse error: {msg}. Only SELECT / ASK / CONSTRUCT / DESCRIBE are allowed; " + "prefixes cm:, cmr:, cmp:, rdf:, rdfs:, owl:, xsd: are predefined." + ) from None + algebra = prepared.algebra + if algebra.get("datasetClause"): + raise InvalidArgumentsError( + "FROM / FROM NAMED clauses are not allowed (they would load " + "external data). Query the dataset's graph directly." + ) + if _FORBIDDEN_NODES & set(_walk(algebra)): + raise InvalidArgumentsError( + "SERVICE (federated queries to remote endpoints) is not allowed." + ) + return prepared + + +def collect( + rows: Iterable[Any], + limit: int, + convert: Callable[[Any], dict[str, Any]], + deadline: float, +) -> tuple[list[dict[str, Any]], bool, bool]: + """Read at most ``limit`` converted rows before ``deadline`` + (``time.monotonic()``). Returns (rows, truncated, timed_out).""" + out: list[dict[str, Any]] = [] + for row in rows: + if len(out) >= limit: + return out, True, False + if time.monotonic() > deadline: + return out, True, True + out.append(convert(row)) + return out, False, False + + +def _kind(val: Any) -> str: + from rdflib import BNode, Literal + + if isinstance(val, Literal): + return "literal" + if isinstance(val, BNode): + return "bnode" + return "uri" + + +def run_query(g: Any, prepared: Any, limit: int, compact_uris: bool) -> dict[str, Any]: + """Execute a vetted query with the row cap and deadline described in + the module docstring.""" + res = g.query(prepared) + qtype = prepared.algebra.name + if qtype == "AskQuery": + return { + "query_type": qtype, + "answer": bool(res.askAnswer), + "rows": [], + "returned": 0, + "truncated": False, + } + fmt = compact if compact_uris else str + deadline = time.monotonic() + QUERY_DEADLINE_S + kinds: dict[str, str] = {} + if qtype == "SelectQuery": + names = [str(v) for v in res.vars] + + def convert(row: Any) -> dict[str, Any]: + for name, val in zip(names, row): + if val is not None: + kinds.setdefault(name, _kind(val)) + return {n: (fmt(v) if v is not None else None) for n, v in zip(names, row)} + + rows, truncated, timed_out = collect(res, limit, convert, deadline) + else: # CONSTRUCT / DESCRIBE -> triples + kinds = {"subject": "uri", "predicate": "uri", "object": "term"} + rows, truncated, timed_out = collect( + res, + limit, + lambda t: {"subject": fmt(t[0]), "predicate": fmt(t[1]), "object": fmt(t[2])}, + deadline, + ) + out: dict[str, Any] = { + "query_type": qtype, + "rows": rows, + "returned": len(rows), + "truncated": truncated, + "variable_kinds": kinds, + } + if timed_out: + out["hint"] = f"Stopped after {QUERY_DEADLINE_S:.0f}s. Narrow the patterns or add LIMIT." + elif truncated: + out["hint"] = "More rows exist: add LIMIT / OFFSET to the query or raise limit." + return out diff --git a/cloudg/mcp/catalog/compliance.py b/cloudg/mcp/catalog/compliance.py new file mode 100644 index 0000000..c31e282 --- /dev/null +++ b/cloudg/mcp/catalog/compliance.py @@ -0,0 +1,212 @@ +"""Compliance tools: framework catalog, per-framework posture, controls, +one control's status, and the assets behind compliance gaps. + +Two sources are combined. The rulesets shipped with cloudg in ``cloudg/rules`` +(CIS, NIST 800-53, PCI DSS, ISO 27001, SOC 2, HIPAA, GDPR, MITRE +ATT&CK...) describe the controls; the dataset's ``ComplianceResult`` list +and each finding's ``compliance_frameworks`` say which controls fail and +on what. +""" + +from __future__ import annotations + +from typing import Annotated, Any, Literal + +from pydantic import Field + +from cloudg.mcp.catalog._common import ( + DEFAULT_LIMIT, + Catalog, + ComplianceSummaryOut, + Cursor, + DatasetArg, + Limit, + asset_brief, + bounded, + fingerprint, + paginate, + parse_severity, + schema, + ws_dataset, +) +from cloudg.mcp.catalog._compliance import ( + briefs, + control_definition, + control_findings, + control_results, + control_state, + dataset_control_rows, + framework_rollup, + gap_rows, + missing_control_error, + ruleset_catalog, + ruleset_control_rows, +) +from cloudg.mcp.catalog._compliance import fw_match as _fw_match +from cloudg.mcp.catalog._compliance import rules_dir as _rules_dir +from cloudg.mcp.core import Registry, Sensitivity +from cloudg.mcp.state import SEVERITY_RANK + +__all__ = ["CATALOG", "_fw_match", "_rules_dir", "framework_rollup", "register", "ruleset_catalog"] + +CATEGORY = "compliance" + +ControlStatus = Literal["", "PASS", "FAIL", "NOT_APPLICABLE", "MANUAL"] + + +CATALOG = Catalog() + +_RO = dict(read_only=True, idempotent=True, open_world=False, category=CATEGORY) + + +@CATALOG.tool( + title="List compliance frameworks", sensitivity=Sensitivity.PUBLIC, capabilities=(), **_RO +) +def list_frameworks( + ctx: Any, + provider: Annotated[str, Field(description="aws, azure or gcp.")] = "", +) -> dict: + """Compliance frameworks cloudg can map findings to (from its + shipped rulesets): name, versions, providers and control counts. + Needs no dataset.""" + cat = ruleset_catalog(_rules_dir(ctx)) + items = [] + for fw, e in sorted(cat.items()): + if ( + provider + and provider.lower() not in {p.lower() for p in e["providers"]} + and e["providers"] + ): + continue + items.append( + { + "framework": fw, + "versions": sorted(e["versions"]), + "providers": sorted(e["providers"]), + "controls": len(e["controls"]), + "files": e["files"], + } + ) + return {"total": len(items), "items": items} + + +@CATALOG.tool( + title="Compliance summary", + sensitivity=Sensitivity.INTERNAL, + output_schema=schema(ComplianceSummaryOut), + **_RO, +) +def compliance_summary( + ctx: Any, + framework: Annotated[ + str, Field(description="Framework name or substring (e.g. 'CIS', 'PCI'); empty = all.") + ] = "", + dataset: DatasetArg = "", +) -> dict: + """Compliance posture per framework: controls evaluated / failing / + passing, pass rate, open findings by severity and affected assets. + Drill down with list_controls / control_status / compliance_gaps.""" + ds = ws_dataset(ctx, dataset) + rows = framework_rollup(ds, framework) + out: dict[str, Any] = {"dataset": ds.name, "frameworks": rows} + if not rows: + out["hint"] = ( + "No compliance data matched. Findings need compliance_frameworks (run " + "normalise_findings after ingesting scanner reports); list_frameworks shows " + "the framework names." + ) + for r in rows[:5]: + ctx.link(r["uri"], f"{r['framework']} compliance") + return out + + +@CATALOG.tool(title="List controls", sensitivity=Sensitivity.INTERNAL, **_RO) +def list_controls( + ctx: Any, + framework: Annotated[str, Field(min_length=1, description="Framework name or substring.")], + status: ControlStatus = "", + source: Annotated[ + Literal["dataset", "ruleset"], + Field( + description="dataset = controls evaluated for this dataset; ruleset = every " + "control the framework defines." + ), + ] = "dataset", + limit: Limit = DEFAULT_LIMIT, + cursor: Cursor = "", + dataset: DatasetArg = "", +) -> dict: + """Controls of a framework: either the ones evaluated against the + dataset (with status and finding counts) or the full ruleset + definition (ids, titles, mapped scanner checks).""" + if source == "ruleset": + matches, rows = ruleset_control_rows(ruleset_catalog(_rules_dir(ctx)), framework) + page, env = paginate(rows, limit, cursor, fingerprint(None, f=framework, s="rs")) + return {"source": "ruleset", "frameworks": matches, **env, "items": page} + ds = ws_dataset(ctx, dataset) + rows = dataset_control_rows(ds, framework, status) + page, env = paginate(rows, limit, cursor, fingerprint(ds, f=framework, s=status)) + return {"dataset": ds.name, "source": "dataset", **env, "items": page} + + +@CATALOG.tool(title="Control status", sensitivity=Sensitivity.CONFIDENTIAL, **_RO) +def control_status( + ctx: Any, + framework: Annotated[str, Field(min_length=1)], + control_id: Annotated[str, Field(min_length=1)], + limit: Annotated[int, Field(ge=1, le=200)] = 25, + dataset: DatasetArg = "", +) -> dict: + """One control: its definition (from the ruleset), status in the + dataset, the findings that fail it and the affected assets.""" + ds = ws_dataset(ctx, dataset) + results = control_results(ds, framework, control_id) + definition = control_definition(ruleset_catalog(_rules_dir(ctx)), framework, control_id) + if not results and definition is None: + raise missing_control_error(ds, framework, control_id) + open_f = control_findings(ds, results) + assets = {ds.finding_asset_id(f) for f in open_f} - {None} + page, env = bounded(open_f, limit) + return { + "dataset": ds.name, + "framework": framework, + "control_id": control_id, + "definition": definition, + "status": control_state(results, open_f), + "results": [ + { + "framework": c.framework, + "control_id": c.control_id, + "control_title": c.control_title, + "status": c.status.value, + } + for c in results + ], + "open_findings": env["total"], + "findings": briefs(ds, page), + "affected_assets": [asset_brief(ds, ds.by_id[a]) for a in sorted(assets)][:limit], + "truncated": env["truncated"], + } + + +@CATALOG.tool(title="Compliance gaps", sensitivity=Sensitivity.CONFIDENTIAL, **_RO) +def compliance_gaps( + ctx: Any, + framework: Annotated[str, Field(description="Framework name or substring; empty = all.")] = "", + min_severity: str = "", + limit: Limit = DEFAULT_LIMIT, + cursor: Cursor = "", + dataset: DatasetArg = "", +) -> dict: + """Assets that fail compliance, worst first: per asset the frameworks + and controls it fails, open findings, max severity and exposure.""" + ds = ws_dataset(ctx, dataset) + min_rank = SEVERITY_RANK[parse_severity(min_severity).value] if min_severity else -1 + rows = gap_rows(ds, framework, min_rank) + page, env = paginate(rows, limit, cursor, fingerprint(ds, f=framework, m=min_severity)) + return {"dataset": ds.name, "framework": framework or None, **env, "items": page} + + +def register(reg: Registry) -> None: + """Add this module's tools to ``reg``.""" + CATALOG.register(reg) diff --git a/cloudg/mcp/catalog/export.py b/cloudg/mcp/catalog/export.py new file mode 100644 index 0000000..fafedc8 --- /dev/null +++ b/cloudg/mcp/catalog/export.py @@ -0,0 +1,180 @@ +"""Export tools: Terraform recreation preview / export and report files. + +Every file is written under the workspace output directory (see +``workspace_status.output_dir``); paths that would escape it are refused. +""" + +from __future__ import annotations + +from typing import Annotated, Any, Literal + +from pydantic import Field + +from cloudg.mcp.catalog._common import Catalog, DatasetArg, parse_enums, ws_dataset +from cloudg.mcp.core import Capability, Registry, Sensitivity +from cloudg.mcp.state import Dataset +from cloudg.schema.models import AssetType, CloudAsset, ScanResult + +CATEGORY = "export" +R, FS = Capability.READ_STATE, Capability.WRITE_FS + +ReportFormat = Literal["json", "html", "inventory", "graphml", "asset_map"] +ReportFormatArg = Annotated[ + ReportFormat, + Field( + description=( + "json: findings.json (assets, findings, compliance, D3 graph); html: interactive " + "report.html; inventory: inventory-map.json + .graphml + graph + dependencies; " + "graphml: the graph only; asset_map: asset-map.json + compliance-map.json " + "(inventory x findings)." + ) + ), +] +SubDir = Annotated[ + str, + Field( + pattern=r"^[A-Za-z0-9_.\-/]{0,200}$", + description="Sub-directory of the workspace output directory ('' = the directory itself).", + ), +] + + +def _select(ds: Dataset, refs: list[str] | None, asset_types: list[str] | None) -> list[CloudAsset]: + types = parse_enums(asset_types, AssetType, "asset type") + if refs: + chosen = [ds.resolve_asset(r) for r in refs] + else: + chosen = list(ds.assets) + return [a for a in chosen if not types or a.asset_type in types] + + +def _scan_result(ds: Dataset) -> ScanResult: + return ScanResult( + assets=ds.assets, edges=ds.edges, findings=ds.findings, compliance=ds.compliance + ) + + +CATALOG = Catalog() + + +@CATALOG.tool( + title="Terraform preview", + category=CATEGORY, + sensitivity=Sensitivity.INTERNAL, + read_only=True, + idempotent=True, + open_world=False, +) +def terraform_preview( + ctx: Any, + refs: Annotated[list[str] | None, Field(description="Only these assets (refs).")] = None, + asset_types: list[str] | None = None, + dataset: DatasetArg = "", +) -> dict: + """What a Terraform recreation of the assets would contain: resource + counts per Terraform type and the asset types that have no Terraform + mapping. Writes nothing.""" + from cloudg.renderers.terraform_export import TerraformExporter + + ds = ws_dataset(ctx, dataset) + assets = _select(ds, refs, asset_types) + return { + "dataset": ds.name, + "assets_considered": len(assets), + **TerraformExporter().preview(assets), + } + + +@CATALOG.tool( + title="Export Terraform", + category=CATEGORY, + sensitivity=Sensitivity.CONFIDENTIAL, + capabilities={R, FS}, + read_only=False, + destructive=False, + idempotent=True, + open_world=False, +) +def export_terraform( + ctx: Any, + subdir: SubDir = "terraform", + refs: Annotated[list[str] | None, Field(description="Only these assets (refs).")] = None, + asset_types: list[str] | None = None, + dataset: DatasetArg = "", +) -> dict: + """Write a Terraform (.tf.json) recreation of the assets (provider, + variables, main and an import_commands.sh) into the output + directory. Overwrites files of the same name there.""" + from cloudg.renderers.terraform_export import TerraformExporter + + ds = ws_dataset(ctx, dataset) + assets = _select(ds, refs, asset_types) + ids = {a.id for a in assets} + out = ctx.workspace.output_path(subdir) + paths = TerraformExporter(output_dir=out).export( + assets, [e for e in ds.edges if e.source_id in ids or e.target_id in ids] + ) + return { + "dataset": ds.name, + "output_dir": str(out), + "resources": len(assets), + "files": {k: str(v) for k, v in paths.items()}, + } + + +@CATALOG.tool( + title="Export report", + category=CATEGORY, + sensitivity=Sensitivity.CONFIDENTIAL, + capabilities={R, FS}, + read_only=False, + destructive=False, + idempotent=True, + open_world=False, +) +def export_report( + ctx: Any, + format: ReportFormatArg = "json", # pylint: disable=redefined-builtin # MCP argument name + subdir: SubDir = "", + dataset: DatasetArg = "", +) -> dict: + """Write the dataset as a cloudg report file set into the output + directory and return the file paths.""" + ds = ws_dataset(ctx, dataset) + out = ctx.workspace.output_path(subdir) + out.mkdir(parents=True, exist_ok=True) + files: dict[str, str] = {} + if format == "json": + from cloudg.renderers.json_export import JSONExporter + + p = JSONExporter(output_dir=str(out)).export( + _scan_result(ds), graph_json=ds.builder.to_d3_json() + ) + files["json"] = str(p) + elif format == "html": + from cloudg.renderers.html_report import HTMLReportGenerator + + p = HTMLReportGenerator(output_dir=str(out)).generate( + _scan_result(ds), graph_json=ds.builder.to_d3_json() + ) + files["html"] = str(p) + elif format == "inventory": + files.update({k: str(v) for k, v in ds.inventory_view().export(out).items()}) + elif format == "graphml": + files["graphml"] = str(ds.builder.save_graphml(out / f"{ds.name}.graphml")) + else: + from cloudg.inventory.mapper import InventoryMapper + + mapper = InventoryMapper(ctx.workspace.config) + files.update( + { + k: str(v) + for k, v in mapper.export_merged(ds.inventory_view(), ds.findings, out).items() + } + ) + return {"dataset": ds.name, "format": format, "output_dir": str(out), "files": files} + + +def register(reg: Registry) -> None: + """Add this module's tools to ``reg``.""" + CATALOG.register(reg) diff --git a/cloudg/mcp/catalog/findings.py b/cloudg/mcp/catalog/findings.py new file mode 100644 index 0000000..bcdbde2 --- /dev/null +++ b/cloudg/mcp/catalog/findings.py @@ -0,0 +1,499 @@ +"""Findings tools: browse, summarise, prioritise, suppress and ingest +security findings. + +Findings come from scanners (Prowler, ScoutSuite, Checkov, Trivy, the IAM +linter), from cloudg's own graph reachability analysis, or from a loaded +findings.json. Each finding is matched to an asset by id / ARN / name. +""" + +from __future__ import annotations + +from typing import Annotated, Any, Literal + +from pydantic import BaseModel, Field + +from cloudg.mcp.catalog._common import ( + DEFAULT_LIMIT, + FINDING_FIELDS, + Catalog, + Cursor, + DatasetArg, + FindingPageOut, + Limit, + RefArg, + TopRisksOut, + asset_brief, + asset_uri, + bounded, + check_fields, + finding_brief, + finding_uri, + fingerprint, + paginate, + project, + schema, + ws_dataset, +) +from cloudg.mcp.catalog._findings import ( + FindingFilter, + check_report_tools, + filter_findings, + merge_compliance, + parse_reports, + renormalise, + risk_ranking, + sort_findings, + summarise_findings, +) +from cloudg.mcp.core import Capability, Registry, Sensitivity +from cloudg.mcp.state import SEVERITY_RANK, Dataset + +__all__ = [ + "CATALOG", + "FindingFilter", + "filter_findings", + "merge_compliance", + "register", + "renormalise", + "risk_ranking", +] + +CATEGORY = "findings" +R, W, FS = Capability.READ_STATE, Capability.WRITE_STATE, Capability.READ_FS + +FindingSort = Literal["risk", "severity", "detected_at", "title", "source_tool"] +FindingGroup = Literal["severity", "source_tool", "framework", "asset_type", "account", "asset"] + + +SeveritiesArg = Annotated[ + list[str] | None, Field(description="Any of CRITICAL, HIGH, MEDIUM, LOW, INFO.") +] +MinSeverityArg = Annotated[str, Field(description="At or above this severity.")] +SourceToolArg = Annotated[ + str, + Field(description="Substring of the producing tool, e.g. 'prowler', 'cloudg-reachability'."), +] +FrameworkArg = Annotated[ + str, Field(description="Substring of a compliance framework, e.g. 'CIS', 'PCI'.") +] +ResourceArg = Annotated[str, Field(description="Only findings on this asset (ref).")] +QueryArg = Annotated[str, Field(description="Substring of title or description.")] +FieldsArg = Annotated[ + list[str] | None, + Field(description=f"Project items to these fields. Valid: {', '.join(FINDING_FIELDS)}"), +] +ReportsArg = Annotated[ + dict[str, list[str]], + Field( + description="Tool -> report paths (files or directories inside an allowed root), " + "e.g. {'prowler': ['out/prowler/'], 'trivy': ['scan.json']}. Tools: prowler, " + "scoutsuite, checkov, trivy." + ), +] +IngestDatasetArg = Annotated[ + str, + Field( + description="Dataset to add the findings to; empty = active. Created if new_dataset " + "is true." + ), +] +NewDatasetArg = Annotated[ + bool, Field(description="Create a new findings-only dataset named `dataset` instead.") +] +NormaliseArg = Annotated[ + bool, + Field(description="Deduplicate across scanners and map to compliance frameworks afterwards."), +] +ActivateArg = Annotated[ + bool, + Field(description="With new_dataset: make the new dataset the active one (like load_dataset)."), +] +ReplaceArg = Annotated[ + bool, Field(description="With new_dataset: overwrite an existing dataset of that name.") +] + + +class ListFindingsArgs(BaseModel): + """list_findings arguments (the MCP wire schema).""" + + severities: SeveritiesArg = None + min_severity: MinSeverityArg = "" + source_tool: SourceToolArg = "" + framework: FrameworkArg = "" + resource: ResourceArg = "" + query: QueryArg = "" + include_suppressed: bool = False + sort_by: FindingSort = "risk" + fields: FieldsArg = None + limit: Limit = DEFAULT_LIMIT + cursor: Cursor = "" + dataset: DatasetArg = "" + + +CATALOG = Catalog() + +_RO = dict(read_only=True, idempotent=True, open_world=False, category=CATEGORY) + + +@CATALOG.tool( + title="List findings", + sensitivity=Sensitivity.CONFIDENTIAL, + output_schema=schema(FindingPageOut), + tags={"start-here"}, + **_RO, +) +def list_findings(ctx: Any, args: ListFindingsArgs) -> dict: + """Search findings by severity, tool, compliance framework, asset or + text, sorted by risk (default), severity, date, title or tool, with + cursor pagination. Suppressed findings are hidden unless + include_suppressed=true. Use get_finding for evidence and + remediation.""" + ds = ws_dataset(ctx, args.dataset) + proj = check_fields(args.fields, FINDING_FIELDS) + flt = FindingFilter.from_args(dict(args)) + items = sort_findings(filter_findings(ds, flt), args.sort_by) + fp = fingerprint(ds, f=flt.key(), o=args.sort_by) + page, env = paginate(items, args.limit, args.cursor, fp) + sev: dict[str, int] = {} + for f in items: + sev[f.severity.value] = sev.get(f.severity.value, 0) + 1 + items_out = [project(finding_brief(ds, f), proj) for f in page] + return {"dataset": ds.name, **env, "severity_breakdown": sev, "items": items_out} + + +@CATALOG.tool(title="Get finding", sensitivity=Sensitivity.CONFIDENTIAL, **_RO) +def get_finding( + ctx: Any, + finding_id: Annotated[ + str, Field(min_length=1, description="Finding id (or the scanner's source_finding_id).") + ], + dataset: DatasetArg = "", +) -> dict: + """One finding in full: description, evidence, remediation, CVSS, + compliance frameworks and the controls it fails, the affected asset + and whether it is suppressed (and why).""" + ds = ws_dataset(ctx, dataset) + f = ds.get_finding(finding_id) + out = { + **finding_brief(ds, f), + "description": f.description, + "evidence": f.evidence, + "remediation": f.remediation, + "source_finding_id": f.source_finding_id, + } + out["controls"] = [ + { + "framework": c.framework, + "control_id": c.control_id, + "control_title": c.control_title, + "status": c.status.value, + } + for c in ds.compliance + if f.id in c.finding_ids + ] + if f.is_suppressed: + out["suppression_reason"] = ds.suppression_reasons.get(f.id, "") + aid = ds.finding_asset_id(f) + if aid: + out["asset"] = asset_brief(ds, ds.by_id[aid]) + ctx.link(asset_uri(aid), out["asset"]["name"], title="Affected asset") + ctx.link(finding_uri(f.id), f.title) + return out + + +@CATALOG.tool(title="Findings summary", sensitivity=Sensitivity.INTERNAL, **_RO) +def findings_summary( + ctx: Any, + group_by: FindingGroup = "severity", + include_suppressed: bool = False, + top: Annotated[int, Field(ge=1, le=200)] = 25, + dataset: DatasetArg = "", +) -> dict: + """Finding counts grouped by severity, tool, compliance framework, + asset type, account or asset, each group with its severity + breakdown. Findings not matched to any mapped asset are counted + under 'unmapped'.""" + ds = ws_dataset(ctx, dataset) + groups = summarise_findings(ds, group_by, include_suppressed) + ordered = sorted( + groups.items(), + key=lambda kv: ( + -SEVERITY_RANK.get(kv[0], -1) if group_by == "severity" else -kv[1]["total"], + kv[0], + ), + ) + return { + "dataset": ds.name, + "group_by": group_by, + "total": sum(1 for f in ds.findings if include_suppressed or not f.is_suppressed), + "suppressed": sum(1 for f in ds.findings if f.is_suppressed), + "groups": dict(ordered[:top]), + "truncated": len(groups) > top, + } + + +@CATALOG.tool( + title="Findings for asset", + sensitivity=Sensitivity.CONFIDENTIAL, + output_schema=schema(FindingPageOut), + **_RO, +) +def findings_for_asset( + ctx: Any, + ref: RefArg, + include_suppressed: bool = False, + limit: Limit = DEFAULT_LIMIT, + dataset: DatasetArg = "", +) -> dict: + """All findings on one asset, highest risk first.""" + ds = ws_dataset(ctx, dataset) + a = ds.resolve_asset(ref) + fs = sort_findings( + [ + f + for f in ds.findings_by_asset.get(a.id, []) + if include_suppressed or not f.is_suppressed + ], + "risk", + ) + page, env = bounded(fs, limit) + return { + "dataset": ds.name, + "asset": asset_brief(ds, a), + **env, + "items": [finding_brief(ds, f) for f in page], + } + + +@CATALOG.tool( + title="Top risks", + sensitivity=Sensitivity.CONFIDENTIAL, + output_schema=schema(TopRisksOut), + **_RO, +) +def top_risks( + ctx: Any, + top: Annotated[int, Field(ge=1, le=100)] = 10, + min_severity: str = "", + dataset: DatasetArg = "", +) -> dict: + """Assets to fix first. Score = worst open finding's risk (0-10) x + exposure (1.5 if internet-exposed or reachable from 0.0.0.0/0 in the + graph) x blast radius (1-2, log of + transitive dependents) x volume (up to 1.5 for many findings). Each + item shows the score components and its top three findings.""" + ds = ws_dataset(ctx, dataset) + rows = risk_ranking(ds, min_severity) + return { + "dataset": ds.name, + "items": rows[:top], + "total_candidates": len(rows), + "formula": "max_risk * exposure * blast * volume", + } + + +@CATALOG.tool( + title="Suppress findings", + category=CATEGORY, + sensitivity=Sensitivity.CONFIDENTIAL, + capabilities={R, W}, + read_only=False, + destructive=False, + idempotent=True, + open_world=False, +) +def suppress_findings( + ctx: Any, + finding_ids: Annotated[list[str], Field(min_length=1, max_length=500)], + reason: Annotated[ + str, + Field(min_length=3, max_length=500, description="Why (accepted risk, false positive...)."), + ], + dataset: DatasetArg = "", +) -> dict: + """Mark findings as suppressed (accepted risk / false positive) in + the in-memory dataset. Suppressed findings drop out of list_findings, + summaries, top_risks and prompts; reverse with unsuppress_findings. + Nothing is written to disk or to the cloud.""" + ds = ws_dataset(ctx, dataset) + res = ds.set_suppressed(finding_ids, True, reason) + if res["changed"]: + ctx.workspace.mutated(ds) + return { + "dataset": ds.name, + "suppressed": res["changed"], + "already_suppressed": res["unchanged"], + "not_found": res["not_found"], + } + + +@CATALOG.tool( + title="Unsuppress findings", + category=CATEGORY, + sensitivity=Sensitivity.CONFIDENTIAL, + capabilities={R, W}, + read_only=False, + destructive=False, + idempotent=True, + open_world=False, +) +def unsuppress_findings( + ctx: Any, + finding_ids: Annotated[list[str], Field(min_length=1, max_length=500)], + dataset: DatasetArg = "", +) -> dict: + """Restore suppressed findings so they count again.""" + ds = ws_dataset(ctx, dataset) + res = ds.set_suppressed(finding_ids, False) + if res["changed"]: + ctx.workspace.mutated(ds) + return { + "dataset": ds.name, + "unsuppressed": res["changed"], + "not_suppressed": res["unchanged"], + "not_found": res["not_found"], + } + + +@CATALOG.tool( + title="Ingest scanner reports", + category=CATEGORY, + sensitivity=Sensitivity.CONFIDENTIAL, + capabilities={R, W, FS}, + read_only=False, + destructive=False, + idempotent=False, + open_world=False, +) +def ingest_reports( + ctx: Any, + reports: ReportsArg, + dataset: IngestDatasetArg = "", + new_dataset: NewDatasetArg = False, + normalise: NormaliseArg = True, + activate: ActivateArg = True, + replace: ReplaceArg = False, +) -> dict: + """Parse existing Prowler (ASFF) / ScoutSuite / Checkov / Trivy output + (no scanner is run, no cloud access) and add the findings to a + dataset, matching them to its assets by ARN / id. With + new_dataset=true the findings go into a new dataset, which becomes + the active one unless activate=false; the result's `active` field + says which dataset is active afterwards. Tip: snapshot_dataset first + to diff before / after.""" + ws = ctx.workspace + check_report_tools(reports) + # Validate the destination before parsing anything + if new_dataset: + target: Any = ws.check_name(dataset or ws.unique_name("ingested"), replace=replace) + else: + target = ws_dataset(ctx, dataset) + per_path, errors, new = parse_reports(ws, reports) + if new_dataset: + target = ws.add( + Dataset(name=target, source="ingest", kind="report"), activate=activate, replace=replace + ) + before = len(target.findings) + _store_findings(ws, target, new, normalise) + return { + "dataset": target.name, + "parsed": len(new), + "per_path": per_path, + "errors": errors, + "findings_before": before, + "findings_after": len(target.findings), + "matched_to_assets": sum(1 for f in new if target.finding_asset_id(f)), + "normalised": normalise, + "active": ws.active_name, + } + + +def _store_findings(ws: Any, ds: Dataset, new: list, normalise: bool) -> None: + if normalise: + renormalise(ws, ds, new) + else: + ds.add_findings(new) + ws.mutated(ds) + + +@CATALOG.tool( + title="Normalise findings", + category=CATEGORY, + sensitivity=Sensitivity.INTERNAL, + capabilities={R, W}, + read_only=False, + destructive=False, + idempotent=True, + open_world=False, +) +def normalise_findings(ctx: Any, dataset: DatasetArg = "") -> dict: + """Re-run cloudg's normaliser on a dataset's findings: deduplicate + within and across scanners (check-equivalence rules), score, and + rebuild the compliance-control mapping.""" + ds = ws_dataset(ctx, dataset) + before, before_c = len(ds.findings), len(ds.compliance) + renormalise(ctx.workspace, ds) + ctx.workspace.mutated(ds) + return { + "dataset": ds.name, + "findings_before": before, + "findings_after": len(ds.findings), + "compliance_results_before": before_c, + "compliance_results_after": len(ds.compliance), + "frameworks": sorted({c.framework for c in ds.compliance}), + } + + +@CATALOG.tool( + title="Reachability findings", + category=CATEGORY, + sensitivity=Sensitivity.CONFIDENTIAL, + capabilities={R, W}, + read_only=False, + destructive=False, + idempotent=False, + open_world=False, +) +def reachability_findings( + ctx: Any, + add_to_dataset: Annotated[ + bool, Field(description="Append them to the dataset's findings (otherwise preview only).") + ] = False, + limit: Limit = DEFAULT_LIMIT, + dataset: DatasetArg = "", +) -> dict: + """Run cloudg's graph reachability analysis: internet-exposed data + stores (CRITICAL), unexpected internet exposure (HIGH) and sensitive + ports open to 0.0.0.0/0 (CRITICAL). Preview, or add them to the + dataset (skipping ones already present). Ids come from the analyser and + are stable across scans (derived from the rule and the asset's ARN or + id), so a previewed id is the id that gets added.""" + from cloudg.graph.reachability import ReachabilityAnalyzer + + ds = ws_dataset(ctx, dataset) + found = ReachabilityAnalyzer(ds.graph.copy()).generate_findings() + have_ids = {f.id for f in ds.findings} + have_keys = {(f.source_tool, f.title, f.resource_id) for f in ds.findings} + fresh = [ + f + for f in found + if f.id not in have_ids and (f.source_tool, f.title, f.resource_id) not in have_keys + ] + added = 0 + if add_to_dataset and fresh: + added = ds.add_findings(fresh) + ctx.workspace.mutated(ds) + page, env = bounded(fresh, limit) + return { + "dataset": ds.name, + "generated": len(found), + "new": len(fresh), + "added": added, + **env, + "items": [finding_brief(ds, f) for f in page], + } + + +def register(reg: Registry) -> None: + """Add this module's tools to ``reg``.""" + CATALOG.register(reg) diff --git a/cloudg/mcp/catalog/graph.py b/cloudg/mcp/catalog/graph.py new file mode 100644 index 0000000..9d7e72c --- /dev/null +++ b/cloudg/mcp/catalog/graph.py @@ -0,0 +1,384 @@ +"""Graph and relationship tools: neighbours, paths, exposure, attack and +lateral-movement paths, dependency / blast-radius analysis, cross-account +edges, centrality, sub-graph export and security-service coverage. + +Three views of the same edges are used, each for the question it answers: + +- the relationship graph (``GraphBuilder``): edges as collected, read + "source verb target"; +- the flow graph (:attr:`Dataset.flow_graph`): the network-flow hops of + :func:`cloudg.graph.reachability.network_flow_graph` (the same rules as + the internet-exposure analysis) plus the identity edges ASSUMES_ROLE, + IAM_TRUST and GRANTS_ACCESS, so a path may continue from a compromised + workload through its role to what the role can reach; +- the dependency graph (``DependencyGraph``): every edge turned into a + "dependent -> dependency" arrow, which depends_on / dependents / blast + radius walk (availability and change impact). +""" + +from __future__ import annotations + +from typing import Annotated, Any, Literal + +from pydantic import BaseModel, Field + +from cloudg.mcp.catalog._common import ( + DEFAULT_LIMIT, + Catalog, + Cursor, + DatasetArg, + EdgePageOut, + Limit, + NeighborsOut, + PathsOut, + RefArg, + asset_brief, + asset_uri, + bounded, + edge_brief, + fingerprint, + node_brief, + paginate, + parse_enums, + schema, + ws_dataset, +) +from cloudg.mcp.catalog._graph import ( + GRAPH_RO, + edge_type_values, + CROWN_JEWELS, + IDENTITY_EDGES, + INTERNET, + SENSITIVE_PORTS, + AttackSearch, + EdgeQuery, + Walk, + allows_port, + attack_graph, + attack_paths_found, + bfs, + brief, + describe_path, + exposure_evidence, + identity_graph, + lateral_chains, + lateral_starts, + path_graph, + resolve_node, + simple_paths, + with_internet, +) +from cloudg.mcp.catalog._graph_impact import CATALOG as IMPACT_CATALOG +from cloudg.mcp.core import Capability, Registry, Sensitivity +from cloudg.schema.models import AssetType + +__all__ = [ + "CATALOG", + "CROWN_JEWELS", + "IDENTITY_EDGES", + "SENSITIVE_PORTS", + "allows_port", + "register", +] + +R = Capability.READ_STATE + +Direction = Literal["out", "in", "both"] +PathMode = Literal["flow", "directed", "undirected"] + +EdgeTypesArg = Annotated[ + list[str] | None, + Field(description="Only these EdgeType values, e.g. ['ASSUMES_ROLE','GRANTS_ACCESS']."), +] + +SourceArg = Annotated[str, Field(description="Source asset ref (or 'internet').")] +TargetArg = Annotated[str, Field(description="Target asset ref.")] +RelationshipArg = Annotated[str, Field(description="Exact relationship / RelationType.")] +InternetOnlyArg = Annotated[ + bool, + Field( + description="Only ingress edges from the internet (0.0.0.0/0, ::/0, or the Azure " + "Internet / Any / * sources)." + ), +] +PortArg = Annotated[int | None, Field(ge=0, le=65535, description="Edges allowing this port.")] + + +class GetEdgesArgs(BaseModel): + """get_edges arguments (the MCP wire schema).""" + + source: SourceArg = "" + target: TargetArg = "" + edge_types: list[str] | None = None + relationship: RelationshipArg = "" + cross_account_only: bool = False + internet_only: InternetOnlyArg = False + port: PortArg = None + limit: Limit = DEFAULT_LIMIT + cursor: Cursor = "" + dataset: DatasetArg = "" + + +CATALOG = Catalog() + + +@CATALOG.tool( + title="Neighbors", + sensitivity=Sensitivity.CONFIDENTIAL, + output_schema=schema(NeighborsOut), + **GRAPH_RO, +) +def neighbors( + ctx: Any, + ref: RefArg, + direction: Annotated[ + Direction, Field(description="out = edges from the asset, in = edges to it.") + ] = "both", + edge_types: EdgeTypesArg = None, + depth: Annotated[int, Field(ge=1, le=4)] = 1, + limit: Annotated[int, Field(ge=1, le=500)] = 100, + dataset: DatasetArg = "", +) -> dict: + """Assets connected to one asset, up to `depth` hops, with the edges + between them (type, relationship, ports, CIDR). Filter by direction + and edge types. Example: what a Lambda can reach through its role: + ref='api-handler', direction='out', edge_types=['ASSUMES_ROLE', + 'GRANTS_ACCESS'], depth=2.""" + ds = ws_dataset(ctx, dataset) + start = resolve_node(ds, ref) + walk = Walk(direction, edge_type_values(edge_types) or None, depth, limit) + nodes, edges, truncated = bfs(ds, start, walk) + ctx.link(asset_uri(start) + "/neighbors", "neighbors") + return { + "dataset": ds.name, + "asset": node_brief(ds, start), + "nodes": [node_brief(ds, n) for n in nodes], + "edges": [edge_brief(ds, e) for e in edges], + "node_count": len(nodes), + "edge_count": len(edges), + "truncated": truncated, + } + + +@CATALOG.tool( + title="Get edges", + sensitivity=Sensitivity.CONFIDENTIAL, + output_schema=schema(EdgePageOut), + **GRAPH_RO, +) +def get_edges(ctx: Any, args: GetEdgesArgs) -> dict: + """Search edges by endpoint, edge type, relationship, cross-account, + internet origin or port, with pagination. Egress rules never count as + internet-sourced. Example: every rule that opens SSH to the internet: + internet_only=true, port=22.""" + ds = ws_dataset(ctx, args.dataset) + q = EdgeQuery( + source=resolve_node(ds, args.source, virtual_internet=True) if args.source else None, + target=resolve_node(ds, args.target) if args.target else None, + edge_types=edge_type_values(args.edge_types), + relationship=args.relationship, + internet_only=args.internet_only, + port=args.port, + ) + xa = args.cross_account_only + out = [ + b + for b in (edge_brief(ds, e) for e in ds.edges if q.matches(e)) + if not xa or b.get("cross_account") + ] + fp = fingerprint(ds, q=vars(q) | {"edge_types": sorted(q.edge_types or ())}, x=xa) + page, env = paginate(out, args.limit, args.cursor, fp) + return {"dataset": ds.name, **env, "items": page} + + +@CATALOG.tool( + title="Find paths", + sensitivity=Sensitivity.CONFIDENTIAL, + output_schema=schema(PathsOut), + **GRAPH_RO, +) +def find_paths( + ctx: Any, + source: Annotated[str, Field(min_length=1, description="Start asset ref, or 'internet'.")], + target: Annotated[str, Field(min_length=1, description="End asset ref.")], + max_depth: Annotated[int, Field(ge=1, le=12)] = 6, + max_paths: Annotated[int, Field(ge=1, le=50)] = 5, + mode: Annotated[ + PathMode, + Field( + description="flow: network-flow hops (as the reachability analysis walks them: " + "security groups admit traffic to attached resources) plus identity pivots " + "(ASSUMES_ROLE, IAM_TRUST, GRANTS_ACCESS); directed: edges as collected; " + "undirected: any connection." + ), + ] = "flow", + dataset: DatasetArg = "", +) -> dict: + """Shortest connection paths between two assets (shortest first), + each with the node briefs and the edge type of every hop. Use + source='internet' to see how something is reachable from any internet + entry point (0.0.0.0/0, ::/0, Azure Internet sources).""" + ds = ws_dataset(ctx, dataset) + s = resolve_node(ds, source, virtual_internet=True) + t = resolve_node(ds, target) + g = path_graph(ds, mode) + if s == INTERNET: + g = with_internet(ds, g) + paths = simple_paths(g, s, t, max_depth, max_paths) + out = { + "dataset": ds.name, + "source": brief(ds, s), + "target": node_brief(ds, t), + "mode": mode, + "paths": [describe_path(ds, g, p) for p in paths[:max_paths]], + "total_found": min(len(paths), max_paths), + "truncated": len(paths) > max_paths, + } + if not paths: + out["hint"] = ( + f"No path within {max_depth} hops in '{mode}' mode. Try mode='undirected' or a " + "larger max_depth, or check neighbors() of each end." + ) + return out + + +@CATALOG.tool( + title="Attack paths", + sensitivity=Sensitivity.CONFIDENTIAL, + output_schema=schema(PathsOut), + **GRAPH_RO, +) +def attack_paths( + ctx: Any, + target: Annotated[ + str, + Field( + description="Asset to reach; empty = every crown-jewel " + "asset (databases, buckets, secrets, keys, roles)." + ), + ] = "", + max_depth: Annotated[int, Field(ge=1, le=12)] = 8, + max_paths: Annotated[int, Field(ge=1, le=100)] = 20, + dataset: DatasetArg = "", +) -> dict: + """Shortest routes from the internet (internet-sourced ingress rules + and internet-exposed assets) to sensitive assets over the flow graph: + network-flow hops as the reachability analysis walks them, plus + identity pivots (ASSUMES_ROLE, IAM_TRUST, GRANTS_ACCESS) from a reached + workload. Ranked by target sensitivity, open findings along the path + and length. Each path lists every hop and edge type; investigate the + entry point and each hop with get_asset / findings_for_asset.""" + ds = ws_dataset(ctx, dataset) + g = attack_graph(ds) + if g.out_degree(INTERNET) == 0: + return { + "dataset": ds.name, + "paths": [], + "total_found": 0, + "truncated": False, + "hint": "Nothing in this dataset is internet-exposed.", + } + if target: + targets = {resolve_node(ds, target): 10} + else: + targets = { + a.id: CROWN_JEWELS[a.asset_type] for a in ds.assets if a.asset_type in CROWN_JEWELS + } + found = attack_paths_found(AttackSearch(ds, targets, max_depth), g) + page, env = bounded(found, max_paths) + return { + "dataset": ds.name, + "paths": page, + "total_found": env["total"], + "truncated": env["truncated"], + "targets_considered": len(targets), + } + + +@CATALOG.tool( + title="Lateral movement paths", + sensitivity=Sensitivity.CONFIDENTIAL, + output_schema=schema(PathsOut), + **GRAPH_RO, +) +def lateral_movement_paths( + ctx: Any, + start: Annotated[ + str, + Field( + description="Start from this asset; empty = every " + "internet-exposed asset and external account." + ), + ] = "", + max_depth: Annotated[int, Field(ge=1, le=8)] = 4, + max_paths: Annotated[int, Field(ge=1, le=100)] = 25, + dataset: DatasetArg = "", +) -> dict: + """Identity pivot chains: from a foothold (internet-exposed compute + or an external / cross-account principal) along ASSUMES_ROLE, + IAM_TRUST and GRANTS_ACCESS edges to the roles and resources it can + reach. Shows how a compromise propagates through IAM.""" + ds = ws_dataset(ctx, dataset) + starts = [resolve_node(ds, start)] if start else lateral_starts(ds) + g = identity_graph(ds) + chains = lateral_chains(ds, g, starts, max_depth) + page, env = bounded(chains, max_paths) + return { + "dataset": ds.name, + "paths": [describe_path(ds, g, p) for p in page], + "total_found": env["total"], + "truncated": env["truncated"], + "starting_points": len(set(starts)), + } + + +@CATALOG.tool(title="Internet exposure", sensitivity=Sensitivity.CONFIDENTIAL, **GRAPH_RO) +def internet_exposure( + ctx: Any, + asset_types: list[str] | None = None, + include_reachable: Annotated[ + bool, + Field( + description="Also list assets transitively reachable from the internet in the graph." + ), + ] = False, + limit: Limit = DEFAULT_LIMIT, + cursor: Cursor = "", + dataset: DatasetArg = "", +) -> dict: + """Internet-exposed assets with the evidence: the internet-sourced + rules / INTERNET_EXPOSED edges reaching them, sensitive ports open + (SSH, RDP, databases...), whether a WAF protects them, and their open + findings. Sorted with unprotected sensitive-port exposures first.""" + ds = ws_dataset(ctx, dataset) + types = parse_enums(asset_types, AssetType, "asset type") + items = [ + {**asset_brief(ds, a), **exposure_evidence(ds, a.id)} + for a in ds.assets + if a.is_internet_exposed and (not types or a.asset_type in types) + ] + items.sort( + key=lambda i: ( + -len(i["sensitive_ports_open"]), + i["protected_by_waf"], + -(i["open_findings"] or 0), + i["name"], + ) + ) + page, env = paginate(items, limit, cursor, fingerprint(ds, t=sorted(x.value for x in types))) + out: dict[str, Any] = {"dataset": ds.name, **env, "items": page} + if include_reachable: + reach = sorted(n for n in ds.internet_reachable() if n in ds.by_id) + sample, renv = bounded(reach, limit) + out["graph_reachable"] = { + "total": renv["total"], + "truncated": renv["truncated"], + "items": [node_brief(ds, n) for n in sample], + } + return out + + +def register(reg: Registry) -> None: + """Add this module's tools to ``reg``.""" + CATALOG.register(reg) + IMPACT_CATALOG.register(reg) diff --git a/cloudg/mcp/catalog/inventory.py b/cloudg/mcp/catalog/inventory.py new file mode 100644 index 0000000..3e3c29c --- /dev/null +++ b/cloudg/mcp/catalog/inventory.py @@ -0,0 +1,387 @@ +"""Inventory tools: search, inspect and aggregate the assets of a dataset. + +These answer "what do we have": filtered / sorted / paginated asset +search, one asset in detail, counts by any dimension, accounts, regions, +tags, collection coverage, unresolved references and the organization +(AWS Organizations / Control Tower, Azure management groups, GCP folders) +topology. +""" + +from __future__ import annotations + +from typing import Annotated, Any, Literal + +from pydantic import BaseModel, Field + +from cloudg.inventory.mapper_result import _service_of +from cloudg.mcp.catalog._common import ( + ASSET_FIELDS, + DEFAULT_LIMIT, + AssetDetailOut, + AssetPageOut, + Catalog, + CountOut, + Cursor, + DatasetArg, + Limit, + RefArg, + asset_brief, + asset_uri, + bounded, + check_fields, + finding_brief, + fingerprint, + paginate, + project, + schema, + ws_dataset, +) +from cloudg.mcp.catalog._inventory import ( + AssetFilter, + account_rows, + filter_assets, + group_key, + org_from_assets, + org_from_discovery, + relation_view, + sort_assets, + truncate_metadata, +) +from cloudg.mcp.core import Capability, Registry, Sensitivity + +CATEGORY = "inventory" +R = Capability.READ_STATE + +SortKey = Literal["name", "type", "region", "account", "risk", "findings"] +GroupBy = Literal["type", "provider", "region", "account", "service", "exposure", "severity"] + +QueryArg = Annotated[ + str, Field(description="Case-insensitive substring of name, ARN, id or a tag value.") +] +ProviderArg = Annotated[str, Field(description="aws, azure or gcp.")] +AssetTypesArg = Annotated[ + list[str] | None, + Field(description="AssetType values, e.g. ['EC2','S3_BUCKET'] (describe_schema)."), +] +RegionArg = Annotated[str, Field(description="Exact region, e.g. us-east-1.")] +AccountArg = Annotated[str, Field(description="Exact account / subscription / project.")] +TagArg = Annotated[str, Field(description="Tag filter: 'key' or 'key=value'.")] +ExposedArg = Annotated[bool | None, Field(description="Only exposed (true) / not exposed (false).")] +HasFindingsArg = Annotated[ + bool | None, Field(description="Only assets with (true) / without (false) open findings.") +] +MinSeverityArg = Annotated[ + str, Field(description="Only assets with an open finding at or above this severity.") +] +SortArg = Annotated[SortKey, Field(description="'risk' ranks by worst finding + exposure.")] +FieldsArg = Annotated[ + list[str] | None, + Field(description=f"Project each item to these fields. Valid: {', '.join(ASSET_FIELDS)}"), +] + + +class FindAssetsArgs(BaseModel): + """find_assets arguments (the MCP wire schema).""" + + query: QueryArg = "" + provider: ProviderArg = "" + asset_types: AssetTypesArg = None + region: RegionArg = "" + account_id: AccountArg = "" + tag: TagArg = "" + internet_exposed: ExposedArg = None + has_findings: HasFindingsArg = None + min_severity: MinSeverityArg = "" + sort_by: SortArg = "name" + descending: bool = False + fields: FieldsArg = None + limit: Limit = DEFAULT_LIMIT + cursor: Cursor = "" + dataset: DatasetArg = "" + + +class CountAssetsArgs(BaseModel): + """count_assets arguments (the MCP wire schema).""" + + group_by: GroupBy = "type" + provider: str = "" + asset_types: list[str] | None = None + region: str = "" + account_id: str = "" + internet_exposed: bool | None = None + top: Annotated[int, Field(ge=1, le=500)] = 50 + dataset: DatasetArg = "" + + +_RO = dict(category=CATEGORY, read_only=True, idempotent=True, open_world=False) + +CATALOG = Catalog() + + +@CATALOG.tool( + title="Find assets", + sensitivity=Sensitivity.CONFIDENTIAL, + output_schema=schema(AssetPageOut), + tags={"start-here"}, + **_RO, +) +def find_assets(ctx: Any, args: FindAssetsArgs) -> dict: + """Search assets with filters (provider, type, region, account, tag, + exposure, findings, severity, free text), sorting and cursor + pagination. Returns compact briefs; call get_asset (or read the + item's uri) for details. Example: find internet-exposed EC2 with + HIGH findings: asset_types=['EC2'], internet_exposed=true, + min_severity='HIGH'.""" + ds = ws_dataset(ctx, args.dataset) + proj = check_fields(args.fields, ASSET_FIELDS) + flt = AssetFilter.from_args(dict(args)) + matched = filter_assets(ds, flt) + sort_assets(ds, matched, args.sort_by, args.descending) + fp = fingerprint(ds, f=flt.key(), s=args.sort_by, d=args.descending) + page, env = paginate(matched, args.limit, args.cursor, fp) + items = [project(asset_brief(ds, a, tags="tags" in (proj or [])), proj) for a in page] + out = {"dataset": ds.name, **env, "items": items} + if not matched and ds.assets: + out["hint"] = ( + "No assets matched. Loosen the filters, or use list_asset_types / " + "list_accounts / list_regions to see valid values." + ) + return out + + +@CATALOG.tool( + title="Get asset", + sensitivity=Sensitivity.CONFIDENTIAL, + output_schema=schema(AssetDetailOut), + **_RO, +) +def get_asset( + ctx: Any, + ref: RefArg, + include_findings: bool = True, + max_relations: Annotated[int, Field(ge=0, le=200)] = 20, + dataset: DatasetArg = "", +) -> dict: + """One asset in detail: identity, tags, exposure, relation counts by + edge type with the first neighbours in / out, direct dependency + counts, its open findings, and the metadata keys available via + get_asset_metadata. Use neighbors / dependency_tree to go further.""" + ds = ws_dataset(ctx, dataset) + a = ds.resolve_asset(ref) + out = asset_brief(ds, a, tags=True) + out["dataset"] = ds.name + out["service"] = _service_of(a) + g = ds.graph + out["relations"] = relation_view(ds, a, max_relations) + needs, needed_by = ds.dependency_graph().direct_counts(a.id) + out["dependencies"] = {"direct_depends_on": needs, "direct_dependents": needed_by} + out["degree"] = {"in": g.in_degree(a.id), "out": g.out_degree(a.id)} if a.id in g else {} + if include_findings: + fs = sorted(ds.open_findings(a.id), key=lambda f: -f.risk_score) + out["findings"] = [finding_brief(ds, f) for f in fs[:25]] + out["metadata_keys"] = sorted(a.metadata)[:100] + out["collected_at"] = a.collected_at.isoformat() if a.collected_at else None + ctx.link(asset_uri(a.id), a.name, title=f"{a.asset_type.value} {a.name}") + ctx.link(asset_uri(a.id) + "/neighbors", f"{a.name} neighbors") + return out + + +@CATALOG.tool(title="Get asset metadata", sensitivity=Sensitivity.RESTRICTED, **_RO) +def get_asset_metadata( + ctx: Any, + ref: RefArg, + keys: Annotated[ + list[str] | None, + Field(description="Only these metadata keys (see get_asset.metadata_keys)."), + ] = None, + max_chars: Annotated[int, Field(ge=1000, le=200_000)] = 40_000, + dataset: DatasetArg = "", +) -> dict: + """Raw collector metadata of one asset (configuration, policies, + rules, declared relations). May hold sensitive values such as policy + documents or environment settings, so it is RESTRICTED. Ask for + specific keys to keep the result small.""" + ds = ws_dataset(ctx, dataset) + a = ds.resolve_asset(ref) + meta = a.metadata + missing: list[str] = [] + if keys: + missing = [k for k in keys if k not in meta] + meta = {k: meta[k] for k in keys if k in meta} + out: dict[str, Any] = {"id": a.id, "name": a.name, "type": a.asset_type.value} + out.update(truncate_metadata(meta, max_chars)) + if missing: + out["missing_keys"] = missing + return out + + +@CATALOG.tool( + title="Count assets", + sensitivity=Sensitivity.INTERNAL, + output_schema=schema(CountOut), + **_RO, +) +def count_assets(ctx: Any, args: CountAssetsArgs) -> dict: + """Count assets grouped by type, provider, region, account, cloud + service, exposure or worst open-finding severity, with optional + filters. A cheap aggregate view with no identifiers beyond group keys.""" + ds = ws_dataset(ctx, args.dataset) + flt = AssetFilter.from_args(dict(args)) + matched = filter_assets(ds, flt) + groups: dict[str, int] = {} + for a in matched: + k = group_key(ds, a, args.group_by) + groups[k] = groups.get(k, 0) + 1 + ordered = sorted(groups.items(), key=lambda kv: (-kv[1], kv[0])) + return { + "dataset": ds.name, + "group_by": args.group_by, + "total": len(matched), + "group_count": len(groups), + "groups": dict(ordered[: args.top]), + "truncated": len(groups) > args.top, + } + + +@CATALOG.tool(title="List asset types", sensitivity=Sensitivity.INTERNAL, **_RO) +def list_asset_types(ctx: Any, dataset: DatasetArg = "") -> dict: + """Asset types present in the dataset with counts, exposed counts + and open-finding counts. These are the valid values for asset_types filters.""" + ds = ws_dataset(ctx, dataset) + stats: dict[str, dict[str, int]] = {} + for a in ds.assets: + s = stats.setdefault( + a.asset_type.value, {"count": 0, "internet_exposed": 0, "open_findings": 0} + ) + s["count"] += 1 + s["internet_exposed"] += int(a.is_internet_exposed) + s["open_findings"] += len(ds.open_findings(a.id)) + items = [{"type": t, **s} for t, s in sorted(stats.items(), key=lambda kv: -kv[1]["count"])] + return {"dataset": ds.name, "total_types": len(items), "items": items} + + +@CATALOG.tool(title="List accounts", sensitivity=Sensitivity.CONFIDENTIAL, **_RO) +def list_accounts( + ctx: Any, limit: Limit = 100, cursor: Cursor = "", dataset: DatasetArg = "" +) -> dict: + """Accounts / subscriptions / projects in the dataset with provider, + display name, asset count, regions, exposed assets, open findings + and whether the account is external (only referenced, not mapped).""" + ds = ws_dataset(ctx, dataset) + page, env = paginate(account_rows(ds), limit, cursor, fingerprint(ds, t="accounts")) + return {"dataset": ds.name, **env, "items": page} + + +@CATALOG.tool(title="List regions", sensitivity=Sensitivity.INTERNAL, **_RO) +def list_regions(ctx: Any, dataset: DatasetArg = "") -> dict: + """Regions in the dataset with providers, asset and account counts, + plus the regions each provider was configured / scanned for.""" + ds = ws_dataset(ctx, dataset) + regions: dict[str, dict[str, Any]] = {} + for a in ds.assets: + r = regions.setdefault( + a.region, {"region": a.region, "providers": set(), "assets": 0, "accounts": set()} + ) + r["providers"].add(a.provider.value) + r["assets"] += 1 + if a.account_id: + r["accounts"].add(a.account_id) + items = [ + {**r, "providers": sorted(r["providers"]), "accounts": len(r["accounts"])} + for r in sorted(regions.values(), key=lambda r: -r["assets"]) + ] + return {"dataset": ds.name, "total": len(items), "items": items, "scanned_regions": ds.regions} + + +@CATALOG.tool(title="List tags", sensitivity=Sensitivity.CONFIDENTIAL, **_RO) +def list_tags( + ctx: Any, + key: Annotated[str, Field(description="Only this tag key (shows its values).")] = "", + top: Annotated[int, Field(ge=1, le=200)] = 30, + dataset: DatasetArg = "", +) -> dict: + """Tag keys used across assets with how many assets carry each and + their most common values. Pass key= to list one key's values. Use + the result with find_assets(tag='key=value').""" + ds = ws_dataset(ctx, dataset) + keys: dict[str, dict[str, int]] = {} + for a in ds.assets: + for k, v in a.tags.items(): + if key and k != key: + continue + vals = keys.setdefault(k, {}) + vals[v] = vals.get(v, 0) + 1 + items = [] + for k, vals in sorted(keys.items(), key=lambda kv: -sum(kv[1].values())): + top_vals = sorted(vals.items(), key=lambda kv: -kv[1])[: (top if key else 5)] + items.append( + { + "key": k, + "assets": sum(vals.values()), + "distinct_values": len(vals), + "top_values": dict(top_vals), + } + ) + page, env = bounded(items, top) + return {"dataset": ds.name, **env, "items": page} + + +@CATALOG.tool(title="Coverage report", sensitivity=Sensitivity.INTERNAL, **_RO) +def coverage_report(ctx: Any, dataset: DatasetArg = "") -> dict: + """Which collectors / services succeeded or failed per provider, + account and region during collection, with error messages for + failures. Explains gaps (e.g. AccessDenied) before trusting the map. + Only live collections carry coverage records.""" + ds = ws_dataset(ctx, dataset) + if not ds.coverage: + return { + "dataset": ds.name, + "records": [], + "note": "This dataset has no coverage records (they exist only for live " + "collections via map_inventory / collect_assets).", + } + records = [c.to_summary() for c in ds.coverage] + total = sum(r["total_services"] for r in records) + ok = sum(r["successful"] for r in records) + return { + "dataset": ds.name, + "records": records, + "total_services": total, + "successful": ok, + "failed": sum(r["failed"] for r in records), + "coverage_pct": round(ok / total * 100, 1) if total else 0.0, + } + + +@CATALOG.tool(title="Unresolved references", sensitivity=Sensitivity.CONFIDENTIAL, **_RO) +def unresolved_references( + ctx: Any, limit: Limit = DEFAULT_LIMIT, cursor: Cursor = "", dataset: DatasetArg = "" +) -> dict: + """Relations the linker could not resolve to a mapped asset (e.g. a + role ARN in a service that was not collected). They are blind spots in + the relationship graph.""" + ds = ws_dataset(ctx, dataset) + page, env = paginate(ds.unresolved_references, limit, cursor, fingerprint(ds, t="unres")) + return {"dataset": ds.name, **env, "items": page} + + +@CATALOG.tool(title="Organization topology", sensitivity=Sensitivity.CONFIDENTIAL, **_RO) +def organization_topology( + ctx: Any, + include_policies: bool = False, + max_accounts: Annotated[int, Field(ge=1, le=1000)] = 200, + dataset: DatasetArg = "", +) -> dict: + """The organization hierarchy: organization, OUs / management groups + / folders, member accounts with their OU path, Control Tower landing + zone and governed regions, and (optionally) the SCPs / org policies + and what they target.""" + ds = ws_dataset(ctx, dataset) + if ds.organization: + return org_from_discovery(ds, ds.organization, include_policies, max_accounts) + return org_from_assets(ds, include_policies, max_accounts) + + +def register(reg: Registry) -> None: + """Add this module's tools to ``reg``.""" + CATALOG.register(reg) diff --git a/cloudg/mcp/catalog/live.py b/cloudg/mcp/catalog/live.py new file mode 100644 index 0000000..cfcc0a7 --- /dev/null +++ b/cloudg/mcp/catalog/live.py @@ -0,0 +1,449 @@ +"""Live tools: collect from cloud APIs, run scanners, run the full pipeline. + +These are the only tools that touch the outside world (``open_world``): +they call provider APIs with the configured credentials +(``cloud_access``) and, for scanners, spawn external binaries (``exec``). +They are long-running, report progress per pipeline phase, and land their +results in the workspace as a new dataset. The tool result is a summary +plus resource links, never the full dump. + +The engine is looked up as ``cloudg.api.CloudGEngine`` at call time so it +can be swapped (tests monkeypatch it; embedders can subclass it). +""" + +from __future__ import annotations + +import asyncio +import time +from typing import Annotated, Any + +from pydantic import BaseModel, Field + +from cloudg.mcp.catalog._common import Catalog, DatasetArg, ws_dataset +from cloudg.mcp.catalog._live import ( + _CHECKS, + FORCE_ROLES, + PHASES, + PROVIDERS, + SCANNERS, + LiveOperationRefused, + LiveRun, + ScanJob, + _bypass, + _check_aws, + _check_azure, + _check_gcp, + _check_images, + _config, + _guarded, + _latest_live_dataset, + _name, + _run_scan, + _run_throttling, + _scanner_config, + _throttling_summary, + engine_errors, + live_scopes, + run_live, +) +from cloudg.mcp.core import ( + Capability, + MCPLayerError, + Registry, + Sensitivity, +) +from cloudg.mcp.state import Dataset + +__all__ = [ + "CATALOG", + "FORCE_ROLES", + "PHASES", + "PREFLIGHT_TIMEOUT_S", + "PROVIDERS", + "SCANNERS", + "LiveOperationRefused", + "PreflightError", + "_check_aws", + "_check_azure", + "_check_gcp", + "_latest_live_dataset", + "_run_throttling", + "_throttling_summary", + "credential_problems", + "live_scopes", + "rate_limit_snapshot", + "register", +] + +CATEGORY = "live" +R, W, FS = Capability.READ_STATE, Capability.WRITE_STATE, Capability.WRITE_FS +CLOUD, EXEC = Capability.CLOUD_ACCESS, Capability.EXEC + +Providers = Annotated[ + list[str] | None, + Field(description="Subset of aws, azure, gcp; empty = the configured providers."), +] +Regions = Annotated[ + list[str] | None, + Field( + description="Regions to collect for every selected provider; empty = configured regions." + ), +] +NameArg = Annotated[ + str, Field(description="Dataset name for the result; default '-'.") +] + + +class PreflightError(MCPLayerError): + """The configured credentials for a provider cannot work, so a live + collection would only burn minutes of failing API calls.""" + + +# Bound on the whole credential check: the AWS default chain may probe the +# instance metadata service and Azure fetches one AAD token. +PREFLIGHT_TIMEOUT_S = 25.0 + + +def credential_problems(cfg: Any) -> dict[str, str]: + """Provider -> why its credentials cannot work, for every configured + provider whose credentials fail a cheap local check (no cloud API + calls beyond one Azure token request).""" + problems: dict[str, str] = {} + for provider in cfg.providers: + check = _CHECKS.get(provider) + if check is None: + continue + problem = check(getattr(cfg, provider)) + if problem: + problems[provider] = problem + return problems + + +async def _preflight(cfg: Any) -> None: + """Fail fast instead of running every collector against credentials + that cannot work (each one would retry, back off and fail).""" + try: + problems = await asyncio.wait_for( + asyncio.to_thread(credential_problems, cfg), PREFLIGHT_TIMEOUT_S + ) + except asyncio.TimeoutError: + raise PreflightError( + f"Credential check did not finish within {PREFLIGHT_TIMEOUT_S:.0f}s; the cloud " + "identity endpoints look unreachable. Pass preflight=false to try anyway." + ) from None + if problems: + detail = "; ".join(f"{p}: {msg}" for p, msg in problems.items()) + raise PreflightError( + f"Cannot collect live data. {detail}. Fix the server's credentials, choose other " + "providers, or work offline with load_dataset. Pass preflight=false to skip this " + "check.", + data={"providers": problems}, + ) + + +Preflight = Annotated[ + bool, + Field( + description="Check the configured credentials before calling any cloud API, so a missing " + "profile fails in seconds instead of after every collector has retried." + ), +] + + +Force = Annotated[ + bool, + Field( + description="Skip the live cooldown for this call. Only honoured for callers with the " + "admin or operator role; others get an access error." + ), +] +Replace = Annotated[ + bool, + Field( + description="Overwrite an existing dataset with the same name (otherwise a name clash " + "is an error, reported before any cloud call)." + ), +] + + +CATALOG = Catalog() + +_LIVE = dict( + category=CATEGORY, read_only=False, destructive=False, idempotent=False, open_world=True +) + + +ServicesArg = Annotated[ + list[str] | None, + Field( + description="Service families: all, network, compute, containers, kubernetes, " + "serverless, integration, data, storage, identity, security, dns, iac, logging." + ), +] +TaggingArg = Annotated[ + bool | None, + Field(description="AWS Resource Groups Tagging API sweep; default from config."), +] +ScannersArg = Annotated[ + list[str] | None, + Field(description="Subset of prowler, scoutsuite, checkov, trivy, iam; empty = configured."), +] +IacArg = Annotated[str, Field(description="IaC directory for Checkov (inside an allowed root).")] +ImagesArg = Annotated[list[str] | None, Field(description="Container images for Trivy.")] +SubdirArg = Annotated[ + str, + Field(pattern=r"^[A-Za-z0-9_.\-/]{0,200}$", description="Output sub-directory for reports."), +] + + +class MapInventoryArgs(BaseModel): + """map_inventory arguments (the MCP wire schema).""" + + providers: Providers = None + regions: Regions = None + services: ServicesArg = None + tagging_sweep: TaggingArg = None + name: NameArg = "" + replace: Replace = False + preflight: Preflight = True + force: Force = False + + +async def _admit(ctx: Any, cfg: Any, force: bool, preflight: bool) -> None: + """Role check for force, then the credential preflight.""" + _bypass(ctx, force) + if preflight: + await _preflight(cfg) + + +@CATALOG.tool( + title="Map inventory (live)", + sensitivity=Sensitivity.CONFIDENTIAL, + capabilities={R, W, CLOUD}, + timeout_seconds=3600, + **_LIVE, +) +async def map_inventory(ctx: Any, args: MapInventoryArgs) -> dict: + """Map the complete live inventory (no scanners) with read-only cloud + API calls using the server's configured credentials: deep collectors, + catch-all sweeps and relationship linking, across accounts when + organization discovery is configured. Takes minutes. The result + becomes the active dataset; this returns a summary and links. + Guarded: identical concurrent calls share one run, a scope collected + recently is refused with retry_after (use the existing dataset), and + only a few live operations run at once.""" + started = time.monotonic() + cfg = _config(ctx, args.providers, args.regions, services=args.services) + ds_name = _name(ctx, args.name, "inventory", args.replace) + await _admit(ctx, cfg, args.force, args.preflight) + + def extra_fields(result: Any, engine: Any, ds: Dataset) -> dict[str, Any]: + cov = [c.to_summary() for c in ds.coverage] + failures = [f for c in cov for f in c["failures"]][:20] + return {"collection_failures": failures, "errors": engine_errors(engine)} + + run = LiveRun( + "map_inventory", + 2.0, + started, + args.replace, + call=lambda engine: engine.map_inventory(tagging_sweep=args.tagging_sweep), + build=lambda result: Dataset.from_inventory(result, ds_name, source="live"), + extra_fields=extra_fields, + note="mapping inventory", + ) + key = (args.regions, args.services, args.tagging_sweep) + return await _guarded( + ctx, "map", cfg, key, live_scopes(cfg), args.force, lambda: run_live(ctx, cfg, run) + ) + + +@CATALOG.tool( + title="Collect assets (live)", + sensitivity=Sensitivity.CONFIDENTIAL, + capabilities={R, W, CLOUD}, + timeout_seconds=3600, + **_LIVE, +) +async def collect_assets( + ctx: Any, + providers: Providers = None, + regions: Regions = None, + name: NameArg = "", + replace: Replace = False, + preflight: Preflight = True, + force: Force = False, +) -> dict: + """Run the standard multi-provider asset collection (the `cloudg + collect` phase: core services, network edges, IAM). It is lighter than + map_inventory. The result becomes the active dataset. Guarded like + map_inventory (shared runs, per-scope cooldown, concurrency caps).""" + started = time.monotonic() + cfg = _config(ctx, providers, regions) + ds_name = _name(ctx, name, "collect", replace) + await _admit(ctx, cfg, force, preflight) + run = LiveRun( + "collect_assets", + 2.0, + started, + replace, + call=lambda engine: engine.collect(), + build=lambda result: Dataset.from_collection(result, ds_name), + extra_fields=lambda result, engine, ds: {"errors": engine_errors(engine)}, + ) + return await _guarded( + ctx, "collect", cfg, (regions,), live_scopes(cfg), force, lambda: run_live(ctx, cfg, run) + ) + + +@CATALOG.tool( + title="Run scanners", + sensitivity=Sensitivity.CONFIDENTIAL, + capabilities={R, W, Capability.READ_FS, FS, CLOUD, EXEC}, + timeout_seconds=7200, + **_LIVE, +) +async def run_scanners( + ctx: Any, + scanners: ScannersArg = None, + iac_dir: IacArg = "", + images: ImagesArg = None, + normalise: bool = True, + dataset: DatasetArg = "", + preflight: Preflight = True, + force: Force = False, +) -> dict: + """Run security scanners (external binaries, cloud credentials) + plus cloudg's reachability analysis against a dataset's assets and + add the findings to it (deduplicated and compliance-mapped when + normalise=true). Scanner output files go to /scans. + Image references must not start with '-'. When Prowler or ScoutSuite + run, the call is guarded like the collection tools (shared runs, + per-scope cooldown, concurrency caps).""" + started = time.monotonic() + ds = ws_dataset(ctx, dataset) + cfg = _scanner_config(ctx, scanners) + _check_images(images) + iac = ctx.workspace.check_path(iac_dir) if iac_dir else None + scan = ScanJob(ds, cfg, iac, images, ctx.workspace.output_path("scans"), normalise, started) + # Only the cloud scanners need credentials (and cloud scopes); + # checkov / trivy / iam run offline + cloud = bool({"prowler", "scoutsuite"} & {s.lower() for s in cfg.scanners.enabled}) + await _admit(ctx, cfg, force, preflight and cloud) + key = (ds.name, sorted(cfg.scanners.enabled), str(iac) if iac else None, images, normalise) + return await _guarded( + ctx, + "scan", + cfg, + key, + live_scopes(cfg) if cloud else [], + force, + lambda: _run_scan(ctx, scan), + ) + + +@CATALOG.tool( + title="Run full pipeline", + sensitivity=Sensitivity.CONFIDENTIAL, + capabilities={R, W, FS, CLOUD, EXEC}, + timeout_seconds=7200, + **_LIVE, +) +async def run_pipeline( + ctx: Any, + providers: Providers = None, + regions: Regions = None, + subdir: SubdirArg = "pipeline", + name: NameArg = "", + replace: Replace = False, + preflight: Preflight = True, + force: Force = False, +) -> dict: + """The complete cloudg pipeline: collect -> scan -> analyse (graph, + ontology, RAG, Terraform) -> normalise -> JSON + HTML reports. The + result becomes the active dataset; report paths are returned. + Guarded like map_inventory (shared runs, per-scope cooldown, + concurrency caps).""" + started = time.monotonic() + cfg = _config(ctx, providers, regions) + ds_name = _name(ctx, name, "pipeline", replace) + out_dir = ctx.workspace.output_path(subdir) + await _admit(ctx, cfg, force, preflight) + + def extra_fields(result: Any, engine: Any, ds: Dataset) -> dict[str, Any]: + return { + "report_paths": {k: str(v) for k, v in result.report_paths.items()}, + "attack_paths_found": len(result.attack_paths), + "errors": list(result.errors) + engine_errors(engine), + } + + run = LiveRun( + "run_pipeline", + 6.0, + started, + replace, + call=lambda engine: engine.run_pipeline(output_dir=out_dir), + build=lambda result: Dataset.from_pipeline(result, ds_name), + extra_fields=extra_fields, + ) + return await _guarded( + ctx, + "pipeline", + cfg, + (regions, subdir), + live_scopes(cfg), + force, + lambda: run_live(ctx, cfg, run), + ) + + +@CATALOG.tool( + title="Rate limit status", + category=CATEGORY, + sensitivity=Sensitivity.INTERNAL, + capabilities={R}, + read_only=True, + idempotent=True, + open_world=False, +) +def rate_limit_status(ctx: Any) -> dict: + """Live-operation guard and cloud API throttling state, without + calling any cloud API: live operations in flight, scopes cooling + down (seconds left before a new live collection is accepted), the + cooldown and concurrency settings, and the throttling governor's + tripped circuit breakers, slowed rate buckets and recent throttling. + Check it before map_inventory / collect_assets to avoid a refusal.""" + return rate_limit_snapshot(ctx.workspace) + + +@CATALOG.resource( + "cloudg://ratelimit", + title="Rate limit status", + category=CATEGORY, + sensitivity=Sensitivity.INTERNAL, + description="Live-operation guard and cloud API throttling state.", +) +def ratelimit_resource(ctx: Any) -> dict: + return rate_limit_snapshot(ctx.workspace) + + +def rate_limit_snapshot(workspace: Any) -> dict[str, Any]: + from cloudg.resilience import get_governor + + guard = workspace.live_guard + return { + "live_guard": { + **guard.status(), + "cooldown_seconds": guard.cooldown_seconds, + "provider_cooldowns": dict(guard.provider_cooldowns), + "max_concurrent_per_scope": guard.max_concurrent_per_scope, + "max_concurrent_total": guard.max_concurrent_total, + "caller_max_operations": guard.caller_max_operations, + }, + "throttling": get_governor().summary(), + } + + +def register(reg: Registry) -> None: + """Add this module's tools to ``reg``.""" + CATALOG.register(reg) diff --git a/cloudg/mcp/catalog/meta.py b/cloudg/mcp/catalog/meta.py new file mode 100644 index 0000000..8399d69 --- /dev/null +++ b/cloudg/mcp/catalog/meta.py @@ -0,0 +1,293 @@ +"""Meta tools: what this server can do and what cloudg's vocabulary means. + +Static, tenant-free (PUBLIC) answers: the tool catalog by category with +recommended workflows, the schema enums (asset types, edge types, +severities, relation types and groups), and explanations of one asset type +or edge type: how cloudg classifies it, its ontology class, Terraform +mapping, dependency semantics and security relevance. +""" + +from __future__ import annotations + +import functools +import inspect +import re +from typing import Annotated, Any, Literal + +from pydantic import Field + +from cloudg.mcp.catalog._common import Catalog, parse_enum +from cloudg.mcp.core import Registry, Sensitivity +from cloudg.schema.models import ( + AssetType, + CloudProvider, + ComplianceStatus, + EdgeType, + Severity, +) + +CATEGORY = "meta" + +SchemaSection = Literal[ + "all", + "asset_types", + "edge_types", + "severities", + "relation_types", + "relation_groups", + "compliance_statuses", + "providers", +] + +EDGE_READS: dict[str, tuple[str, str]] = { + "REFERENCES": ("uses / points at", "function -> secret, table -> KMS key, queue -> DLQ"), + "ATTACHED_TO": ( + "is attached to", + "instance -> security group, volume -> instance, ENI -> instance", + ), + "ROUTE": ( + "routes through / forwards to", + "route table -> gateway, DNS record -> load balancer, ingress -> service", + ), + "PEERING": ("peers with", "VPC -> peering connection -> VPC"), + "USES_IMAGE": ("runs the image of", "task definition / workload / function -> registry"), + "ASSUMES_ROLE": ( + "runs as", + "function, task, node group, instance profile -> role / service account", + ), + "LOGS_TO": ("sends logs to", "trail, flow log, LB, function -> log destination"), + "LOAD_BALANCER_TARGET": ("sends traffic to", "load balancer -> target group -> instance"), + "GRANTS_ACCESS": ("is granted access to", "principal -> resource its policy names"), + "IAM_POLICY_ATTACHMENT": ("has policy", "user / group / role -> managed policy"), + "IAM_TRUST": ("may assume", "principal / account / provider -> role"), + "CONTAINS": ("contains", "VPC -> subnet, subnet -> instance, account -> resource"), + "INVOKES": ("triggers / calls", "bucket -> function, queue -> function, rule -> target"), + "PROTECTS": ("protects", "WAF -> ALB / API / distribution, firewall -> VPC"), + "MONITORS": ("monitors", "Inspector -> instance, GuardDuty -> account, alarm -> resource"), + "MANAGES": ("manages", "stack -> resource, ASG -> instance, landing zone -> account"), + "GOVERNS": ("governs", "SCP / control / org policy -> OU, account, folder, project"), + "SECURITY_GROUP_RULE": ("allows traffic", "CIDR -> SG (ingress), SG -> CIDR (egress)"), + "NACL_RULE": ("allows traffic", "network ACL rules"), + "INTERNET_EXPOSED": ("exposes", "0.0.0.0/0 -> exposed asset"), +} + +# Edge types whose ontology relations come from a rule over the edge's +# details rather than a fixed list (cloudg.graph.ontology_rules). +RULE_RELATIONS: dict[str, str] = { + "SECURITY_GROUP_RULE": "inferred from ports, protocol, CIDR and direction " + "(INGRESS_ALLOWED / EGRESS_ALLOWED, ONLY_SSH, ONLY_HTTPS, ALL_TRAFFIC, " + "INTERNET_REACHABLE...)", + "CONTAINS": "inferred from both endpoint types (VPC_CONTAINS_SUBNET, " + "SUBNET_CONTAINS_INSTANCE, CLUSTER_CONTAINS_SERVICE, ORG_CONTAINS_ACCOUNT), else the " + "generic CONTAINS", + "ATTACHED_TO": "PROTECTED_BY_SG when the target is a security group or NSG, " + "PROTECTED_BY_NACL for a NACL, else DEPENDS_ON", + "IAM_TRUST": "ROLE_ASSUMES_ROLE, plus CROSS_ACCOUNT_TRUST when the accounts differ", + "NACL_RULE": "none; protection relations come from ATTACHED_TO edges", + "IAM_POLICY_ATTACHMENT": "inferred from the principal type (USER_HAS_POLICY, " + "ROLE_HAS_POLICY, GROUP_HAS_POLICY)", +} + +WORKFLOWS = { + "orient": ["workspace_status", "dataset_summary", "count_assets", "findings_summary"], + "triage risk": ["top_risks", "get_asset", "findings_for_asset", "blast_radius"], + "attack surface": [ + "internet_exposure", + "attack_paths", + "lateral_movement_paths", + "find_paths(source='internet', target=...)", + ], + "change impact": ["dependents", "dependency_tree", "shared_dependencies", "blast_radius"], + "compliance": ["compliance_summary", "list_controls", "control_status", "compliance_gaps"], + "drift": ["snapshot_dataset", "map_inventory / load_dataset", "diff_datasets"], + "semantic": ["ontology_stats", "relation_groups", "ontology_neighbourhood", "sparql_query"], +} + + +# ``NAME = "VALUE"`` (optionally parenthesised) plus an optional trailing +# comment. Each optional part starts with a distinct character, so the +# match stays linear in the line length. +_MEMBER_LINE = re.compile(r'([A-Z_0-9]+)\s*=\s*(?:\(\s*)?"[A-Z_0-9]+"(?:\s*\))?\s*(?:#(.*))?$') +_SUPPRESSION_MARKERS = ("nosec", "nosemgrep") + + +def _strip_suppression(note: str) -> str: + """``note`` without a trailing scanner suppression marker.""" + cuts = [i for i in (note.find(m) for m in _SUPPRESSION_MARKERS) if i >= 0] + return (note[: min(cuts)] if cuts else note).strip().strip(" -#") + + +@functools.lru_cache(maxsize=1) +def enum_comments() -> dict[str, dict[str, str]]: + """Section headings and inline comments of the AssetType / EdgeType + enums, read from their source (they document the taxonomy).""" + out: dict[str, dict[str, str]] = {"family": {}, "asset_note": {}, "edge_note": {}} + for enum, note_key in ((AssetType, "asset_note"), (EdgeType, "edge_note")): + try: + src = inspect.getsource(enum) + except (OSError, TypeError): # pragma: no cover (source unavailable) + continue + section = "Other" + for line in src.splitlines(): + s = line.strip() + if s.startswith("#") and s[1:].strip() and enum is AssetType: + section = s[1:].strip() + continue + m = _MEMBER_LINE.match(s) + if m: + name, note = m.group(1), _strip_suppression(m.group(2) or "") + if enum is AssetType: + out["family"][name] = section + if note: + out[note_key][name] = note + return out + + +def _relation_types() -> dict[str, list[str]]: + from cloudg.graph.ontology_rules import RelationGroup, get_relations_for_group + + return {g.value: [r.value for r in get_relations_for_group(g)] for g in RelationGroup} + + +CATALOG = Catalog() + +_RO = dict( + category=CATEGORY, + sensitivity=Sensitivity.PUBLIC, + capabilities=(), + read_only=True, + idempotent=True, + open_world=False, +) + + +@CATALOG.tool(title="List capabilities", tags={"start-here"}, **_RO) +def list_capabilities(ctx: Any) -> dict: + """What this server offers: every tool you may call grouped by + category (with read-only / sensitivity / required capabilities), the + resource and prompt counts, and recommended tool sequences for + common jobs (orienting, risk triage, attack surface, change impact, + compliance, drift, semantic queries).""" + layer = ctx.layer + cats: dict[str, list[dict[str, Any]]] = {} + for spec in layer.list_tools(ctx.principal): + cats.setdefault(spec.category, []).append( + { + "name": layer.exposed_name(spec.name), + "title": spec.title, + "read_only": spec.annotations.read_only, + "open_world": spec.annotations.open_world, + "sensitivity": spec.sensitivity.value, + "capabilities": sorted(c.value for c in spec.capabilities), + "summary": (spec.description.split(". ")[0].strip()[:160]), + } + ) + return { + "server": layer.name, + "version": layer.version, + "categories": {k: sorted(v, key=lambda t: t["name"]) for k, v in sorted(cats.items())}, + "tool_count": sum(len(v) for v in cats.values()), + "resources": len(layer.list_resources(ctx.principal)), + "resource_templates": len(layer.list_resource_templates(ctx.principal)), + "prompts": [layer.exposed_name(p.name) for p in layer.list_prompts(ctx.principal)], + "workflows": WORKFLOWS, + "conventions": { + "dataset": "empty = active dataset", + "ref": "asset id, ARN, unique name or unique ARN tail", + "pagination": "pass next_cursor back as cursor; cursors expire when the " + "dataset changes", + }, + } + + +@CATALOG.tool(title="Describe schema", **_RO) +def describe_schema(ctx: Any, section: SchemaSection = "all") -> dict: + """cloudg's vocabulary: asset types (grouped by family), edge types + (with how each reads), severities, ontology relation types by group, + compliance statuses and providers. Use these exact values in filters.""" + notes = enum_comments() + out: dict[str, Any] = {} + if section in ("all", "asset_types"): + fam: dict[str, list[str]] = {} + for t in AssetType: + fam.setdefault(notes["family"].get(t.name, "Other"), []).append(t.value) + out["asset_types"] = fam + if section in ("all", "edge_types"): + out["edge_types"] = {e.value: EDGE_READS.get(e.value, ("", ""))[0] for e in EdgeType} + if section in ("all", "severities"): + out["severities"] = [s.value for s in Severity] + if section in ("all", "relation_types", "relation_groups"): + rel = _relation_types() + out["relation_groups"] = list(rel) if section == "relation_groups" else rel + if section in ("all", "compliance_statuses"): + out["compliance_statuses"] = [s.value for s in ComplianceStatus] + if section in ("all", "providers"): + out["providers"] = [p.value for p in CloudProvider] + return out + + +@CATALOG.tool(title="Explain asset type", **_RO) +def explain_asset_type( + ctx: Any, + asset_type: Annotated[ + str, Field(min_length=1, description="An AssetType value, e.g. LAMBDA_FUNCTION.") + ], +) -> dict: + """How cloudg models one asset type: family, notes, OWL ontology + class, Terraform resource type, whether it is a sensitive data store + / crown jewel, whether internet exposure is expected, and how many + the active dataset holds.""" + from cloudg.graph.ontology import _ASSET_TYPE_CLASSES + from cloudg.graph.reachability import EXPECTED_EXPOSED_TYPES, SENSITIVE_ASSET_TYPES + from cloudg.mcp.catalog.graph import CROWN_JEWELS + from cloudg.renderers.terraform_export import _TF_RESOURCE_MAP + + t = parse_enum(asset_type, AssetType, "asset type") + notes = enum_comments() + out: dict[str, Any] = { + "asset_type": t.value, + "family": notes["family"].get(t.name, "Other"), + "note": notes["asset_note"].get(t.name), + "ontology_class": f"cm:{_ASSET_TYPE_CLASSES.get(t, 'CloudResource')}", + "terraform_type": _TF_RESOURCE_MAP.get(t), + "sensitive_data_store": t in SENSITIVE_ASSET_TYPES, + "crown_jewel_weight": CROWN_JEWELS.get(t, 0), + "internet_exposure_expected": t in EXPECTED_EXPOSED_TYPES, + } + ws = ctx.workspace + if ws.active_name: + ds = ws.get() + out["in_active_dataset"] = sum(1 for a in ds.assets if a.asset_type == t) + return out + + +@CATALOG.tool(title="Explain edge type", **_RO) +def explain_edge_type( + ctx: Any, + edge_type: Annotated[ + str, Field(min_length=1, description="An EdgeType value, e.g. ASSUMES_ROLE.") + ], +) -> dict: + """How to read one edge type: direction ("source verb target"), + typical endpoints, its dependency direction (forward: source depends + on target; reverse: target depends on source; none), and the + ontology relations inferred from it.""" + from cloudg.graph.ontology_rules import _SIMPLE_EDGE_RELATIONS + from cloudg.inventory.dependencies import DEPENDENCY_DIRECTION + + e = parse_enum(edge_type, EdgeType, "edge type") + reads, typical = EDGE_READS.get(e.value, ("", "")) + return { + "edge_type": e.value, + "reads_as": f"source {reads} target" if reads else None, + "typical": typical or None, + "note": enum_comments()["edge_note"].get(e.name), + "dependency_direction": DEPENDENCY_DIRECTION.get(e.value, "none"), + "ontology_relations": [r.value for r in _SIMPLE_EDGE_RELATIONS.get(e, [])] + or RULE_RELATIONS.get(e.value, "none inferred from the edge type alone"), + } + + +def register(reg: Registry) -> None: + """Add this module's tools to ``reg``.""" + CATALOG.register(reg) diff --git a/cloudg/mcp/catalog/ontology.py b/cloudg/mcp/catalog/ontology.py new file mode 100644 index 0000000..17fa35f --- /dev/null +++ b/cloudg/mcp/catalog/ontology.py @@ -0,0 +1,352 @@ +"""Ontology and RAG tools: the RDF/OWL knowledge graph cloudg builds from a +dataset (~60 typed relations in 7 groups), read-only SPARQL over it, an +asset's semantic neighbourhood, relation-group views, RAG-ready chunks and +ontology export. + +The ontology is built lazily on first use and cached per dataset; on large +maps the first call can take a while. SPARQL vetting and its result bounds +live in :mod:`cloudg.mcp.catalog._sparql`. +""" + +from __future__ import annotations + +from collections import deque +from typing import Annotated, Any, Literal + +from pydantic import Field + +from cloudg.mcp.catalog._common import ( + Catalog, + Cursor, + DatasetArg, + RefArg, + SparqlOut, + asset_uri, + fingerprint, + paginate, + parse_enum, + parse_severity, + schema, + ws_dataset, +) +from cloudg.mcp.catalog._sparql import compact, namespaces, prepare_readonly, run_query +from cloudg.mcp.core import Capability, Registry, Sensitivity +from cloudg.mcp.state import SEVERITY_RANK + +__all__ = ["CATALOG", "compact", "namespaces", "prepare_readonly", "register"] + +CATEGORY = "ontology" +R, FS = Capability.READ_STATE, Capability.WRITE_FS + +OntologyFormat = Literal["turtle", "json-ld", "xml", "nt"] +ChunkType = Literal["entity", "community", "relation_group"] +FORMAT_EXT = {"turtle": "ttl", "json-ld": "jsonld", "xml": "rdf", "nt": "nt"} +FORMAT_MIME = { + "turtle": "text/turtle", + "json-ld": "application/ld+json", + "xml": "application/rdf+xml", + "nt": "application/n-triples", +} + + +def _ontology(ctx: Any, ds: Any) -> Any: + """The dataset's ontology, reporting progress while building it.""" + if not ds.has_ontology(): + sync = getattr(ctx, "report_progress_sync", None) + if sync: + sync(0, 1, f"building ontology for {len(ds.assets)} assets") + onto = ds.ontology() + if sync: + sync(1, 1, "ontology built") + return onto + return ds.ontology() + + +def _neighbour_edges(g: Any, node: Any) -> list[tuple[Any, Any, Any, Any]]: + """(subject, predicate, object, other end) of every relation at ``node``.""" + from rdflib import URIRef + + # the schema's data properties are literals, not relations + from cloudg.graph.ontology import _DATA_PROPERTIES, CMP + + cmp_ns = str(CMP) + data_props = frozenset(_DATA_PROPERTIES) + edges = [(node, p, o, o) for p, o in g.predicate_objects(node)] + edges += [(s, p, node, s) for s, p in g.subject_predicates(node)] + return [ + (s, p, o, other) + for s, p, o, other in edges + if str(p).startswith(cmp_ns) + and str(p)[len(cmp_ns) :] not in data_props + and isinstance(other, URIRef) + ] + + +def _neighbourhood(g: Any, start: Any, hops: int, limit: int) -> tuple[list[dict], bool]: + """Breadth-first relation triples around ``start`` (at most ``limit``).""" + from cloudg.graph.ontology import CMP + + cmp_len = len(str(CMP)) + name_p = CMP["hasName"] + + def name(node: Any) -> str: + n = g.value(node, name_p) + return str(n) if n is not None else compact(node) + + seen = {start} + triples: list[dict[str, Any]] = [] + queue: deque[tuple[Any, int]] = deque([(start, 0)]) + while queue: + node, depth = queue.popleft() + if depth >= hops: + continue + for s, p, o, other in _neighbour_edges(g, node): + if len(triples) >= limit: + return triples, True + triples.append( + { + "subject": name(s), + "predicate": str(p)[cmp_len:], + "object": name(o), + "depth": depth + 1, + "subject_id": compact(s), + "object_id": compact(o), + } + ) + if other not in seen: + seen.add(other) + queue.append((other, depth + 1)) + return triples, False + + +CATALOG = Catalog() + +_RO = dict(read_only=True, idempotent=True, open_world=False, category=CATEGORY) + + +@CATALOG.tool(title="Ontology statistics", sensitivity=Sensitivity.INTERNAL, **_RO) +def ontology_stats(ctx: Any, dataset: DatasetArg = "") -> dict: + """Size and shape of the dataset's RDF ontology: triples, classes, + individuals, and counts per relation type and relation group (builds + the ontology on first call).""" + ds = ws_dataset(ctx, dataset) + built = not ds.has_ontology() + stats = _ontology(ctx, ds).stats() + ctx.link("cloudg://ontology/turtle", "ontology (turtle)", mime_type="text/turtle") + return {"dataset": ds.name, "built_now": built, **stats} + + +# RESTRICTED: rows are keyed by the query's own variable names, which the +# privacy layer cannot classify, so strict ceilings hide the tool. +@CATALOG.tool( + title="SPARQL query", + sensitivity=Sensitivity.RESTRICTED, + output_schema=schema(SparqlOut), + timeout_seconds=120, + **_RO, +) +def sparql_query( + ctx: Any, + query: Annotated[ + str, + Field( + min_length=1, + max_length=20_000, + description=( + "Read-only SPARQL (SELECT / ASK / CONSTRUCT / DESCRIBE). Predefined prefixes: " + "cm: (classes, e.g. cm:RelationalDatabase), cmp: (relations / properties, e.g. " + "cmp:INTERNET_REACHABLE, cmp:hasName), cmr: (individuals, cmr:), rdf:, " + "rdfs:, owl:, xsd:. FROM, SERVICE and UPDATE are rejected." + ), + ), + ], + limit: Annotated[int, Field(ge=1, le=1000)] = 100, + compact_uris: Annotated[bool, Field(description="Shorten URIs to prefix:name.")] = True, + dataset: DatasetArg = "", +) -> dict: + """Run a read-only SPARQL query against the dataset's ontology. + Example (internet-reachable databases): + SELECT ?r ?name WHERE { ?s cmp:INTERNET_REACHABLE ?r . ?r a + cm:RelationalDatabase ; cmp:hasName ?name } + + At most `limit` rows are read; a query whose patterns multiply out + stops there. variable_kinds says whether each column holds URIs or + literals.""" + prepared = prepare_readonly(query) + ds = ws_dataset(ctx, dataset) + g = _ontology(ctx, ds).graph + return {"dataset": ds.name, **run_query(g, prepared, limit, compact_uris)} + + +@CATALOG.tool(title="Ontology neighbourhood", sensitivity=Sensitivity.CONFIDENTIAL, **_RO) +def ontology_neighbourhood( + ctx: Any, + ref: RefArg, + hops: Annotated[int, Field(ge=1, le=3)] = 1, + limit: Annotated[int, Field(ge=1, le=500)] = 100, + dataset: DatasetArg = "", +) -> dict: + """Semantic relations around an asset in the ontology (e.g. + PROTECTED_BY_SG, ENCRYPTED_BY_KMS, INTERNET_REACHABLE, + FINDING_AFFECTS, TAGGED_WITH), up to `hops` away, as + subject-predicate-object triples with names.""" + from cloudg.graph.ontology import CMR + + ds = ws_dataset(ctx, dataset) + a = ds.resolve_asset(ref) + g = _ontology(ctx, ds).graph + triples, truncated = _neighbourhood(g, CMR[a.id], hops, limit) + preds: dict[str, int] = {} + for t in triples: + preds[t["predicate"]] = preds.get(t["predicate"], 0) + 1 + ctx.link(asset_uri(a.id), a.name) + return { + "dataset": ds.name, + "asset": {"id": a.id, "name": a.name, "type": a.asset_type.value}, + "predicates": preds, + "triples": triples, + "returned": len(triples), + "truncated": truncated, + } + + +@CATALOG.tool(title="Relation groups", sensitivity=Sensitivity.CONFIDENTIAL, **_RO) +def relation_groups( + ctx: Any, + group: Annotated[ + str, + Field( + description="NETWORK, CONTAINMENT, IAM, DATA_FLOW, " + "SECURITY, COMPUTE or GOVERNANCE; empty = overview." + ), + ] = "", + limit: Annotated[int, Field(ge=1, le=500)] = 100, + cursor: Cursor = "", + dataset: DatasetArg = "", +) -> dict: + """Without group: counts of every semantic relation group and type in + the dataset's ontology. With group: the triples of that group (e.g. + every IAM relation: ROLE_ASSUMES_ROLE, CROSS_ACCOUNT_TRUST...), + paginated.""" + from cloudg.graph.ontology import RelationGroup, get_relations_for_group + + ds = ws_dataset(ctx, dataset) + onto = _ontology(ctx, ds) + if not group: + stats = ds.cached("ontology_stats", onto.stats) + groups = {} + for g in RelationGroup: + types = { + rt.value: stats["relation_type_counts"].get(rt.value, 0) + for rt in get_relations_for_group(g) + } + groups[g.value] = { + "total": stats["relation_group_counts"].get(g.value, 0), + "types": {k: v for k, v in types.items() if v}, + } + return {"dataset": ds.name, "groups": groups} + grp = parse_enum(group, RelationGroup, "relation group") + rows = ds.cached(f"relgroup:{grp.value}", lambda: onto.query_by_relation_group(grp)) + rows = [{k: compact(v) if k in ("src", "tgt") else v for k, v in r.items()} for r in rows] + page, env = paginate(rows, limit, cursor, fingerprint(ds, g=grp.value)) + return { + "dataset": ds.name, + "group": grp.value, + "relation_types": [rt.value for rt in get_relations_for_group(grp)], + **env, + "items": page, + } + + +@CATALOG.tool(title="RAG chunks", sensitivity=Sensitivity.CONFIDENTIAL, **_RO) +def rag_chunks( + ctx: Any, + chunk_type: ChunkType = "entity", + ref: Annotated[str, Field(description="Entity chunks: only this asset's chunk.")] = "", + query: Annotated[str, Field(description="Substring of chunk content.")] = "", + min_severity: Annotated[ + str, Field(description="Entity chunks: severity_max at or above.") + ] = "", + limit: Annotated[int, Field(ge=1, le=100)] = 10, + cursor: Cursor = "", + dataset: DatasetArg = "", +) -> dict: + """Retrieval-ready text chunks with metadata: entity (one per asset + with its relations and findings), community (Louvain clusters of + tightly-connected resources with risk), relation_group (all triples + of a semantic group). Good context for summarising an area.""" + ds = ws_dataset(ctx, dataset) + q = query.lower() + aid = ds.resolve_asset(ref).id if ref else None + rank = SEVERITY_RANK[parse_severity(min_severity).value] if min_severity else None + sel = [ + c + for c in ds.rag_chunks(chunk_type) + if (not aid or c.chunk_id == f"entity::{aid}") + and (not q or q in c.content.lower()) + and (rank is None or SEVERITY_RANK.get(c.metadata.get("severity_max", ""), -1) >= rank) + ] + page, env = paginate( + sel, limit, cursor, fingerprint(ds, t=chunk_type, r=aid, q=q, s=min_severity) + ) + return { + "dataset": ds.name, + "chunk_type": chunk_type, + **env, + "items": [ + { + "chunk_id": c.chunk_id, + "chunk_type": c.chunk_type, + "content": c.content, + "metadata": c.metadata, + } + for c in page + ], + } + + +@CATALOG.tool( + title="Export ontology", + category=CATEGORY, + sensitivity=Sensitivity.CONFIDENTIAL, + capabilities={R, FS}, + read_only=False, + destructive=False, + idempotent=True, + open_world=False, +) +def export_ontology( + ctx: Any, + format: OntologyFormat = "turtle", # pylint: disable=redefined-builtin # MCP argument name + filename: Annotated[ + str, + Field( + pattern=r"^[A-Za-z0-9_.\-]{1,100}$", + description="File name (no directories); extension added if missing.", + ), + ] = "ontology", + dataset: DatasetArg = "", +) -> dict: + """Write the dataset's ontology to the workspace output directory as + Turtle, JSON-LD, RDF/XML or N-Triples. Returns the path; read the + cloudg://ontology/{format} resource instead to get the content + without writing a file.""" + ds = ws_dataset(ctx, dataset) + ext = FORMAT_EXT[format] + name = filename if filename.endswith(f".{ext}") else f"{filename}.{ext}" + path = ctx.workspace.output_path(name) + onto = _ontology(ctx, ds) + onto.save(path, fmt=format) + ctx.link(f"cloudg://ontology/{format}", f"ontology ({format})", mime_type=FORMAT_MIME[format]) + return { + "dataset": ds.name, + "path": str(path), + "format": format, + "triples": onto.triple_count, + "bytes": path.stat().st_size, + } + + +def register(reg: Registry) -> None: + """Add this module's tools to ``reg``.""" + CATALOG.register(reg) diff --git a/cloudg/mcp/catalog/privacy.py b/cloudg/mcp/catalog/privacy.py new file mode 100644 index 0000000..df06b77 --- /dev/null +++ b/cloudg/mcp/catalog/privacy.py @@ -0,0 +1,424 @@ +"""Privacy tools and resources: inspect the active policy, preview what it +does to data, list detectors, read the audit trail, and (for privileged +roles only) reverse pseudonyms. + +===================== ============= ============================================== +primitive sensitivity notes +===================== ============= ============================================== +``privacy_status`` internal active profile, denied capabilities, counts +``preview_transform`` internal run the policy over sample JSON / text, with + a throw-away vault (hidden by ``strict``) +``list_detectors`` public detectors + the strategy each gets here +``privacy_audit_log`` restricted recent access decisions (hashed arguments); + only the caller's own entries unless it has + the ``admin`` or ``privacy-admin`` role +``reveal_token`` restricted needs ``Capability.REVEAL``; denied by + ``standard`` except admin roles, always by + ``strict``; every reveal is audited +``cloudg://policy`` internal ``policy.describe()`` (never the vault key) +``cloudg://privacy/detectors`` public +===================== ============= ============================================== + +Handlers reach the policy through ``ctx.layer.policy``. +""" + +from __future__ import annotations + +import json +import logging +from typing import Annotated, Any + +from pydantic import Field + +from cloudg.mcp.core import Capability, Registry, Sensitivity +from cloudg.mcp.policy.audit import log_safe + +CATEGORY = "privacy" +audit_logger = logging.getLogger("cloudg.mcp.audit") +#: Roles that may read every principal's audit entries. +AUDIT_ADMIN_ROLES = frozenset({"admin", "privacy-admin"}) + + +def _policy(ctx: Any) -> Any: + return ctx.layer.policy + + +def _redactor_for(policy: Any, principal: Any) -> Any: + from cloudg.mcp.transforms import Redactor + + for t in policy._generic_pipeline(principal).transforms: + if isinstance(t, Redactor): + return t + return None + + +def detector_table( + policy: Any, principal: Any = None, category: str | None = None +) -> dict[str, Any]: + """Detectors and key rules with the strategy the policy applies.""" + from cloudg.mcp.transforms.detectors import ENTITY_PARENTS + from cloudg.mcp.transforms.redaction import _chain + + red = _redactor_for(policy, principal) + rows = [] + for d in policy.detectors(): + if category and d.category != category: + continue + row = { + "name": d.name, + "entity": d.entity, + "category": d.category, + "confidence": d.confidence, + "description": d.description, + "enabled": d.enabled, + } + if red is not None: + row["strategy"] = red.resolve(_chain(d.entity) + (d.name, d.category)).kind + # a list, so entity names are values rather than keys that key + # rules (account, name...) would act on + refined = [ + { + "entity": e, + "strategy": red.resolve(_chain(e) + (d.entity, d.name, d.category)).kind, + } + for e, p in ENTITY_PARENTS.items() + if p == d.entity + ] + if refined: + row["refined"] = refined + row["active"] = d.name in red.active_detectors + rows.append(row) + key_rules = [] + for r in policy.detectors().key_rules: + if category and r.category != category: + continue + item = { + "name": r.name, + "entity": r.entity, + "category": r.category, + "key_pattern": r.pattern, + } + if red is not None: + item["strategy"] = red._rule_strategy(r).kind + key_rules.append(item) + return {"policy": policy.name, "detectors": rows, "key_rules": key_rules} + + +# --------------------------------------------------------------------------- +# Handlers +# --------------------------------------------------------------------------- + + +def privacy_status(ctx: Any) -> dict[str, Any]: + """Show the active privacy policy for the calling principal: profile, + denied capabilities, sensitivity ceiling, what the output pipeline + does (redaction / pseudonymisation / projection), pseudonym vault + counts and access-decision counters. Identifiers in results may be + pseudonyms: pass them back to tools unchanged.""" + policy = _policy(ctx) + cfg = policy.config + principal = ctx.principal + pipeline = policy._generic_pipeline(principal) + steps = [] + for t in pipeline.transforms: + desc = t.describe() if hasattr(t, "describe") else {"type": t.name} + steps.append(desc if isinstance(desc, dict) else {"type": t.name}) + vault = policy.vault.stats() + return { + "policy": policy.name, + "description": cfg.description, + "loaded_from": policy.source, + "extends": policy.extends_chain, + "principal": {"id": principal.id, "roles": sorted(principal.roles)}, + "max_sensitivity": cfg.max_sensitivity.value if cfg.max_sensitivity else None, + "denied_capabilities": sorted(c.value for c in policy._base.deny_caps), + "visible_tools": len(ctx.layer.list_tools(principal)), + "output_pipeline": steps, + "pseudonymisation": { + "active": any(getattr(t, "reversible", False) for t in pipeline.transforms), + "vault_scope": vault["scope"], + "vault_entries": vault["entries"], + "by_entity": vault["by_entity"], + "key_source": vault["key_source"], + }, + "audit_enabled": cfg.audit, + "counters": dict(policy.counters), + } + + +def preview_transform( + ctx: Any, + data: Annotated[ + Any, + Field(description="JSON value (object / array / string) or text to transform"), + ], + tool: Annotated[ + str | None, + Field( + description="Preview with the pipeline of this tool (default: the policy's " + "generic pipeline)" + ), + ] = None, + parse_json: Annotated[ + bool, Field(description="Parse `data` as JSON when it is a JSON string") + ] = True, +) -> dict[str, Any]: + """Show how the active policy would transform some data for you: + which values get redacted, masked, hashed, pseudonymised, projected + or fenced as untrusted, plus the transform report. Pseudonyms in a + preview come from a separate throw-away vault, so they never match the + pseudonyms in real results; nothing is stored.""" + policy = _policy(ctx) + spec = ctx.layer.get_tool(tool, ctx.principal) if tool else None + value = data + if parse_json and isinstance(data, str) and data.strip()[:1] in ("{", "["): + try: + value = json.loads(data) + except ValueError: + value = data + transformed, report = policy.preview(value, spec, ctx.principal) + return { + "policy": policy.name, + "tool": tool, + "transformed": transformed, + "report": report, + } + + +def list_detectors( + ctx: Any, + category: Annotated[ + str | None, + Field( + description="Only this category: secret, credential, identifier, network, " + "pii, temporal, free_text" + ), + ] = None, +) -> dict[str, Any]: + """List the sensitive-entity detectors (content regexes and key-name + rules) and the strategy the active policy applies to each + (redact / mask / hash / pseudonymize / generalize / drop / keep).""" + return detector_table(_policy(ctx), ctx.principal, category) + + +def privacy_audit_log( + ctx: Any, + limit: Annotated[int, Field(ge=1, le=1000, description="Most recent entries")] = 50, + decision: Annotated[ + str | None, Field(description="Filter: allowed, denied or rate_limited") + ] = None, +) -> dict[str, Any]: + """Recent access decisions recorded by the policy (denials and + rate-limit hits always; every call when the policy enables audit). + Arguments appear as field names with keyed hashes, never values. Callers + without the admin or privacy-admin role see only their own entries.""" + policy = _policy(ctx) + principal = ctx.principal + see_all = bool(AUDIT_ADMIN_ROLES & set(getattr(principal, "roles", ()) or ())) + entries = list(policy.audit_log) + if not see_all: + pid = str(getattr(principal, "id", "anonymous")) + entries = [e for e in entries if e.get("principal") == pid] + if decision: + entries = [e for e in entries if e.get("decision") == decision] + return { + "policy": policy.name, + "audit_enabled": policy.config.audit, + "scope": "all" if see_all else "own", + "total": len(entries), + "entries": entries[-limit:], + } + + +def _alias_maps(ctx: Any, policy: Any) -> list[Any]: + """Alias tables in the output pipelines of the tools the caller can + see, i.e. the aliases that may appear in values they hold.""" + from cloudg.mcp.transforms.substitution import AliasMap, Substitution + + found: dict[int, Any] = {} + for spec in ctx.layer.list_tools(ctx.principal): + for t in policy.output_pipeline(spec, ctx.principal).transforms: + maps = ( + t.alias_maps + if isinstance(t, Substitution) + else ([t] if isinstance(t, AliasMap) and t.aliases else []) + ) + for m in maps: + found.setdefault(id(m), m) + return list(found.values()) + + +def _log_reveal(ctx: Any, vault: Any, token: str, detail: str) -> None: + """Audit one reveal: who, what kind, and a keyed hash of the input. + Neither the pseudonym nor the revealed value is logged.""" + audit_logger.warning( + "reveal_token principal=%s %s token_hash=%s", + log_safe(ctx.principal.id), + log_safe(detail), + vault.hash_value(token, entity_type="audit"), + ) + + +def _reveal_aliased(ctx: Any, vault: Any, token: str, aliases: list[Any]) -> dict[str, Any]: + """Aliased values (soc-analyst account names) wrap vault tokens: + reverse aliases and tokens together, as tool inputs do.""" + from cloudg.mcp.transforms.base import TransformContext + from cloudg.mcp.transforms.substitution import Depseudonymizer + + tctx = TransformContext(principal=ctx.principal, direction="input", vault=vault) + real = Depseudonymizer(vault, aliases=aliases).apply(token, tctx) + n = int(tctx.report.get("depseudonymized", 0)) + _log_reveal(ctx, vault, token, f"reversed={n}") + return {"pseudonym": token, "found": n > 0, "value": real if n else None, "replaced": n} + + +def reveal_token( + ctx: Any, + token: Annotated[ + str, + Field( + description="A pseudonym (e.g. an account ID, ARN or IP from a result) or " + "text containing pseudonyms", + min_length=1, + max_length=20000, + ), + ], +) -> dict[str, Any]: + """Reverse pseudonymisation: return the real value behind a pseudonym + (or replace every pseudonym inside a text). Privileged: it requires the + `reveal` capability, and every call is audited.""" + policy = _policy(ctx) + vault = policy.vault + ns = vault.namespace_for(ctx.principal) + entry = vault.entry(token, namespace=ns) + if entry is None: + aliases = _alias_maps(ctx, policy) + if aliases: + return _reveal_aliased(ctx, vault, token, aliases) + text, n = vault.detokenize_text_count(token, namespace=ns) + _log_reveal(ctx, vault, token, f"embedded_tokens={n}") + return {"pseudonym": token, "found": n > 0, "value": text if n else None, "replaced": n} + _log_reveal(ctx, vault, token, f"entity={entry['entity_type']}") + return { + "pseudonym": token, + "found": True, + "value": entry["value"], + "entity_type": entry["entity_type"], + } + + +def policy_resource(ctx: Any) -> dict[str, Any]: + return _policy(ctx).describe() + + +def detectors_resource(ctx: Any) -> dict[str, Any]: + return detector_table(_policy(ctx), ctx.principal) + + +# --------------------------------------------------------------------------- +# Registration +# --------------------------------------------------------------------------- + +_READ = {"read_only": True, "open_world": False} + +#: (handler, registration options) of every privacy tool. +_TOOLS: tuple[tuple[Any, dict[str, Any]], ...] = ( + ( + privacy_status, + { + "name": "privacy_status", + "title": "Privacy policy status", + "tags": ("privacy", "policy"), + "sensitivity": Sensitivity.INTERNAL, + "idempotent": True, + }, + ), + ( + preview_transform, + { + "name": "preview_transform", + "title": "Preview privacy transforms", + "tags": ("privacy", "policy"), + "sensitivity": Sensitivity.INTERNAL, + "idempotent": True, + # The result is already transformed by the policy; do not + # transform it again, and do not reverse pseudonyms in the sample. + "transform_hints": {"skip": ["*"], "depseudonymize": False, "input_guard": False}, + }, + ), + ( + list_detectors, + { + "name": "list_detectors", + "title": "List sensitive-data detectors", + "tags": ("privacy",), + "sensitivity": Sensitivity.PUBLIC, + "idempotent": True, + }, + ), + ( + privacy_audit_log, + { + "name": "privacy_audit_log", + "title": "Privacy audit log", + "tags": ("privacy", "audit"), + "sensitivity": Sensitivity.RESTRICTED, + }, + ), + ( + reveal_token, + { + "name": "reveal_token", + "title": "Reveal a pseudonym", + "tags": ("privacy", "reveal"), + "sensitivity": Sensitivity.RESTRICTED, + "capabilities": (Capability.REVEAL,), + "destructive": False, + "idempotent": True, + # Must see the token itself and return the real value untouched + "transform_hints": { + "pseudonymize": False, + "depseudonymize": False, + "input_guard": False, + "skip": ["substitute", "alias"], + }, + }, + ), +) + +#: (handler, URI, registration options) of every privacy resource. +_RESOURCES: tuple[tuple[Any, str, dict[str, Any]], ...] = ( + ( + policy_resource, + "cloudg://policy", + { + "name": "privacy_policy", + "title": "Active privacy policy", + "description": "The active access / transform policy (profiles, rules, transforms, " + "rate limits, vault statistics; never the vault key).", + "sensitivity": Sensitivity.INTERNAL, + }, + ), + ( + detectors_resource, + "cloudg://privacy/detectors", + { + "name": "privacy_detectors", + "title": "Sensitive-data detectors", + "description": "Detectors and key rules with the strategy the policy applies to each.", + "sensitivity": Sensitivity.PUBLIC, + }, + ), +) + + +def register(registry: Registry) -> Registry: + """Add the privacy tools and resources to ``registry``.""" + for fn, options in _TOOLS: + registry.tool(category=CATEGORY, **_READ, **options)(fn) + for fn, uri, options in _RESOURCES: + registry.resource(uri, category=CATEGORY, **options)(fn) + return registry + + +__all__ = ["AUDIT_ADMIN_ROLES", "CATEGORY", "detector_table", "register"] diff --git a/cloudg/mcp/catalog/prompts.py b/cloudg/mcp/catalog/prompts.py new file mode 100644 index 0000000..ef80b8f --- /dev/null +++ b/cloudg/mcp/catalog/prompts.py @@ -0,0 +1,587 @@ +"""Prompts: packaged security-analysis jobs. + +Each prompt returns a user message with the task and the exact tools to +call next, then the relevant workspace context so the model starts +informed without a dozen tool calls: embedded resources whose URI serves +exactly the embedded content (dataset summary, asset, finding, framework) +and, for computed context that has no resource (top risks, diffs), a +second text message with a JSON block. With no dataset loaded the prompt +says how to load one instead of failing. + +The task text never contains workspace values (dataset, asset, finding or +framework names, account ids, the caller's own arguments): it points at +the context instead. Every context message carries its data in +``meta["cloudg/data"]`` with the rendering rule in ``meta["cloudg/render"]`` +(see :func:`context_message` and :func:`embedded_message`), so the layer +can run its key-aware privacy transforms on the data and re-render the +text from the transformed data. +""" + +from __future__ import annotations + +import json +from typing import Any, Callable +from urllib.parse import quote + +from cloudg.mcp.catalog._common import ( + SUMMARY_CORE_KEYS, + asset_brief, + asset_uri, + finding_uri, + parse_severity, +) +from cloudg.mcp.catalog.resources import ( + complete_asset, + complete_dataset, + complete_finding, + complete_framework, + complete_severity, + finding_detail, + framework_detail, +) +from cloudg.mcp.core import ( + EmbeddedResource, + PromptArgument, + PromptMessage, + PromptResult, + PromptSpec, + Registry, + Sensitivity, + TextContent, + TextResourceContents, +) +from cloudg.mcp.state import Dataset, NoDatasetError, ReferenceNotFoundError + +CATEGORY = "prompts" +DATA_META = "cloudg/data" +RENDER_META = "cloudg/render" + +NO_DATA = ( + "No cloudg dataset is loaded yet. First call `workspace_status`, then load data with " + "`load_dataset(path=...)` (an inventory-map.json, findings.json or scanner report inside " + "an allowed directory) or collect it with `map_inventory`. Then run this prompt again." +) + + +# --------------------------------------------------------------------------- +# Message builders +# --------------------------------------------------------------------------- + + +def _plain(data: Any) -> Any: + """``data`` as plain JSON types (what the rendering serialises).""" + return json.loads(json.dumps(data, default=str)) + + +def render_json_block(data: Any) -> str: + """The ``json-block`` rendering of a context message.""" + return "Context for this task (JSON):\n```json\n" + json.dumps(data, indent=1) + "\n```" + + +def render_json(data: Any) -> str: + """The ``json`` rendering of an embedded resource.""" + return json.dumps(data, indent=2) + + +def context_message(data: Any) -> PromptMessage: + """A text message holding ``data`` as a JSON block, with the data and + its rendering rule in ``meta`` for the layer's transforms.""" + plain = _plain(data) + meta = {DATA_META: plain, RENDER_META: "json-block"} + return PromptMessage("user", TextContent(render_json_block(plain), meta=meta)) + + +def embedded_message(uri: str, data: Any) -> PromptMessage: + """An embedded resource holding ``data`` as JSON (same marker).""" + plain = _plain(data) + meta = {DATA_META: plain, RENDER_META: "json"} + return PromptMessage( + "user", + EmbeddedResource(TextResourceContents(uri, render_json(plain), "application/json", meta)), + ) + + +def _result( + description: str, + task: str, + context: dict[str, Any] | None = None, + resources: list[tuple[str, Any]] | None = None, +) -> PromptResult: + msgs = [PromptMessage("user", TextContent(task))] + if context is not None: + msgs.append(context_message(context)) + for uri, data in resources or []: + msgs.append(embedded_message(uri, data)) + return PromptResult(msgs, description) + + +def _dataset(ctx: Any, name: str | None) -> Dataset | None: + try: + return ctx.workspace.get(name or None) + except NoDatasetError: + if name: + raise + return None + + +def _summary_res(ds: Dataset) -> tuple[str, Any]: + """``cloudg://datasets/{name}/summary`` and its exact content.""" + return f"cloudg://datasets/{ds.name}/summary", ds.summary() + + +def _is_active(ctx: Any, ds: Dataset) -> bool: + """Asset / finding / compliance resources read the active dataset, so + they are only embedded when the prompt targets it.""" + return ctx.workspace.active_name == ds.name + + +def _overview(ds: Dataset) -> dict[str, Any]: + s = ds.summary() + keep = ( + "dataset", + "providers", + "total_assets", + "total_edges", + *SUMMARY_CORE_KEYS, + "compliance_frameworks", + "assets_by_type", + ) + return {k: s[k] for k in keep if k in s} + + +def _top_risks(ds: Dataset, n: int = 5, min_severity: str = "") -> list[dict[str, Any]]: + from cloudg.mcp.catalog.findings import risk_ranking + + return [ + { + "asset": r["asset"]["name"], + "id": r["asset"]["id"], + "type": r["asset"]["type"], + "score": r["score"], + "internet_exposed": r["asset"]["internet_exposed"], + "top_finding": r["top_findings"][0]["title"] if r["top_findings"] else None, + } + for r in risk_ranking(ds, min_severity)[:n] + ] + + +def _asset_ctx(ds: Dataset, ref: str) -> dict[str, Any]: + from cloudg.mcp.catalog.resources import asset_detail + + return asset_detail(ds, ref) + + +# --------------------------------------------------------------------------- +# Posture +# --------------------------------------------------------------------------- + + +def security_posture_review(ctx: Any, dataset: str = "") -> PromptResult: + ds = _dataset(ctx, dataset) + if ds is None: + return _result("Security posture review", NO_DATA) + from cloudg.mcp.catalog.compliance import framework_rollup + + ctx_data = { + "dataset": ds.name, + "top_risks": _top_risks(ds), + "compliance": framework_rollup(ds)[:5], + } + task = """Review the security posture of the cloud estate in the dataset named in the +context. The embedded JSON has the overview, the five riskiest assets and compliance posture. + +Work through, calling cloudg tools as needed: +1. `top_risks(top=10)`: the assets to fix first and why (score components). +2. `internet_exposure()` and `attack_paths()` for what an outside attacker can reach. +3. `lateral_movement_paths()` and `cross_account_edges(external_only=true)` (identity pivots). +4. `security_coverage()`: missing GuardDuty / Inspector / WAF etc. +5. `compliance_summary()` to find the weakest frameworks. + +Deliver: an overall rating (critical / poor / fair / good) with justification, the top 5 +risks with evidence (asset, finding, exposure, blast radius), quick wins, and structural +improvements. Cite asset names / ids exactly as tools return them.""" + return _result("Security posture review", task, ctx_data, [_summary_res(ds)]) + + +def executive_summary(ctx: Any, dataset: str = "") -> PromptResult: + ds = _dataset(ctx, dataset) + if ds is None: + return _result("Executive summary", NO_DATA) + task = """Write a one-page executive summary of the cloud security state of the dataset +in the context for non-technical leadership, from the embedded figures (call +`findings_summary`, `compliance_summary` and `top_risks(top=3)` for anything missing). + +Structure: headline (one sentence), scale of the estate, risk level with the 3 issues that +matter most in business terms, compliance standing, and 3 recommended decisions with rough +effort. No identifiers or jargon; numbers rounded.""" + return _result("Executive summary", task, _overview(ds), [_summary_res(ds)]) + + +def attack_surface_report(ctx: Any, dataset: str = "") -> PromptResult: + ds = _dataset(ctx, dataset) + if ds is None: + return _result("Attack surface report", NO_DATA) + exposed = [a for a in ds.assets if a.is_internet_exposed] + by_type: dict[str, int] = {} + for a in exposed: + by_type[a.asset_type.value] = by_type.get(a.asset_type.value, 0) + 1 + ctx_data = { + "dataset": ds.name, + "internet_exposed_total": len(exposed), + "exposed_by_type": by_type, + "exposed_sample": [asset_brief(ds, a) for a in exposed[:15]], + } + task = """Map the external attack surface of the dataset named in the context. + +1. `internet_exposure(limit=100)`: every exposed asset with the rules exposing it, sensitive + ports and WAF protection. +2. `get_edges(internet_only=true, port=22)` and `port=3389` for admin ports open to the world. +3. `attack_paths(max_paths=20)`: routes from the internet to data stores, secrets, keys, roles. +4. For the worst 3 paths, `find_paths(source='internet', target=)` and + `findings_for_asset` on each hop. + +Report: entry points ranked by risk, exposed sensitive services, the critical paths with +every hop explained, and concrete fixes (close rule X, add WAF to Y, move Z private).""" + return _result("Attack surface report", task, ctx_data, [_summary_res(ds)]) + + +# --------------------------------------------------------------------------- +# One asset +# --------------------------------------------------------------------------- + + +def investigate_asset(ctx: Any, ref: str, dataset: str = "") -> PromptResult: + ds = _dataset(ctx, dataset) + if ds is None: + return _result("Investigate asset", NO_DATA) + data = _asset_ctx(ds, ref) + task = """Investigate the asset described in the context (use its `id` as `ref` in the +calls below). Its current detail is embedded. + +1. `get_asset(ref=)` and `get_asset_metadata` (specific keys) for configuration. +2. `neighbors(ref, depth=2)` for what it connects to, then `depends_on` / `dependents`. +3. `findings_for_asset(ref)` then `get_finding` on the serious ones. +4. `blast_radius(ref)` and, if it is exposed, `find_paths(source='internet', target=ref)`. +5. `ontology_neighbourhood(ref)` for semantic relations (encryption, protection, ownership). + +Report what it is, who owns it (tags), how it is exposed, what it can reach, what depends on +it, its findings, and prioritised recommendations.""" + if _is_active(ctx, ds): + return _result( + "Investigate asset", task, {"dataset": ds.name}, [(asset_uri(data["id"]), data)] + ) + return _result("Investigate asset", task, {"dataset": ds.name, **data}) + + +def blast_radius_assessment(ctx: Any, ref: str, dataset: str = "") -> PromptResult: + ds = _dataset(ctx, dataset) + if ds is None: + return _result("Blast radius assessment", NO_DATA) + a = ds.resolve_asset(ref) + deps = ds.dependency_graph().dependents(a.id, 10) + data = { + "dataset": ds.name, + "asset": asset_brief(ds, a), + "transitive_dependents": len(deps), + "accounts_affected": sorted({ds.by_id[d.asset_id].account_id or "" for d in deps} - {""}), + "direct_dependents": [asset_brief(ds, ds.by_id[d.asset_id]) for d in deps if d.depth == 1][ + :20 + ], + } + task = """Assess the blast radius of the asset in the context (`asset.id`) if it is +(a) compromised and (b) deleted or unavailable. + +1. `blast_radius(ref=)`: dependency and network reach, sensitive stores reachable. +2. `dependency_tree(ref, direction='down')` shows what breaks, layer by layer. +3. `lateral_movement_paths(start=)` for identity pivots from it. +4. `findings_for_asset` on the most critical dependents. + +Report both scenarios with affected services / accounts / data, severity, and containment +and resilience recommendations (least privilege, redundancy, segmentation).""" + res = [(asset_uri(a.id), _asset_ctx(ds, a.id))] if _is_active(ctx, ds) else [] + return _result("Blast radius assessment", task, data, res) + + +def change_impact_analysis(ctx: Any, ref: str, change: str = "", dataset: str = "") -> PromptResult: + ds = _dataset(ctx, dataset) + if ds is None: + return _result("Change impact analysis", NO_DATA) + a = ds.resolve_asset(ref) + dg = ds.dependency_graph() + data = { + "dataset": ds.name, + "planned_change": change or "a configuration change", + "asset": asset_brief(ds, a), + "depends_on": [asset_brief(ds, ds.by_id[d.asset_id]) for d in dg.depends_on(a.id, 1)][:20], + "dependents": [asset_brief(ds, ds.by_id[d.asset_id]) for d in dg.dependents(a.id, 1)][:20], + } + task = """Plan the change described in the context (`planned_change`) to the asset in the +context (`asset.id`). Direct upstream / downstream dependencies are embedded. + +1. `dependents(ref=, max_depth=5)` and `dependency_tree(ref, direction='down')`. +2. `shared_dependencies()`: is it a single point of failure? +3. `cross_account_edges(account_id=)` for other accounts affected. +4. `get_asset_metadata(ref)` for the settings being changed. + +Deliver: affected components ranked by impact, risks, pre-checks, a rollout and rollback +plan, and how to verify (e.g. snapshot_dataset now, re-collect after, diff_datasets).""" + res = [(asset_uri(a.id), _asset_ctx(ds, a.id))] if _is_active(ctx, ds) else [] + return _result("Change impact analysis", task, data, res) + + +# --------------------------------------------------------------------------- +# Identity / compliance / remediation +# --------------------------------------------------------------------------- + + +def cross_account_trust_review(ctx: Any, dataset: str = "") -> PromptResult: + ds = _dataset(ctx, dataset) + if ds is None: + return _result("Cross-account trust review", NO_DATA) + from cloudg.inventory.dependencies import cross_account_edges + + rows = ds.cached("cross_account", lambda: cross_account_edges(ds.assets, ds.edges)) + pairs: dict[str, int] = {} + for r in rows: + k = f"{r['source_account']} -> {r['target_account']} ({r['edge_type']})" + pairs[k] = pairs.get(k, 0) + 1 + data = { + "dataset": ds.name, + "cross_account_edges": len(rows), + "external": sum(1 for r in rows if r["external"]), + "by_account_pair": pairs, + } + task = """Review trust relationships across account boundaries in the dataset named in the +context. + +1. `cross_account_edges(external_only=true)` then all of them by account pair. +2. `get_edges(edge_types=['IAM_TRUST'])` and `get_edges(edge_types=['GRANTS_ACCESS'], + cross_account_only=true)`. +3. `lateral_movement_paths()`: chains a compromised external principal could follow. +4. `organization_topology(include_policies=true)` for SCPs that limit the blast radius. + +Flag unknown / external accounts, overly broad trust (account root, wildcards), missing +external-id conditions, and recommend tightening with evidence for each.""" + return _result("Cross-account trust review", task, data, [_summary_res(ds)]) + + +def compliance_gap_analysis(ctx: Any, framework: str, dataset: str = "") -> PromptResult: + ds = _dataset(ctx, dataset) + if ds is None: + return _result("Compliance gap analysis", NO_DATA) + from cloudg.mcp.catalog.compliance import framework_rollup + + data = {"dataset": ds.name, "framework": framework, "posture": framework_rollup(ds, framework)} + task = """Run a compliance gap analysis for the framework named in the context +(`framework`, called below). + +1. `compliance_summary(framework=)`. +2. `list_controls(framework=, status='FAIL')` and + `list_controls(framework=, source='ruleset')` for the full control set. +3. `control_status` on the failing controls with the most / worst findings. +4. `compliance_gaps(framework=)` to list the assets behind the gaps. + +Deliver: pass rate and failing controls grouped by domain, the assets responsible, a +remediation backlog ordered by severity x effort, and evidence references (finding ids).""" + try: + detail = framework_detail(ds, framework) + except ReferenceNotFoundError: + return _result("Compliance gap analysis", task, data) + if _is_active(ctx, ds): + uri = "cloudg://compliance/" + quote(framework, safe="") + return _result( + "Compliance gap analysis", + task, + {"dataset": ds.name, "framework": framework}, + [(uri, detail)], + ) + return _result("Compliance gap analysis", task, {**detail, "framework": framework}) + + +def remediation_plan(ctx: Any, severity: str = "HIGH", dataset: str = "") -> PromptResult: + ds = _dataset(ctx, dataset) + if ds is None: + return _result("Remediation plan", NO_DATA) + sev = parse_severity(severity or "HIGH").value + data = {"dataset": ds.name, "min_severity": sev, "top_risks": _top_risks(ds, 10, sev)} + task = f"""Build a remediation plan for open findings at {sev} or above in the dataset named +in the context. + +1. `list_findings(min_severity='{sev}', limit=100)` (page with next_cursor if needed). +2. `top_risks(top=20, min_severity='{sev}')` to order work by real risk. +3. `get_finding` for remediation text; `blast_radius` on assets where a fix may disrupt. +4. Group fixes that share a root cause (same security group, role, policy, account setting). + +Deliver: phased plan (now / this sprint / next quarter) with owner hints from tags, each +item with finding ids, affected assets, the fix, its risk reduction and change risk. Note +accepted risks that could be suppressed (suppress_findings) with the reason.""" + return _result("Remediation plan", task, data, [_summary_res(ds)]) + + +def incident_triage(ctx: Any, finding_id: str, dataset: str = "") -> PromptResult: + ds = _dataset(ctx, dataset) + if ds is None: + return _result("Incident triage", NO_DATA) + f = ds.get_finding(finding_id) + data = finding_detail(ds, f) + aid = ds.finding_asset_id(f) + target = {"dataset": ds.name, "target_ref": aid or f.resource_arn or f.resource_id} + task = f"""Triage the finding in the context ({f.severity.value}) as a potential +incident. `id` is the finding id and `target_ref` the affected asset ( below). + +1. `get_finding(finding_id=)` for evidence and mapped controls. +2. `get_asset(ref=)` and `internet_exposure()`: is the affected asset reachable? +3. `find_paths(source='internet', target=)` and `blast_radius(ref=)`. +4. `lateral_movement_paths(start=)` to see where an attacker goes next. + +Deliver: is it exploitable (likely / possible / unlikely) and why, impact if exploited, +immediate containment steps, evidence to collect, the permanent fix, and severity +re-rating if justified.""" + if _is_active(ctx, ds): + return _result("Incident triage", task, target, [(finding_uri(f.id), data)]) + return _result("Incident triage", task, {**data, **target}) + + +def drift_review(ctx: Any, base: str, target: str = "") -> PromptResult: + diff = ctx.workspace.diff(base, target or None, limit=25) + task = """Review what changed between the datasets `base` (before) and `target` (after) +named in the context. The diff is embedded (counts plus the first items). + +Call `diff_datasets(base=, target=, limit=200)` for more, and `get_asset` / +`findings_for_asset` on anything suspicious. Report: new or newly exposed assets, removed +controls (security groups, WAF, logging), new trust edges, findings introduced vs resolved, +and whether each change looks intended or like drift.""" + return _result("Drift review", task, diff) + + +# --------------------------------------------------------------------------- +# Registration +# --------------------------------------------------------------------------- + + +def _ds_arg() -> PromptArgument: + return PromptArgument( + "dataset", "Dataset name; empty = active dataset.", completion=complete_dataset + ) + + +def _ref_arg() -> PromptArgument: + return PromptArgument("ref", "Asset id, ARN or unique name.", True, complete_asset) + + +_C, _I = Sensitivity.CONFIDENTIAL, Sensitivity.INTERNAL + +# (handler, title, description, arguments factory, sensitivity) +_PROMPTS: list[tuple[Callable[..., Any], str, str, Callable[[], list], Sensitivity]] = [ + ( + security_posture_review, + "Security posture review", + "Whole-estate security review: risks, exposure, identity pivots, coverage gaps, " + "compliance.", + lambda: [_ds_arg()], + _C, + ), + ( + executive_summary, + "Executive summary", + "Non-technical one-page summary of the estate's security state.", + lambda: [_ds_arg()], + _I, + ), + ( + attack_surface_report, + "Attack surface report", + "External attack surface: exposed assets, open admin ports, internet-to-crown-jewel paths.", + lambda: [_ds_arg()], + _C, + ), + ( + investigate_asset, + "Investigate asset", + "Deep dive on one asset: config, relations, exposure, findings, blast radius.", + lambda: [_ref_arg(), _ds_arg()], + _C, + ), + ( + blast_radius_assessment, + "Blast radius assessment", + "Impact of compromise or loss of one asset.", + lambda: [_ref_arg(), _ds_arg()], + _C, + ), + ( + change_impact_analysis, + "Change impact analysis", + "What a planned change to one asset affects, with rollout and rollback plan.", + lambda: [_ref_arg(), PromptArgument("change", "The planned change, in words."), _ds_arg()], + _C, + ), + ( + cross_account_trust_review, + "Cross-account trust review", + "IAM trust and grants across account boundaries, external accounts, pivot chains.", + lambda: [_ds_arg()], + _C, + ), + ( + compliance_gap_analysis, + "Compliance gap analysis", + "Failing controls, responsible assets and a remediation backlog for one framework.", + lambda: [ + PromptArgument( + "framework", + "Framework, e.g. CIS-AWS, PCI-DSS, NIST-800-53.", + True, + complete_framework, + ), + _ds_arg(), + ], + _C, + ), + ( + remediation_plan, + "Remediation plan", + "Phased fix plan for findings at or above a severity.", + lambda: [ + PromptArgument( + "severity", "Minimum severity (default HIGH).", False, complete_severity + ), + _ds_arg(), + ], + _C, + ), + ( + incident_triage, + "Incident triage", + "Exploitability, impact and containment for one finding.", + lambda: [PromptArgument("finding_id", "Finding id.", True, complete_finding), _ds_arg()], + _C, + ), + ( + drift_review, + "Drift review", + "Explain the differences between two datasets (before / after).", + lambda: [ + PromptArgument("base", "Baseline dataset.", True, complete_dataset), + PromptArgument("target", "Later dataset; empty = active.", False, complete_dataset), + ], + _C, + ), +] + + +def register(reg: Registry) -> None: + """Add the packaged prompts to ``reg``.""" + for fn, title, description, arguments, sensitivity in _PROMPTS: + reg.add( + PromptSpec( + name=fn.__name__, + handler=fn, + title=title, + description=description, + arguments=arguments(), + category=CATEGORY, + sensitivity=sensitivity, + tags={"workflow"}, + ) + ) diff --git a/cloudg/mcp/catalog/resources.py b/cloudg/mcp/catalog/resources.py new file mode 100644 index 0000000..464d017 --- /dev/null +++ b/cloudg/mcp/catalog/resources.py @@ -0,0 +1,602 @@ +"""Resources and resource templates. + +Resources give clients addressable, cacheable views of the workspace. +The tool results link to them (``resource_link``) instead of inlining big +payloads: + +===================================== ======================================= +``cloudg://workspace`` loaded datasets, active one, roots +``cloudg://datasets`` dataset listing +``cloudg://datasets/{dataset}/summary`` headline numbers of one dataset +``cloudg://assets/{+ref}`` one asset (detail) +``cloudg://assets/{+ref}/neighbors`` its edges and neighbours +``cloudg://assets/{+ref}/findings`` its findings +``cloudg://findings/summary`` finding counts by severity / tool +``cloudg://findings/{finding_id}`` one finding in full +``cloudg://findings/severity/{severity}`` findings of one severity (max 200) +``cloudg://compliance`` posture of every framework +``cloudg://compliance/{framework}`` one framework: controls and gaps +``cloudg://graph/d3``, ``/{format}`` whole graph as D3 / Cytoscape / GraphML +``cloudg://ontology/turtle``, ``/{format}`` ontology (turtle, json-ld, xml, nt) +``cloudg://schema/asset-types`` ... vocabulary (public) +``cloudg://docs``, ``/{topic}`` cloudg documentation sections +===================================== ======================================= + +Asset / dataset / framework / severity / topic variables have completions. +""" + +from __future__ import annotations + +import os +import re +from pathlib import Path +from typing import Any + +from cloudg.mcp.catalog._common import ( + Catalog, + asset_brief, + edge_brief, + finding_brief, + ws_dataset, +) +from cloudg.mcp.catalog._resources import ( + GRAPH_FORMATS, + ONTOLOGY_FORMATS, + asset_detail, + complete_asset, + complete_asset_type, + complete_dataset, + complete_finding, + complete_framework, + complete_graph_format, + complete_ontology_format, + complete_severity, + finding_detail, + framework_detail, +) +from cloudg.mcp.core import Registry, Sensitivity, TextResourceContents +from cloudg.mcp.state import SEVERITY_RANK, ReferenceNotFoundError +from cloudg.schema.models import AssetType, EdgeType + +DOC_TOPICS: dict[str, tuple[str, str, str]] = { + # topic: (file, heading prefix, one-line description) + "pipeline": ("DOCUMENTATION.md", "## The pipeline", "Phases of a cloudg run"), + "cli": ("DOCUMENTATION.md", "## CLI reference", "cloudg command-line reference"), + "inputs": ("DOCUMENTATION.md", "## Input requirements", "Scanner report formats"), + "outputs": ("DOCUMENTATION.md", "## Output contract", "Files cloudg writes"), + "data-models": ("DOCUMENTATION.md", "## Data models", "Finding / CloudAsset / NetworkEdge"), + "configuration": ("DOCUMENTATION.md", "## Configuration", "config.yaml reference"), + "python-api": ("DOCUMENTATION.md", "## Python API", "CloudGEngine and inventory API"), + "recipes": ("DOCUMENTATION.md", "## Integration recipes", "Integration recipes"), + "authentication": ("DOCUMENTATION.md", "## Authentication", "Cloud credentials"), + "identifiers": ( + "INVENTORY_REFERENCE.md", + "## 4. Identifier formats", + "ARN / Azure / GCP / placeholder identifier formats", + ), + "edges": ("INVENTORY_REFERENCE.md", "## 5. NetworkEdge", "NetworkEdge fields"), + "edge-types": ( + "INVENTORY_REFERENCE.md", + "## 6. Edge types and direction", + "Edge types, direction and dependency semantics", + ), + "relations": ( + "INVENTORY_REFERENCE.md", + "## 7. The relation object", + "Declared relations on assets", + ), + "inventory-summary": ("INVENTORY_REFERENCE.md", "## 9. summary", "InventoryResult.summary"), + "dependencies": ( + "INVENTORY_REFERENCE.md", + "## 12. DependencyGraph results", + "depends_on / dependents / tree / blast radius", + ), + "inventory-map": ( + "INVENTORY_REFERENCE.md", + "## 14. inventory-map.json", + "inventory-map.json format", + ), + "organization": ( + "INVENTORY_REFERENCE.md", + "## 17. OrganizationTopology", + "Organization / Control Tower topology", + ), +} + +MCP_GUIDE = """\ +# Using cloudg through MCP + +1. `workspace_status` shows what is loaded. Load data with `load_dataset(path=...)` + (inventory-map.json, findings.json, scanner output) or collect live with + `map_inventory` (needs cloud credentials on the server). +2. Orient: `dataset_summary`, `count_assets(group_by=...)`, `findings_summary`. +3. Prioritise: `top_risks`, `internet_exposure`, `attack_paths`. +4. Drill down: `get_asset(ref)`, `findings_for_asset`, `neighbors`, `blast_radius`, + `dependency_tree`, `find_paths`. +5. Compliance: `compliance_summary`, `list_controls`, `control_status`, `compliance_gaps`. +6. Semantics: `relation_groups`, `ontology_neighbourhood`, `sparql_query` (read-only). +7. Change tracking: `snapshot_dataset` before re-collecting, then `diff_datasets`. + +Asset references accept ids, ARNs, unique names or unique ARN tails. List tools +are paginated: pass `next_cursor` back as `cursor`. Identifiers may be +pseudonymised by the server's privacy policy; pass them back unchanged. +Prompts (`security_posture_review`, `investigate_asset`, ...) package these +steps for common jobs. +""" + + +def docs_dir() -> Path | None: + env = os.environ.get("CLOUDG_DOCS_DIR") + candidates = [Path(env)] if env else [] + import cloudg + + pkg = Path(cloudg.__file__).resolve().parent + # A source checkout keeps them in docs/; wheels ship them in cloudg/_docs + candidates += [pkg.parent / "docs", pkg / "_docs"] + for c in candidates: + if c.is_dir(): + return c + return None + + +def doc_section(topic: str) -> str: + if topic == "mcp": + return MCP_GUIDE + if topic not in DOC_TOPICS: + raise ReferenceNotFoundError( + f"Unknown docs topic. Topics: mcp, {', '.join(DOC_TOPICS)}", data={"value": topic} + ) + fname, heading, _ = DOC_TOPICS[topic] + d = docs_dir() + path = d / fname if d else None + if path is None or not path.exists(): + return ( + f"# {topic}\n\nThe cloudg documentation files are not installed with this " + "package. See https://morpheuslord.github.io/cloudg/ or read the 'mcp' topic." + ) + lines = path.read_text().splitlines() + start = next((i for i, ln in enumerate(lines) if ln.startswith(heading)), None) + if start is None: + return f"# {topic}\n\nSection '{heading}' not found in {fname}." + level = len(heading) - len(heading.lstrip("#")) + end = len(lines) + for i in range(start + 1, len(lines)): + m = re.match(r"^(#+)\s", lines[i]) + if m and len(m.group(1)) <= level and not _in_code(lines, i): + end = i + break + text = "\n".join(lines[start:end]).strip() + return text[:40_000] + + +def _in_code(lines: list[str], idx: int) -> bool: + return sum(1 for ln in lines[:idx] if ln.startswith("```")) % 2 == 1 + + +def complete_topic(ctx: Any, partial: str = "") -> list[str]: + return [t for t in ["mcp", *DOC_TOPICS] if t.startswith(partial.lower())] + + +__all__ = [ + "CATALOG", + "DOC_TOPICS", + "GRAPH_FORMATS", + "ONTOLOGY_FORMATS", + "asset_detail", + "complete_asset", + "complete_asset_type", + "complete_dataset", + "complete_finding", + "complete_framework", + "complete_graph_format", + "complete_ontology_format", + "complete_severity", + "complete_topic", + "doc_section", + "docs_dir", + "finding_detail", + "framework_detail", + "register", +] + +CATALOG = Catalog() + +# -- workspace ------------------------------------------------------ + + +@CATALOG.resource( + "cloudg://workspace", + title="Workspace", + category="workspace", + sensitivity=Sensitivity.INTERNAL, + description="Loaded datasets, the active dataset and allowed directories.", +) +def workspace_resource(ctx: Any) -> dict: + return ctx.workspace.status() + + +@CATALOG.resource( + "cloudg://datasets", + title="Datasets", + category="workspace", + sensitivity=Sensitivity.INTERNAL, + description="Loaded datasets.", +) +def datasets_resource(ctx: Any) -> dict: + ws = ctx.workspace + return {"active_dataset": ws.active_name, "datasets": [d.describe() for d in ws.datasets()]} + + +@CATALOG.resource_template( + "cloudg://datasets/{dataset}/summary", + title="Dataset summary", + category="workspace", + sensitivity=Sensitivity.INTERNAL, + description="Headline numbers of one dataset.", + completions={"dataset": complete_dataset}, +) +def dataset_summary_resource(ctx: Any, dataset: str) -> dict: + return ctx.workspace.get(dataset).summary() + + +# -- assets --------------------------------------------------------- + + +@CATALOG.resource_template( + "cloudg://assets/{+ref}", + title="Asset", + category="inventory", + sensitivity=Sensitivity.CONFIDENTIAL, + description="One asset of the active dataset (ref = id, ARN or unique name, URL-encoded).", + completions={"ref": complete_asset}, +) +def asset_resource(ctx: Any, ref: str) -> dict: + return asset_detail(ws_dataset(ctx), ref) + + +@CATALOG.resource_template( + "cloudg://assets/{+ref}/neighbors", + title="Asset neighbours", + category="graph", + sensitivity=Sensitivity.CONFIDENTIAL, + description="Edges and one-hop neighbours of an asset.", + completions={"ref": complete_asset}, +) +def asset_neighbors_resource(ctx: Any, ref: str) -> dict: + ds = ws_dataset(ctx) + node = ref if (ref in ds.graph and ref not in ds.by_id) else ds.resolve_asset(ref).id + adj = ds.adjacency.get(node, [])[:500] + others = list(dict.fromkeys(o for _, o, _ in adj)) + return { + "asset": node, + "edges": [edge_brief(ds, e) for e, _, _ in adj], + "neighbors": [ + asset_brief(ds, ds.by_id[o]) if o in ds.by_id else {"id": o, "external": True} + for o in others + ], + } + + +@CATALOG.resource_template( + "cloudg://assets/{+ref}/findings", + title="Asset findings", + category="findings", + sensitivity=Sensitivity.CONFIDENTIAL, + description="Open findings of an asset.", + completions={"ref": complete_asset}, +) +def asset_findings_resource(ctx: Any, ref: str) -> dict: + ds = ws_dataset(ctx) + a = ds.resolve_asset(ref) + fs = sorted(ds.open_findings(a.id), key=lambda f: -f.risk_score) + return {"asset": asset_brief(ds, a), "findings": [finding_brief(ds, f) for f in fs]} + + +# -- findings ------------------------------------------------------- + + +@CATALOG.resource( + "cloudg://findings/summary", + title="Findings summary", + category="findings", + sensitivity=Sensitivity.INTERNAL, + description="Open finding counts by severity and tool (active dataset).", +) +def findings_summary_resource(ctx: Any) -> dict: + ds = ws_dataset(ctx) + by_tool: dict[str, int] = {} + for f in ds.open_findings(): + by_tool[f.source_tool] = by_tool.get(f.source_tool, 0) + 1 + s = ds.summary() + return { + "dataset": ds.name, + "open_findings": s["open_findings"], + "suppressed": s["suppressed_findings"], + "severity_breakdown": s["severity_breakdown"], + "by_tool": by_tool, + } + + +@CATALOG.resource_template( + "cloudg://findings/{finding_id}", + title="Finding", + category="findings", + sensitivity=Sensitivity.CONFIDENTIAL, + description="One finding with evidence and remediation.", + completions={"finding_id": complete_finding}, +) +def finding_resource(ctx: Any, finding_id: str) -> dict: + ds = ws_dataset(ctx) + return finding_detail(ds, ds.get_finding(finding_id)) + + +@CATALOG.resource_template( + "cloudg://findings/severity/{severity}", + title="Findings by severity", + category="findings", + sensitivity=Sensitivity.CONFIDENTIAL, + description="Open findings of one severity (first 200, by risk).", + completions={"severity": complete_severity}, +) +def findings_by_severity_resource(ctx: Any, severity: str) -> dict: + ds = ws_dataset(ctx) + sev = severity.upper() + if sev not in SEVERITY_RANK: + raise ReferenceNotFoundError( + f"Unknown severity; use one of {', '.join(SEVERITY_RANK)}", data={"value": severity} + ) + fs = sorted( + (f for f in ds.open_findings() if f.severity.value == sev), key=lambda f: -f.risk_score + ) + return { + "dataset": ds.name, + "severity": sev, + "total": len(fs), + "truncated": len(fs) > 200, + "findings": [finding_brief(ds, f) for f in fs[:200]], + } + + +# -- compliance ----------------------------------------------------- + + +@CATALOG.resource( + "cloudg://compliance", + title="Compliance posture", + category="compliance", + sensitivity=Sensitivity.INTERNAL, + description="Posture of every compliance framework in the active dataset.", +) +def compliance_resource(ctx: Any) -> dict: + from cloudg.mcp.catalog.compliance import framework_rollup + + ds = ws_dataset(ctx) + return {"dataset": ds.name, "frameworks": framework_rollup(ds)} + + +@CATALOG.resource_template( + "cloudg://compliance/{framework}", + title="Framework posture", + category="compliance", + sensitivity=Sensitivity.INTERNAL, + description="One framework: posture and failing controls.", + completions={"framework": complete_framework}, +) +def framework_resource(ctx: Any, framework: str) -> dict: + return framework_detail(ws_dataset(ctx), framework) + + +# -- graph / ontology ----------------------------------------------- + + +def _graph(ctx: Any, fmt: str) -> Any: + """d3 / cytoscape as JSON objects (the layer transforms them key by key); + graphml as text.""" + import networkx as nx + + ds = ws_dataset(ctx) + if fmt == "d3": + return ds.builder.to_d3_json() + if fmt == "cytoscape": + return ds.builder.to_cytoscape_json() + if fmt == "graphml": + return [ + TextResourceContents( + f"cloudg://graph/{fmt}", + "\n".join(nx.generate_graphml(ds.graph)), + "application/graphml+xml", + ) + ] + raise ReferenceNotFoundError( + f"Unknown graph format; use {', '.join(GRAPH_FORMATS)}", data={"value": fmt} + ) + + +@CATALOG.resource( + "cloudg://graph/d3", + title="Graph (D3)", + category="graph", + sensitivity=Sensitivity.CONFIDENTIAL, + description="The active dataset's whole relationship graph as D3 JSON " + "({nodes, links}). Large; prefer subgraph_export for parts.", +) +def graph_d3_resource(ctx: Any) -> Any: + return _graph(ctx, "d3") + + +# RESTRICTED: serves GraphML, an opaque text export the privacy layer can +# only scan as free text, so strict ceilings hide it. +@CATALOG.resource_template( + "cloudg://graph/{format}", + title="Graph export", + category="graph", + sensitivity=Sensitivity.RESTRICTED, + description="Whole graph as d3 or cytoscape (application/json) " + "or graphml (application/graphml+xml); each read returns the " + "format's own mimeType.", + completions={"format": complete_graph_format}, +) +def graph_resource( + ctx: Any, + format: str, # pylint: disable=redefined-builtin # URI variable +) -> Any: + return _graph(ctx, format) + + +def _ontology(ctx: Any, fmt: str) -> list[TextResourceContents]: + if fmt not in ONTOLOGY_FORMATS: + raise ReferenceNotFoundError( + f"Unknown ontology format; use {', '.join(ONTOLOGY_FORMATS)}", data={"value": fmt} + ) + g = ws_dataset(ctx).ontology().graph + return [ + TextResourceContents( + f"cloudg://ontology/{fmt}", g.serialize(format=fmt), ONTOLOGY_FORMATS[fmt] + ) + ] + + +# The RDF serialisations are opaque text: RESTRICTED (see graph/{format}). +@CATALOG.resource( + "cloudg://ontology/turtle", + title="Ontology (Turtle)", + category="ontology", + sensitivity=Sensitivity.RESTRICTED, + mime_type="text/turtle", + description="The active dataset's RDF ontology in Turtle.", +) +def ontology_turtle_resource(ctx: Any) -> list[TextResourceContents]: + return _ontology(ctx, "turtle") + + +@CATALOG.resource_template( + "cloudg://ontology/{format}", + title="Ontology export", + category="ontology", + sensitivity=Sensitivity.RESTRICTED, + mime_type="text/turtle", + description="Ontology as turtle (text/turtle), json-ld " + "(application/ld+json), xml (application/rdf+xml) or nt " + "(application/n-triples); each read returns the format's own " + "mimeType.", + completions={"format": complete_ontology_format}, +) +def ontology_resource( + ctx: Any, + format: str, # pylint: disable=redefined-builtin # URI variable +) -> list[TextResourceContents]: + return _ontology(ctx, format) + + +# -- schema / docs (public) ----------------------------------------- + + +@CATALOG.resource( + "cloudg://schema/asset-types", + title="Asset types", + category="meta", + sensitivity=Sensitivity.PUBLIC, + description="Asset types by family.", +) +def asset_types_resource(ctx: Any) -> dict: + from cloudg.mcp.catalog.meta import enum_comments + + fam: dict[str, list[str]] = {} + for t in AssetType: + fam.setdefault(enum_comments()["family"].get(t.name, "Other"), []).append(t.value) + return fam + + +@CATALOG.resource_template( + "cloudg://schema/asset-types/{asset_type}", + title="Asset type", + category="meta", + sensitivity=Sensitivity.PUBLIC, + description="How cloudg models one asset type.", + completions={"asset_type": complete_asset_type}, +) +def asset_type_resource(ctx: Any, asset_type: str) -> dict: + from cloudg.graph.ontology import _ASSET_TYPE_CLASSES + from cloudg.mcp.catalog.meta import enum_comments + from cloudg.renderers.terraform_export import _TF_RESOURCE_MAP + + try: + t = AssetType(asset_type.upper()) + except ValueError: + raise ReferenceNotFoundError("Unknown asset type.", data={"value": asset_type}) from None + return { + "asset_type": t.value, + "family": enum_comments()["family"].get(t.name), + "note": enum_comments()["asset_note"].get(t.name), + "ontology_class": f"cm:{_ASSET_TYPE_CLASSES.get(t, 'CloudResource')}", + "terraform_type": _TF_RESOURCE_MAP.get(t), + } + + +@CATALOG.resource( + "cloudg://schema/edge-types", + title="Edge types", + category="meta", + sensitivity=Sensitivity.PUBLIC, + description="Edge types with direction and dependency semantics.", +) +def edge_types_resource(ctx: Any) -> dict: + from cloudg.inventory.dependencies import DEPENDENCY_DIRECTION + from cloudg.mcp.catalog.meta import EDGE_READS + + return { + e.value: { + "reads_as": EDGE_READS.get(e.value, ("", ""))[0], + "typical": EDGE_READS.get(e.value, ("", ""))[1], + "dependency_direction": DEPENDENCY_DIRECTION.get(e.value, "none"), + } + for e in EdgeType + } + + +@CATALOG.resource( + "cloudg://schema/relation-types", + title="Ontology relation types", + category="meta", + sensitivity=Sensitivity.PUBLIC, + description="Semantic relation types by relation group.", +) +def relation_types_resource(ctx: Any) -> dict: + from cloudg.mcp.catalog.meta import _relation_types + + return _relation_types() + + +@CATALOG.resource( + "cloudg://docs", + title="Documentation topics", + category="meta", + sensitivity=Sensitivity.PUBLIC, + description="Index of cloudg documentation topics (cloudg://docs/{topic}).", +) +def docs_index_resource(ctx: Any) -> dict: + return { + "available": docs_dir() is not None, + "topics": { + "mcp": "How to use cloudg through MCP", + **{k: v[2] for k, v in DOC_TOPICS.items()}, + }, + } + + +@CATALOG.resource_template( + "cloudg://docs/{topic}", + title="Documentation", + category="meta", + sensitivity=Sensitivity.PUBLIC, + mime_type="text/markdown", + description="One cloudg documentation section (markdown).", + completions={"topic": complete_topic}, +) +def docs_resource(ctx: Any, topic: str) -> str: + return doc_section(topic) + + +def register(reg: Registry) -> None: + """Add this module's tools to ``reg``.""" + CATALOG.register(reg) diff --git a/cloudg/mcp/catalog/workspace.py b/cloudg/mcp/catalog/workspace.py new file mode 100644 index 0000000..bea2475 --- /dev/null +++ b/cloudg/mcp/catalog/workspace.py @@ -0,0 +1,240 @@ +"""Workspace tools: load, list, select, snapshot, diff and unload datasets. + +A dataset is one inventory map / findings report / scanner output (or a +live collection) held in memory. Every other tool reads the active dataset +unless given ``dataset=``. +""" + +from __future__ import annotations + +from typing import Annotated, Any, Literal + +from pydantic import Field + +from cloudg.mcp.catalog._common import ( + Catalog, + DatasetArg, + DiffOut, + SummaryOut, + WorkspaceOut, + schema, + ws_dataset, +) +from cloudg.mcp.core import Capability, Registry, Sensitivity + +CATEGORY = "workspace" +R, W, FS = Capability.READ_STATE, Capability.WRITE_STATE, Capability.READ_FS + +CATALOG = Catalog() + + +@CATALOG.tool( + title="Workspace status", + category=CATEGORY, + sensitivity=Sensitivity.INTERNAL, + read_only=True, + idempotent=True, + open_world=False, + output_schema=schema(WorkspaceOut), + tags={"start-here"}, +) +def workspace_status(ctx: Any) -> dict: + """Show what is loaded: every dataset (asset / edge / finding counts, + source, which analyses are cached), the active dataset, and the + directories file tools may read and write. Call this first; if no + dataset is loaded, use load_dataset or map_inventory.""" + status = ctx.workspace.status() + if not status["datasets"]: + status["next_steps"] = [ + "load_dataset(path='')", + "map_inventory(providers=['aws']) to collect live data (needs credentials)", + ] + else: + ctx.link("cloudg://workspace", "workspace", title="Workspace status") + return status + + +@CATALOG.tool( + title="List datasets", + category=CATEGORY, + sensitivity=Sensitivity.INTERNAL, + read_only=True, + idempotent=True, + open_world=False, +) +def list_datasets(ctx: Any) -> dict: + """List loaded datasets with their size, kind (inventory, report, + generic, live, snapshot, scanner name) and which one is active.""" + ws = ctx.workspace + items = [{**d.describe(), "active": d.name == ws.active_name} for d in ws.datasets()] + return {"active_dataset": ws.active_name, "datasets": items, "total": len(items)} + + +@CATALOG.tool( + title="Load dataset", + category=CATEGORY, + sensitivity=Sensitivity.INTERNAL, + capabilities={R, W, FS}, + read_only=False, + destructive=False, + idempotent=False, + open_world=False, +) +def load_dataset( + ctx: Any, + path: Annotated[ + str, + Field( + min_length=1, + description="File or directory inside an allowed root: inventory-map.json " + "(or its directory), a cloudg findings.json report, a generic " + "{assets, edges, findings} JSON, or native Prowler / ScoutSuite / Checkov / " + "Trivy output.", + ), + ], + name: Annotated[str, Field(description="Dataset name; default derives from the path.")] = "", + kind: Annotated[ + Literal[ + "auto", + "inventory", + "report", + "generic", + "prowler", + "scoutsuite", + "checkov", + "trivy", + ], + Field(description="Content type; 'auto' detects it."), + ] = "auto", + activate: Annotated[bool, Field(description="Make it the active dataset.")] = True, + replace: Annotated[ + bool, + Field( + description="Overwrite an existing dataset with the " + "same name (otherwise a name clash is an error)." + ), + ] = False, +) -> dict: + """Load a file or directory into the workspace as a named dataset. + An inventory directory that also holds findings.json gets those + findings merged in. Scanner output is normalised (deduplicated and + mapped to compliance frameworks); Prowler must be ASFF (-M json-asff), + OCSF output is refused. Returns the dataset summary.""" + ds = ctx.workspace.load(path, name or None, kind=kind, activate=activate, replace=replace) + ctx.link(f"cloudg://datasets/{ds.name}/summary", f"{ds.name} summary", title="Dataset summary") + return {"loaded": ds.name, "active": ctx.workspace.active_name == ds.name, **ds.summary()} + + +@CATALOG.tool( + title="Select dataset", + category=CATEGORY, + sensitivity=Sensitivity.INTERNAL, + capabilities={R, W}, + read_only=False, + destructive=False, + idempotent=True, + open_world=False, +) +def select_dataset( + ctx: Any, name: Annotated[str, Field(min_length=1, description="Dataset to activate.")] +) -> dict: + """Make a loaded dataset the active one (the default for every tool).""" + ds = ctx.workspace.select(name) + return {"active_dataset": ds.name, **ds.describe()} + + +@CATALOG.tool( + title="Unload dataset", + category=CATEGORY, + sensitivity=Sensitivity.INTERNAL, + capabilities={R, W}, + read_only=False, + destructive=True, + idempotent=True, + open_world=False, +) +def unload_dataset( + ctx: Any, name: Annotated[str, Field(min_length=1, description="Dataset to drop.")] +) -> dict: + """Remove a dataset from memory (files on disk are untouched). Any + suppressions or ingested findings held only in memory are lost.""" + ctx.workspace.remove(name) + return { + "unloaded": name, + "active_dataset": ctx.workspace.active_name, + "remaining": ctx.workspace.names(), + } + + +@CATALOG.tool( + title="Snapshot dataset", + category=CATEGORY, + sensitivity=Sensitivity.INTERNAL, + capabilities={R, W}, + read_only=False, + destructive=False, + idempotent=False, + open_world=False, +) +def snapshot_dataset( + ctx: Any, + new_name: Annotated[str, Field(min_length=1, description="Name for the copy.")], + dataset: DatasetArg = "", + replace: Annotated[ + bool, + Field( + description="Overwrite an existing dataset named " + "new_name (otherwise a name clash is an error)." + ), + ] = False, +) -> dict: + """Freeze a copy of a dataset under a new name (e.g. before ingesting + new findings or re-collecting) so diff_datasets can compare later. + The copy does not become active.""" + ds = ctx.workspace.snapshot(dataset or None, new_name, replace=replace) + return {"snapshot": ds.name, **ds.describe()} + + +@CATALOG.tool( + title="Diff datasets", + category=CATEGORY, + sensitivity=Sensitivity.CONFIDENTIAL, + read_only=True, + idempotent=True, + open_world=False, + output_schema=schema(DiffOut), +) +def diff_datasets( + ctx: Any, + base: Annotated[str, Field(min_length=1, description="Baseline dataset (before).")], + target: Annotated[str, Field(description="Dataset to compare (after); empty = active.")] = "", + limit: Annotated[int, Field(ge=1, le=500, description="Max items per section.")] = 50, +) -> dict: + """Compare two datasets: assets added / removed / changed (matched + on ARN), edges added / removed, findings new / resolved / changed + severity, and assets that became internet-exposed. Use for drift and + change review between two collections or reports.""" + return ctx.workspace.diff(base, target or None, limit=limit) + + +@CATALOG.tool( + title="Dataset summary", + category=CATEGORY, + sensitivity=Sensitivity.INTERNAL, + read_only=True, + idempotent=True, + open_world=False, + output_schema=schema(SummaryOut), +) +def dataset_summary(ctx: Any, dataset: DatasetArg = "") -> dict: + """Headline numbers for a dataset: asset / edge / finding counts, + open findings by severity, accounts, regions, internet-exposed + assets, cross-account edges, top asset types, compliance frameworks.""" + ds = ws_dataset(ctx, dataset) + ctx.link(f"cloudg://datasets/{ds.name}/summary", f"{ds.name} summary") + return ds.summary() + + +def register(reg: Registry) -> None: + """Add this module's tools to ``reg``.""" + CATALOG.register(reg) diff --git a/cloudg/mcp/cli.py b/cloudg/mcp/cli.py new file mode 100644 index 0000000..48dc92a --- /dev/null +++ b/cloudg/mcp/cli.py @@ -0,0 +1,471 @@ +"""Command-line front end (``cloudg mcp``) for the cloudg MCP layer. + +Commands:: + + cloudg mcp serve run an MCP server (stdio / streamable HTTP / SSE) + cloudg mcp tools list tools (category, sensitivity, annotations) + cloudg mcp resources list resources and resource templates + cloudg mcp prompts list prompts + cloudg mcp call call a tool in-process and print its result + cloudg mcp read read a resource in-process + cloudg mcp config print a client configuration snippet + +The group is registered on the main ``cloudg`` CLI and also exposed as a +stand-alone entry point (``cloudg-mcp = "cloudg.mcp.cli:main"``) and as +``python -m cloudg.mcp``. + +stdout discipline: ``serve --transport stdio`` writes nothing but protocol +messages to stdout, and every other command writes only its result there +(tables or ``--json``), so output can be piped. Logs and notices go to +stderr. +""" + +from __future__ import annotations + +import asyncio +import json +import logging +import os +import sys +from pathlib import Path +from typing import Any + +import click + +from cloudg.mcp.cli_common import ( + AUTH_TOKENS_ENV, + _build_layer, + _echo_json, + _first_line, + _layer_options, + _principal, + _principal_options, + _setup_cli_logging, + _stdout_console, +) +from cloudg.mcp.cli_config import CLIENTS, client_config, config_cmd + +__all__ = ["CLIENTS", "client_config", "main", "mcp_group"] + + +# --------------------------------------------------------------------------- +# Group +# --------------------------------------------------------------------------- + + +@click.group("mcp") +def mcp_group() -> None: + """Model Context Protocol server: expose cloudg to AI assistants. + + \b + Quick start: + cloudg mcp serve # stdio (Claude Desktop / Code, Cursor...) + cloudg mcp serve --transport http # http://127.0.0.1:8765/mcp + cloudg mcp config --client claude-code # client configuration snippet + """ + + +# --------------------------------------------------------------------------- +# serve +# --------------------------------------------------------------------------- + + +@mcp_group.command("serve") +@click.option( + "--transport", + type=click.Choice(["stdio", "http", "streamable-http", "sse"]), + default="stdio", + show_default=True, + help="MCP transport.", +) +@click.option( + "--host", + default="127.0.0.1", + show_default=True, + help="HTTP bind address (keep 127.0.0.1 unless you also set --auth-token).", +) +@click.option("--port", type=int, default=8765, show_default=True, help="HTTP port.") +@click.option("--path", "http_path", default="/mcp", show_default=True, help="HTTP endpoint.") +@click.option( + "--flavor", + type=click.Choice(["auto", "native", "sdk", "fastmcp"]), + default="auto", + show_default=True, + help="Server implementation: mcp SDK, standalone fastmcp, or the built-in " + "dependency-free server.", +) +@click.option( + "--auth-token", + "auth_tokens", + multiple=True, + metavar="TOKEN[:ROLES[:ID]]", + envvar=AUTH_TOKENS_ENV, + show_envvar=True, + help="Require 'Authorization: Bearer TOKEN' (HTTP). ROLES is a comma list; " + "the principal also gets the 'default' role. TOKEN may be env:VAR. A token " + "containing ':' needs the full form TOKEN:ROLES:ID (fields may be empty). " + "Repeatable.", +) +@click.option( + "--allowed-origin", + "allowed_origins", + multiple=True, + help="Allowed browser Origin pattern, replacing the default (loopback origins). " + "Same-origin requests are accepted when the Host header is checked (loopback " + "bind or --allowed-host); the Origin check is on for every bind and flavor. " + "Repeatable.", +) +@click.option( + "--allowed-host", + "allowed_hosts", + multiple=True, + help="Allowed Host header pattern. Default: loopback names on a loopback bind; " + "on other binds the Host header is not checked unless this is set " + "(Origin still is). Repeatable.", +) +@click.option( + "--cors-origin", + "cors_origins", + multiple=True, + help="Enable CORS for this browser origin (native flavor only). Repeatable.", +) +@click.option("--json-response", is_flag=True, help="Answer HTTP POSTs with JSON, never SSE.") +@click.option("--stateless", is_flag=True, help="Session-less Streamable HTTP.") +@click.option( + "--cache-ttl", + type=float, + default=None, + help="Cache read-only idempotent tool results for N seconds.", +) +@click.option("--max-concurrency", type=int, default=None, help="Max concurrent tool calls.") +@click.option( + "--page-size", + type=int, + default=100, + show_default=True, + help="Items per page for list results (native flavor).", +) +@click.option( + "--log-level", + default="WARNING", + show_default=True, + type=click.Choice(["DEBUG", "INFO", "WARNING", "ERROR"], case_sensitive=False), + help="Server log level (logs go to stderr).", +) +@_layer_options +@click.pass_context +def serve_cmd(ctx: click.Context, **opts: Any) -> None: + """Run an MCP server exposing cloudg's tools, resources and prompts.""" + from cloudg.mcp.server import ( + ServeOptions, + _ensure_stderr_logging, + serve, + validate_serve_options, + ) + + serve_opts = ServeOptions( + flavor=opts.pop("flavor"), + host=opts.pop("host"), + port=opts.pop("port"), + path=opts.pop("http_path"), + auth_tokens=_auth_tokens(opts.pop("auth_tokens")), + allowed_origins=list(opts.pop("allowed_origins")) or None, + allowed_hosts=list(opts.pop("allowed_hosts")) or None, + cors_origins=list(opts.pop("cors_origins")), + json_response=opts.pop("json_response"), + stateless=opts.pop("stateless"), + page_size=opts.pop("page_size"), + log_level=opts.pop("log_level"), + ) + transport = opts.pop("transport") + try: + validate_serve_options( + serve_opts.flavor, + transport, + cors_origins=serve_opts.cors_origins, + json_response=serve_opts.json_response, + stateless=serve_opts.stateless, + ) + except ValueError as exc: + raise click.UsageError(str(exc)) from exc + _ensure_stderr_logging(serve_opts.log_level) + extra = {"cache_ttl": opts.pop("cache_ttl"), "max_concurrency": opts.pop("max_concurrency")} + layer = _build_layer(ctx, opts, **extra) + serve(layer, transport, options=serve_opts) + + +def _auth_tokens(specs: tuple[str, ...]) -> list[str]: + """Token specs from ``--auth-token`` / ``$CLOUDG_MCP_AUTH_TOKENS``, + validated. Given but empty (an unset ``$VAR`` expands to nothing) is an + error: serving without the authentication the operator asked for would + fail open.""" + tokens = [t for spec in specs for t in spec.split()] + env_blank = os.environ.get(AUTH_TOKENS_ENV) is not None and not tokens + if (specs or env_blank) and not tokens: + raise click.BadParameter( + "an empty token was given (is the variable it comes from unset?)", + param_hint="--auth-token", + ) + if tokens: + from cloudg.mcp.native.auth import TokenAuth + + try: + TokenAuth.from_specs(tokens) + except ValueError as exc: + raise click.BadParameter(str(exc), param_hint="--auth-token") from exc + return tokens + + +# --------------------------------------------------------------------------- +# Listing +# --------------------------------------------------------------------------- + + +def _annotation_flags(ann: dict[str, Any]) -> str: + flags = [] + for key, label in ( + ("readOnlyHint", "read-only"), + ("destructiveHint", "destructive"), + ("idempotentHint", "idempotent"), + ("openWorldHint", "open-world"), + ): + if ann.get(key) is True: + flags.append(label) + return ", ".join(flags) + + +@mcp_group.command("tools") +@click.option("--json", "as_json", is_flag=True, help="Print the MCP wire definitions as JSON.") +@click.option("--category", "categories", multiple=True, help="Filter by category.") +@_principal_options +@_layer_options +@click.pass_context +def tools_cmd( + ctx: click.Context, + as_json: bool, + categories: tuple[str, ...], + roles: tuple[str, ...], + principal_id: str | None, + **layer_opts: Any, +) -> None: + """List the tools the active policy exposes.""" + _setup_cli_logging() + layer = _build_layer(ctx, layer_opts) + principal = _principal(roles, principal_id) + tools = [ + t + for t in layer.tools_wire(principal) + if not categories or t.get("_meta", {}).get("cloudg/category") in categories + ] + if as_json: + _echo_json(tools) + return + from rich.table import Table + + table = Table(title=f"cloudg MCP tools ({len(tools)})", show_lines=False) + for col in ("Tool", "Category", "Sensitivity", "Hints", "Description"): + table.add_column(col, overflow="fold") + for t in tools: + meta = t.get("_meta", {}) + table.add_row( + t["name"], + str(meta.get("cloudg/category", "")), + str(meta.get("cloudg/sensitivity", "")), + _annotation_flags(t.get("annotations", {})), + _first_line(t.get("description")), + ) + _stdout_console().print(table) + + +@mcp_group.command("resources") +@click.option("--json", "as_json", is_flag=True, help="Print the MCP wire definitions as JSON.") +@_principal_options +@_layer_options +@click.pass_context +def resources_cmd( + ctx: click.Context, + as_json: bool, + roles: tuple[str, ...], + principal_id: str | None, + **layer_opts: Any, +) -> None: + """List resources and resource templates.""" + _setup_cli_logging() + layer = _build_layer(ctx, layer_opts) + principal = _principal(roles, principal_id) + resources = layer.resources_wire(principal) + templates = layer.resource_templates_wire(principal) + if as_json: + _echo_json({"resources": resources, "resourceTemplates": templates}) + return + from rich.table import Table + + table = Table(title=f"cloudg MCP resources ({len(resources)} + {len(templates)} templates)") + for col in ("URI / template", "Name", "Category", "Sensitivity", "MIME", "Description"): + table.add_column(col, overflow="fold") + for r in resources: + _add_resource_row(table, r["uri"], r) + for r in templates: + _add_resource_row(table, r["uriTemplate"], r) + _stdout_console().print(table) + + +def _add_resource_row(table: Any, uri: str, record: dict[str, Any]) -> None: + """One ``cloudg mcp resources`` table row for a resource or template.""" + meta = record.get("_meta", {}) + table.add_row( + uri, + record["name"], + str(meta.get("cloudg/category", "")), + str(meta.get("cloudg/sensitivity", "")), + record.get("mimeType", ""), + _first_line(record.get("description")), + ) + + +@mcp_group.command("prompts") +@click.option("--json", "as_json", is_flag=True, help="Print the MCP wire definitions as JSON.") +@_principal_options +@_layer_options +@click.pass_context +def prompts_cmd( + ctx: click.Context, + as_json: bool, + roles: tuple[str, ...], + principal_id: str | None, + **layer_opts: Any, +) -> None: + """List prompts.""" + _setup_cli_logging() + layer = _build_layer(ctx, layer_opts) + prompts = layer.prompts_wire(_principal(roles, principal_id)) + if as_json: + _echo_json(prompts) + return + from rich.table import Table + + table = Table(title=f"cloudg MCP prompts ({len(prompts)})") + for col in ("Prompt", "Category", "Arguments", "Description"): + table.add_column(col, overflow="fold") + for p in prompts: + args = ", ".join( + a["name"] + ("*" if a.get("required") else "") for a in p.get("arguments", []) + ) + table.add_row( + p["name"], + str(p.get("_meta", {}).get("cloudg/category", "")), + args, + _first_line(p.get("description")), + ) + _stdout_console().print(table) + + +# --------------------------------------------------------------------------- +# In-process calls +# --------------------------------------------------------------------------- + + +def _parse_args(raw: str | None) -> dict[str, Any]: + if not raw: + return {} + if raw.startswith("@"): + raw = Path(raw[1:]).expanduser().read_text(encoding="utf-8") + elif raw == "-": + raw = sys.stdin.read() + try: + value = json.loads(raw) + except ValueError as exc: + raise click.BadParameter(f"not valid JSON: {exc}", param_hint="--args") from exc + if not isinstance(value, dict): + raise click.BadParameter("must be a JSON object", param_hint="--args") + return value + + +@mcp_group.command("call") +@click.argument("tool") +@click.option( + "--args", + "raw_args", + default=None, + metavar="JSON|@FILE|-", + help="Tool arguments as a JSON object, @file.json, or - for stdin.", +) +@click.option("--raw", is_flag=True, help="Print the full MCP CallToolResult.") +@_principal_options +@_layer_options +@click.pass_context +def call_cmd( + ctx: click.Context, + tool: str, + raw_args: str | None, + raw: bool, + roles: tuple[str, ...], + principal_id: str | None, + **layer_opts: Any, +) -> None: + """Call TOOL in-process (policies and transforms apply) and print the result. + + Exits with status 1 when the tool reports an error. + """ + from cloudg.mcp.core import MCPLayerError + + _setup_cli_logging() + arguments = _parse_args(raw_args) + layer = _build_layer(ctx, layer_opts) + try: + result = asyncio.run( + layer.call_tool(tool, arguments, principal=_principal(roles, principal_id)) + ) + except MCPLayerError as exc: + raise click.ClickException(exc.message) from exc + wire = result.to_wire() + if raw: + _echo_json(wire) + elif result.structured is not None: + _echo_json(result.structured) + else: + for block in wire["content"]: + click.echo(block.get("text") if block.get("type") == "text" else json.dumps(block)) + if result.is_error: + ctx.exit(1) + + +@mcp_group.command("read") +@click.argument("uri") +@_principal_options +@_layer_options +@click.pass_context +def read_cmd( + ctx: click.Context, + uri: str, + roles: tuple[str, ...], + principal_id: str | None, + **layer_opts: Any, +) -> None: + """Read resource URI in-process and print its contents.""" + from cloudg.mcp.core import MCPLayerError, TextResourceContents + + _setup_cli_logging() + layer = _build_layer(ctx, layer_opts) + try: + contents = asyncio.run(layer.read_resource(uri, principal=_principal(roles, principal_id))) + except MCPLayerError as exc: + raise click.ClickException(exc.message) from exc + for item in contents: + if isinstance(item, TextResourceContents): + click.echo(item.text) + else: + _echo_json(item.to_wire()) + + +mcp_group.add_command(config_cmd) + + +def main(argv: list[str] | None = None) -> None: + """Entry point for ``cloudg-mcp`` / ``python -m cloudg.mcp``.""" + logging.captureWarnings(True) + mcp_group.main(args=argv, prog_name="cloudg-mcp") + + +if __name__ == "__main__": # pragma: no cover + main() diff --git a/cloudg/mcp/cli_common.py b/cloudg/mcp/cli_common.py new file mode 100644 index 0000000..fd7131a --- /dev/null +++ b/cloudg/mcp/cli_common.py @@ -0,0 +1,238 @@ +"""Options and helpers shared by the ``cloudg mcp`` commands +(:mod:`cloudg.mcp.cli`): layer-shaping options, ``--registry`` loading, +principal options and output helpers.""" + +from __future__ import annotations + +import json +import re +import sys +from typing import Any, Callable + +import click + +from cloudg.mcp.context import Principal + +AUTH_TOKENS_ENV = "CLOUDG_MCP_AUTH_TOKENS" + + +# --------------------------------------------------------------------------- +# Shared options +# --------------------------------------------------------------------------- + + +# Options that shape the layer, shared by every command. ``click.option`` +# builds a fresh ``Option`` each time a decorator is applied, so one tuple +# can decorate several commands. +_LAYER_OPTIONS = ( + click.option( + "-c", + "--config", + "config_path", + default=None, + help="cloudg config.yaml (defaults to the parent command's config).", + ), + click.option( + "--policy", + default=None, + envvar="CLOUDG_MCP_POLICY", + show_envvar=True, + help="Policy profile name, or a YAML/JSON policy file.", + ), + click.option( + "--dataset", + "datasets", + multiple=True, + metavar="[NAME=]PATH", + help="Preload a dataset (findings.json, inventory-map.json...). Repeatable.", + ), + click.option("--prefix", default="", help="Prefix for tool and prompt names."), + click.option( + "--include-category", + "include_categories", + multiple=True, + help="Only expose these categories. Repeatable.", + ), + click.option( + "--exclude-category", + "exclude_categories", + multiple=True, + help="Hide these categories. Repeatable.", + ), + click.option( + "--include-tool", + "include_tools", + multiple=True, + help="Only expose these tools. Repeatable.", + ), + click.option( + "--exclude-tool", "exclude_tools", multiple=True, help="Hide these tools. Repeatable." + ), + click.option( + "--read-only", + is_flag=True, + help="Drop tools that write files, call cloud APIs, run scanners or are destructive.", + ), + click.option( + "--audit-log", + type=click.Path(dir_okay=False), + default=None, + envvar="CLOUDG_MCP_AUDIT_LOG", + show_envvar=True, + help="Append a JSONL audit trail (argument values are hashed).", + ), + click.option( + "--timeout", + type=float, + default=300.0, + show_default=True, + help="Per-call timeout in seconds.", + ), + click.option( + "--registry", + "registry_ref", + default=None, + metavar="MODULE:ATTR", + help="Serve a custom Registry (or a zero-argument factory returning one) " + "instead of the built-in catalog.", + ), +) + + +def _layer_options(fn: Callable[..., Any]) -> Callable[..., Any]: + """Options that shape the layer (shared by every command).""" + for option in reversed(_LAYER_OPTIONS): + fn = option(fn) + return fn + + +def _parent_config(ctx: click.Context) -> Any: + root = ctx.find_root() + obj = getattr(root, "obj", None) + try: + from cloudg.config import CloudGConfig + + return obj if isinstance(obj, CloudGConfig) else None + except Exception: # pragma: no cover + return None + + +_DOTTED_NAME = re.compile(r"[A-Za-z_][A-Za-z0-9_]*(?:\.[A-Za-z_][A-Za-z0-9_]*)*") + + +def _parse_registry_ref(ref: str) -> tuple[str, list[str]]: + """Split ``MODULE[:ATTR]`` after checking both are dotted Python + identifiers (no relative imports, paths or expressions).""" + module_name, _, attr = ref.partition(":") + attr = attr or "default_registry" + for part in (module_name, attr): + if not _DOTTED_NAME.fullmatch(part): + raise click.BadParameter( + f"{ref!r} is not MODULE:ATTR with dotted Python names", param_hint="--registry" + ) + if any(name.startswith("_") for name in attr.split(".")): + raise click.BadParameter( + f"{ref!r}: the attribute must be a public name", param_hint="--registry" + ) + return module_name, attr.split(".") + + +def _load_registry(ref: str | None) -> Any: + """Load ``--registry MODULE:ATTR``: a :class:`~cloudg.mcp.core.Registry` + or a zero-argument factory returning one. + + This imports and runs operator code, exactly like ``python -m MODULE`` + would: the value comes from the command line or config of whoever starts + the server, never from an MCP client. It is validated as a plain dotted + module name and a public attribute path before anything is imported. + """ + if not ref: + return None + import importlib + + from cloudg.mcp.core import Registry + + module_name, attrs = _parse_registry_ref(ref) + try: + # trusted operator input, validated by _parse_registry_ref above + obj: Any = importlib.import_module(module_name) # nosemgrep - validated operator input + for part in attrs: + obj = getattr(obj, part) + except (ImportError, AttributeError) as exc: + raise click.BadParameter(f"cannot import {ref!r}: {exc}", param_hint="--registry") from exc + if not isinstance(obj, Registry) and callable(obj): + obj = obj() + if not isinstance(obj, Registry): + raise click.BadParameter( + f"{ref!r} is not a cloudg.mcp.core.Registry (or a factory returning one)", + param_hint="--registry", + ) + return obj + + +def _build_layer(ctx: click.Context, opts: dict[str, Any], **extra: Any) -> Any: + from cloudg.mcp.server import create_layer_from_options + + registry = _load_registry(opts.get("registry_ref")) + try: + return create_layer_from_options( + config=None if opts.get("config_path") else _parent_config(ctx), + config_path=opts.get("config_path"), + policy=opts.get("policy"), + datasets=opts.get("datasets") or (), + prefix=opts.get("prefix") or "", + include_categories=opts.get("include_categories") or None, + exclude_categories=opts.get("exclude_categories") or None, + include_tools=opts.get("include_tools") or None, + exclude_tools=opts.get("exclude_tools") or None, + read_only=bool(opts.get("read_only")), + audit_log=opts.get("audit_log"), + default_timeout=opts.get("timeout"), + registry=registry, + **extra, + ) + except (FileNotFoundError, ValueError) as exc: + raise click.ClickException(str(exc)) from exc + + +def _principal(roles: tuple[str, ...], principal_id: str | None) -> Principal: + if not roles and not principal_id: + return Principal.local() + return Principal(id=principal_id or "cli", roles=set(roles) or {"default"}) + + +def _principal_options(fn: Callable[..., Any]) -> Callable[..., Any]: + fn = click.option( + "--role", + "roles", + multiple=True, + help="Call as a principal with this role (repeatable; default: the local user).", + )(fn) + fn = click.option("--principal-id", default=None, help="Principal id to call as.")(fn) + return fn + + +def _stdout_console() -> Any: + from rich.console import Console + + try: + tty = sys.stdout.isatty() + except (AttributeError, ValueError): + tty = False + # piped / captured output: wide enough that URIs and names are not folded + return Console(file=sys.stdout, soft_wrap=False, width=None if tty else 160) + + +def _setup_cli_logging(verbose: bool = False) -> None: + from cloudg.mcp.server import _ensure_stderr_logging + + _ensure_stderr_logging("DEBUG" if verbose else "WARNING") + + +def _echo_json(data: Any) -> None: + click.echo(json.dumps(data, indent=2, ensure_ascii=False, default=str)) + + +def _first_line(text: str | None, width: int = 90) -> str: + line = (text or "").strip().splitlines()[0] if (text or "").strip() else "" + return line if len(line) <= width else line[: width - 1] + "…" diff --git a/cloudg/mcp/cli_config.py b/cloudg/mcp/cli_config.py new file mode 100644 index 0000000..1cebb8a --- /dev/null +++ b/cloudg/mcp/cli_config.py @@ -0,0 +1,148 @@ +"""``cloudg mcp config``: ready-to-paste MCP client configuration snippets +(re-exported by :mod:`cloudg.mcp.cli`, the documented import path of +:func:`client_config`).""" + +from __future__ import annotations + +import shlex +import shutil +import sys +from pathlib import Path +from typing import Any + +import click + +from cloudg.mcp.cli_common import _echo_json + +__all__ = ["CLIENTS", "client_config", "config_cmd"] + +CLIENTS = ("claude-desktop", "claude-code", "cursor", "vscode") + + +# --------------------------------------------------------------------------- +# Client configuration snippets +# --------------------------------------------------------------------------- + + +def _server_command(command: str) -> list[str]: + if command != "auto": + return shlex.split(command) + exe = shutil.which("cloudg-mcp") + if exe: + return [exe, "serve"] + exe = shutil.which("cloudg") + if exe: + return [exe, "mcp", "serve"] + return [sys.executable, "-m", "cloudg.mcp", "serve"] + + +def client_config( + client: str, + *, + name: str = "cloudg", + transport: str = "stdio", + url: str = "http://127.0.0.1:8765/mcp", + command: list[str] | None = None, + serve_args: list[str] | None = None, + token_env: str | None = None, +) -> dict[str, Any]: + """Build the JSON snippet for ``client`` (see ``cloudg mcp config``).""" + if client not in CLIENTS: + raise ValueError(f"Unknown client {client!r}; choose from {CLIENTS}") + argv = list(command or _server_command("auto")) + list(serve_args or []) + exe, args = argv[0], argv[1:] + if transport == "stdio": + if client == "vscode": + return {"servers": {name: {"type": "stdio", "command": exe, "args": args}}} + entry: dict[str, Any] = {"command": exe, "args": args} + if client == "claude-code": + entry = {"type": "stdio", **entry, "env": {}} + return {"mcpServers": {name: entry}} + headers: dict[str, str] = {} + if token_env: + ref = "${" + token_env + "}" if client == "claude-code" else "${env:" + token_env + "}" + headers = {"Authorization": f"Bearer {ref}"} + if client == "claude-desktop": + # Claude Desktop's config file launches local commands; bridge remote + # servers with mcp-remote. + bridge = ["-y", "mcp-remote", url] + if token_env: + bridge += ["--header", f"Authorization:Bearer ${{{token_env}}}"] + return {"mcpServers": {name: {"command": "npx", "args": bridge}}} + remote: dict[str, Any] = {"url": url} + if client in ("claude-code", "vscode"): + remote = {"type": "http", **remote} + if headers: + remote["headers"] = headers + key = "servers" if client == "vscode" else "mcpServers" + return {key: {name: remote}} + + +_CONFIG_HINTS = { + "claude-desktop": "Merge into claude_desktop_config.json " + "(Settings > Developer > Edit Config), then restart Claude Desktop.", + "claude-code": "Save as .mcp.json in your project root, or run: claude mcp add-json " + " ''.", + "cursor": "Merge into ~/.cursor/mcp.json (global) or .cursor/mcp.json (project).", + "vscode": 'Save as .vscode/mcp.json (or add under "mcp" in settings.json).', +} + + +@click.command("config") +@click.option( + "--client", + type=click.Choice(CLIENTS), + default="claude-desktop", + show_default=True, + help="Target MCP client.", +) +@click.option("--name", default="cloudg", show_default=True, help="Server name in the config.") +@click.option( + "--transport", + type=click.Choice(["stdio", "http"]), + default="stdio", + show_default=True, + help="Launch locally (stdio) or connect to a running server.", +) +@click.option( + "--url", + default="http://127.0.0.1:8765/mcp", + show_default=True, + help="Server URL for --transport http.", +) +@click.option( + "--command", + "command", + default="auto", + show_default=True, + help="Server launch command (auto: cloudg-mcp / cloudg / python -m cloudg.mcp).", +) +@click.option("--policy", default=None, help="Add --policy to the launch arguments.") +@click.option("--dataset", "datasets", multiple=True, help="Add --dataset (repeatable).") +@click.option("--read-only", is_flag=True, help="Add --read-only.") +@click.option( + "--token-env", default=None, metavar="VAR", help="For http: send 'Authorization: Bearer $VAR'." +) +def config_cmd(**opts: Any) -> None: + """Print a ready-to-paste MCP client configuration snippet.""" + serve_args: list[str] = [] + if opts["policy"]: + serve_args += ["--policy", opts["policy"]] + for ds in opts["datasets"]: + p = ds.partition("=") + path = str(Path(p[2] if p[1] else ds).expanduser().resolve()) + serve_args += ["--dataset", f"{p[0]}={path}" if p[1] else path] + if opts["read_only"]: + serve_args.append("--read-only") + client = opts["client"] + snippet = client_config( + client, + name=opts["name"], + transport=opts["transport"], + url=opts["url"], + command=_server_command(opts["command"]), + serve_args=serve_args, + token_env=opts["token_env"], + ) + _echo_json(snippet) + click.echo(_CONFIG_HINTS[client], err=True) diff --git a/cloudg/mcp/content.py b/cloudg/mcp/content.py new file mode 100644 index 0000000..36930cd --- /dev/null +++ b/cloudg/mcp/content.py @@ -0,0 +1,275 @@ +"""MCP content blocks, annotations and results for the cloudg MCP layer. + +Plain dataclasses that know nothing about any MCP server implementation. +Their ``to_wire()`` methods produce the MCP specification's camelCase +shapes. Everything here is re-exported by :mod:`cloudg.mcp.core`, which is +the documented import path. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Literal + +__all__ = [ + "BlobResourceContents", + "Content", + "ContentAnnotations", + "EmbeddedResource", + "Icon", + "ImageContent", + "PromptMessage", + "PromptResult", + "ResourceContents", + "ResourceLink", + "Role", + "TextContent", + "TextResourceContents", + "ToolAnnotations", + "ToolResult", +] + +Role = Literal["user", "assistant"] + + +# --------------------------------------------------------------------------- +# Annotations +# --------------------------------------------------------------------------- + + +@dataclass +class ToolAnnotations: + """MCP tool behaviour hints (spec ``ToolAnnotations``).""" + + title: str | None = None + read_only: bool | None = None + destructive: bool | None = None + idempotent: bool | None = None + open_world: bool | None = None + + def to_wire(self) -> dict[str, Any]: + out: dict[str, Any] = {} + if self.title is not None: + out["title"] = self.title + if self.read_only is not None: + out["readOnlyHint"] = self.read_only + if self.destructive is not None: + out["destructiveHint"] = self.destructive + if self.idempotent is not None: + out["idempotentHint"] = self.idempotent + if self.open_world is not None: + out["openWorldHint"] = self.open_world + return out + + +@dataclass +class ContentAnnotations: + """MCP content / resource annotations (spec ``Annotations``).""" + + audience: list[Role] | None = None + priority: float | None = None # 0.0 (optional) .. 1.0 (required) + last_modified: str | None = None # ISO 8601 + + def to_wire(self) -> dict[str, Any]: + out: dict[str, Any] = {} + if self.audience is not None: + out["audience"] = list(self.audience) + if self.priority is not None: + out["priority"] = max(0.0, min(1.0, float(self.priority))) + if self.last_modified is not None: + out["lastModified"] = self.last_modified + return out + + +@dataclass +class Icon: + """An icon for a tool, resource, prompt or the server (MCP 2025-11-25). + Prefer ``data:`` URIs so clients need not fetch anything.""" + + src: str + mime_type: str | None = None + sizes: list[str] | None = None # e.g. ["48x48"] + theme: Literal["light", "dark"] | None = None + + def to_wire(self) -> dict[str, Any]: + out: dict[str, Any] = {"src": self.src} + if self.mime_type: + out["mimeType"] = self.mime_type + if self.sizes: + out["sizes"] = list(self.sizes) + if self.theme: + out["theme"] = self.theme + return out + + +def _icons_wire(icons: list[Icon] | None) -> list[dict[str, Any]] | None: + return [i.to_wire() for i in icons] if icons else None + + +# --------------------------------------------------------------------------- +# Content and results +# --------------------------------------------------------------------------- + + +@dataclass +class TextContent: + text: str + annotations: ContentAnnotations | None = None + meta: dict[str, Any] | None = None + + def to_wire(self) -> dict[str, Any]: + out: dict[str, Any] = {"type": "text", "text": self.text} + _add_common(out, self.annotations, self.meta) + return out + + +@dataclass +class ImageContent: + data: str # base64 + mime_type: str + annotations: ContentAnnotations | None = None + meta: dict[str, Any] | None = None + + def to_wire(self) -> dict[str, Any]: + out: dict[str, Any] = {"type": "image", "data": self.data, "mimeType": self.mime_type} + _add_common(out, self.annotations, self.meta) + return out + + +@dataclass +class ResourceLink: + """A pointer to a resource the client may read (spec ``ResourceLink``).""" + + uri: str + name: str + title: str | None = None + description: str | None = None + mime_type: str | None = None + size: int | None = None + annotations: ContentAnnotations | None = None + meta: dict[str, Any] | None = None + + def to_wire(self) -> dict[str, Any]: + out: dict[str, Any] = {"type": "resource_link", "uri": self.uri, "name": self.name} + if self.title: + out["title"] = self.title + if self.description: + out["description"] = self.description + if self.mime_type: + out["mimeType"] = self.mime_type + if self.size is not None: + out["size"] = self.size + _add_common(out, self.annotations, self.meta) + return out + + +@dataclass +class TextResourceContents: + uri: str + text: str + mime_type: str | None = "application/json" + meta: dict[str, Any] | None = None + + def to_wire(self) -> dict[str, Any]: + out: dict[str, Any] = {"uri": self.uri, "text": self.text} + if self.mime_type: + out["mimeType"] = self.mime_type + if self.meta: + out["_meta"] = self.meta + return out + + +@dataclass +class BlobResourceContents: + uri: str + blob: str # base64 + mime_type: str | None = None + meta: dict[str, Any] | None = None + + def to_wire(self) -> dict[str, Any]: + out: dict[str, Any] = {"uri": self.uri, "blob": self.blob} + if self.mime_type: + out["mimeType"] = self.mime_type + if self.meta: + out["_meta"] = self.meta + return out + + +ResourceContents = TextResourceContents | BlobResourceContents + + +@dataclass +class EmbeddedResource: + resource: ResourceContents + annotations: ContentAnnotations | None = None + meta: dict[str, Any] | None = None + + def to_wire(self) -> dict[str, Any]: + out: dict[str, Any] = {"type": "resource", "resource": self.resource.to_wire()} + _add_common(out, self.annotations, self.meta) + return out + + +Content = TextContent | ImageContent | ResourceLink | EmbeddedResource + + +def _add_common( + out: dict[str, Any], annotations: ContentAnnotations | None, meta: dict[str, Any] | None +) -> None: + if annotations is not None: + wire = annotations.to_wire() + if wire: + out["annotations"] = wire + if meta: + out["_meta"] = meta + + +@dataclass +class ToolResult: + """Framework-agnostic ``CallToolResult``. + + ``structured`` is the JSON object returned as ``structuredContent``; the + layer also renders it into ``content`` as text so clients without + structured-output support still see it. + """ + + content: list[Content] = field(default_factory=list) + structured: dict[str, Any] | None = None + is_error: bool = False + meta: dict[str, Any] = field(default_factory=dict) + + def to_wire(self) -> dict[str, Any]: + out: dict[str, Any] = { + "content": [c.to_wire() for c in self.content], + "isError": self.is_error, + } + if self.structured is not None: + out["structuredContent"] = self.structured + if self.meta: + out["_meta"] = self.meta + return out + + @classmethod + def error(cls, message: str, **meta: Any) -> "ToolResult": + return cls(content=[TextContent(message)], is_error=True, meta=dict(meta)) + + +@dataclass +class PromptMessage: + role: Role + content: Content + + def to_wire(self) -> dict[str, Any]: + return {"role": self.role, "content": self.content.to_wire()} + + +@dataclass +class PromptResult: + messages: list[PromptMessage] + description: str | None = None + + def to_wire(self) -> dict[str, Any]: + out: dict[str, Any] = {"messages": [m.to_wire() for m in self.messages]} + if self.description: + out["description"] = self.description + return out diff --git a/cloudg/mcp/context.py b/cloudg/mcp/context.py new file mode 100644 index 0000000..f120436 --- /dev/null +++ b/cloudg/mcp/context.py @@ -0,0 +1,162 @@ +"""Per-call context handed to every tool, resource and prompt handler.""" + +from __future__ import annotations + +import asyncio +import logging +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Awaitable, Callable + +from cloudg.mcp.core import ContentAnnotations, ResourceLink + +if TYPE_CHECKING: # pragma: no cover + from cloudg.mcp.layer import CloudGMCPLayer + from cloudg.mcp.state import Workspace + +logger = logging.getLogger("cloudg.mcp") + +ProgressCallback = Callable[[float, float | None, str | None], Awaitable[None] | None] +LogCallback = Callable[[str, Any, str | None], Awaitable[None] | None] + +# MCP (RFC 5424) log levels -> stdlib levels +LOG_LEVELS = { + "debug": logging.DEBUG, + "info": logging.INFO, + "notice": logging.INFO + 5, + "warning": logging.WARNING, + "error": logging.ERROR, + "critical": logging.CRITICAL, + "alert": logging.CRITICAL + 5, + "emergency": logging.CRITICAL + 10, +} + + +@dataclass +class Principal: + """Who is calling. Adapters derive it from the transport (HTTP auth + header, stdio = local user); policies key access and transform rules + off ``roles``.""" + + id: str = "anonymous" + roles: set[str] = field(default_factory=lambda: {"default"}) + attributes: dict[str, Any] = field(default_factory=dict) + + @classmethod + def local(cls) -> "Principal": + return cls(id="local", roles={"default", "local"}) + + +@dataclass +class ToolContext: + """Context for one MCP request. + + Handlers use it to reach the shared :class:`~cloudg.mcp.state.Workspace` + (loaded datasets and cached analyses), report progress, emit log + notifications, and attach resource links to their result. + """ + + layer: "CloudGMCPLayer" + principal: Principal = field(default_factory=Principal) + kind: str = "tool" # tool | resource | prompt | completion + name: str = "" + request_id: Any = None + # Adapter-supplied callbacks; no-ops when the transport lacks them + progress_callback: ProgressCallback | None = None + log_callback: LogCallback | None = None + # Opaque adapter session (mcp ServerSession / fastmcp Context / None) + session: Any = None + links: list[ResourceLink] = field(default_factory=list) + meta: dict[str, Any] = field(default_factory=dict) + # Event loop serving the request; lets sync handlers (which run in a + # worker thread) report progress via report_progress_sync + loop: asyncio.AbstractEventLoop | None = None + # The caller's output transform for this call (set by the layer): + # progress messages and log data reach the client through it, like the + # result does + transform_output: Callable[[Any], Any] | None = None + _last_progress: float | None = field(default=None, repr=False) + + @property + def workspace(self) -> "Workspace": + return self.layer.workspace + + @property + def config(self) -> Any: + return self.layer.workspace.config + + async def report_progress( + self, progress: float, total: float | None = None, message: str | None = None + ) -> None: + if self.progress_callback is None: + return + # MCP requires progress to increase with every notification + if self._last_progress is not None and progress <= self._last_progress: + return + self._last_progress = progress + try: + if message: + message = str(self._outgoing(message)) + res = self.progress_callback(progress, total, message) + if hasattr(res, "__await__"): + await res # type: ignore[misc] + except Exception: # progress must never break a tool + logger.debug("progress callback failed", exc_info=True) + + def report_progress_sync( + self, progress: float, total: float | None = None, message: str | None = None + ) -> None: + """Progress from a sync handler running in a worker thread.""" + loop = self.loop + if self.progress_callback is None or loop is None or loop.is_closed(): + return + try: + asyncio.run_coroutine_threadsafe(self.report_progress(progress, total, message), loop) + except RuntimeError: + logger.debug("progress dropped: loop not running", exc_info=True) + + async def log(self, level: str, data: Any, logger_name: str | None = "cloudg") -> None: + """Send an MCP ``notifications/message`` (and log locally). + ``level`` is an MCP level: debug, info, notice, warning, error, + critical, alert or emergency.""" + level = str(level).lower() + if level not in LOG_LEVELS: + level = "info" + logger.log(LOG_LEVELS[level], "%s", data) + if self.log_callback is None: + return + try: + res = self.log_callback(level, self._outgoing(data), logger_name) + if hasattr(res, "__await__"): + await res # type: ignore[misc] + except Exception: # a log notification must never break a call + logger.debug("log notification dropped", exc_info=True) + + def _outgoing(self, value: Any) -> Any: + """``value`` as the client may see it. A transform that refuses the + value (a blocked secret) raises, and the notification is dropped.""" + if self.transform_output is None: + return value + return self.transform_output(value) + + def link( + self, + uri: str, + name: str, + *, + title: str | None = None, + description: str | None = None, + mime_type: str | None = "application/json", + priority: float | None = None, + ) -> None: + """Attach a ``resource_link`` content block to this call's result.""" + ann = ContentAnnotations(priority=priority) if priority is not None else None + self.links.append( + ResourceLink( + uri=uri, + name=name, + title=title, + description=description, + mime_type=mime_type, + annotations=ann, + ) + ) diff --git a/cloudg/mcp/core.py b/cloudg/mcp/core.py new file mode 100644 index 0000000..de515eb --- /dev/null +++ b/cloudg/mcp/core.py @@ -0,0 +1,643 @@ +"""Framework-agnostic MCP primitives for the cloudg MCP layer. + +Everything the layer exposes (tools, resources, resource templates and +prompts) is described here as plain dataclasses that know nothing about +any particular MCP server implementation. Adapters +(:mod:`cloudg.mcp.adapters`) translate these specs into the official +``mcp`` SDK (1.x ``FastMCP`` or 2.x ``MCPServer`` / low-level ``Server``), +the standalone ``fastmcp`` package, or the dependency-free native server. + +The wire shapes produced by the ``to_wire()`` methods follow the MCP +specification (camelCase keys), so adapters can feed them straight into +``mcp.types..model_validate``. +""" + +from __future__ import annotations + +import inspect +import re +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, Awaitable, Callable, Iterable + +# Re-exported: cloudg.mcp.core is the documented import path for these. +from cloudg.mcp.content import ( + BlobResourceContents, + Content, + ContentAnnotations, + EmbeddedResource, + Icon, + ImageContent, + PromptMessage, + PromptResult, + ResourceContents, + ResourceLink, + Role, + TextContent, + TextResourceContents, + ToolAnnotations, + ToolResult, + _icons_wire, +) + +__all__ = [ + "AccessDeniedError", + "BlobResourceContents", + "Capability", + "Content", + "ContentAnnotations", + "EmbeddedResource", + "Handler", + "Icon", + "ImageContent", + "InvalidArgumentsError", + "MCPLayerError", + "META_PREFIX", + "NotFoundError", + "PromptArgument", + "PromptMessage", + "PromptResult", + "PromptSpec", + "RateLimitedError", + "Registry", + "ResourceContents", + "ResourceLink", + "ResourceSpec", + "ResourceTemplateSpec", + "Role", + "SENSITIVITY_ORDER", + "Sensitivity", + "TextContent", + "TextResourceContents", + "ToolAnnotations", + "ToolResult", + "ToolSpec", + "maybe_await", + "render_json", +] + +# --------------------------------------------------------------------------- +# Errors +# --------------------------------------------------------------------------- + + +# JSON-RPC codes. -32768..-32000 belong to JSON-RPC and -32020..-32099 are +# reserved by MCP, so application errors use -31xxx. +META_PREFIX = "cloudg/" + + +class MCPLayerError(Exception): + """Base error for the MCP layer. ``code`` is a JSON-RPC error code.""" + + code: int = -32603 + + def __init__(self, message: str, *, data: Any = None) -> None: + super().__init__(message) + self.message = message + self.data = data + + +class InvalidArgumentsError(MCPLayerError): + """Tool / prompt arguments failed validation.""" + + code = -32602 + + +class NotFoundError(MCPLayerError): + """Unknown tool, prompt, resource URI or dataset. MCP 2026-07-28 maps + unknown resources to Invalid Params; raised inside a tool handler it + becomes an ``isError`` result instead.""" + + code = -32602 + + +class AccessDeniedError(MCPLayerError): + """The active policy forbids this operation for the caller.""" + + code = -31001 + + +class RateLimitedError(MCPLayerError): + """The caller exceeded a configured rate limit.""" + + code = -31029 + + +# --------------------------------------------------------------------------- +# Enums +# --------------------------------------------------------------------------- + + +class Sensitivity(str, Enum): + """How sensitive a primitive's output is. Policies use this to decide + which transforms (redaction, pseudonymisation...) apply and whether the + primitive is exposed at all.""" + + PUBLIC = "public" # schema/catalog data, no tenant information + INTERNAL = "internal" # aggregate counts, summaries + CONFIDENTIAL = "confidential" # asset identifiers, topology, findings + RESTRICTED = "restricted" # raw metadata, policies, anything that may hold secrets + + +SENSITIVITY_ORDER = {s: i for i, s in enumerate(Sensitivity)} + + +class Capability(str, Enum): + """Side effects a tool needs. Policies can deny whole capabilities + (e.g. a read-only deployment denies CLOUD_ACCESS and WRITE_FS).""" + + READ_STATE = "read_state" # reads loaded datasets only + WRITE_STATE = "write_state" # mutates the in-memory workspace + READ_FS = "read_fs" # reads files from disk + WRITE_FS = "write_fs" # writes files to disk + CLOUD_ACCESS = "cloud_access" # calls cloud provider APIs with credentials + EXEC = "exec" # spawns external processes (scanners) + REVEAL = "reveal" # reverses pseudonymisation / reveals hidden values + + +# --------------------------------------------------------------------------- +# Specs +# --------------------------------------------------------------------------- + +Handler = Callable[..., Any | Awaitable[Any]] + +_NAME_RE = re.compile(r"^[A-Za-z0-9_.\-]{1,128}$") + + +@dataclass +class _BaseSpec: + name: str + handler: Handler + title: str | None = None + description: str = "" + category: str = "general" + tags: set[str] = field(default_factory=set) + sensitivity: Sensitivity = Sensitivity.CONFIDENTIAL + capabilities: set[Capability] = field(default_factory=lambda: {Capability.READ_STATE}) + meta: dict[str, Any] = field(default_factory=dict) + icons: list[Icon] | None = None + + def __post_init__(self) -> None: + if not self.description: + self.description = inspect.cleandoc(self.handler.__doc__ or "").strip() + + def _common_wire(self, out: dict[str, Any]) -> dict[str, Any]: + if self.title: + out["title"] = self.title + icons = _icons_wire(self.icons) + if icons: + out["icons"] = icons + out["_meta"] = { + f"{META_PREFIX}category": self.category, + f"{META_PREFIX}sensitivity": self.sensitivity.value, + **({f"{META_PREFIX}tags": sorted(self.tags)} if self.tags else {}), + **self.meta, + } + return out + + +@dataclass +class ToolSpec(_BaseSpec): + """A callable tool. + + ``handler(ctx, **arguments)`` may be sync or async. ``input_schema`` is + derived from the handler signature (see :mod:`cloudg.mcp.schema`) when + not given. The handler returns any JSON-serialisable value, a pydantic + model, or a :class:`ToolResult`; the layer normalises it. + """ + + input_schema: dict[str, Any] | None = None + output_schema: dict[str, Any] | None = None + annotations: ToolAnnotations = field(default_factory=ToolAnnotations) + # Transform profile overrides for this tool (see cloudg.mcp.transforms.policy) + transform_hints: dict[str, Any] = field(default_factory=dict) + timeout_seconds: float | None = None + + def __post_init__(self) -> None: + super().__post_init__() + if not _NAME_RE.match(self.name): + raise ValueError(f"Invalid tool name {self.name!r}") + if self.annotations.title is None and self.title: + self.annotations.title = self.title + self._derive_hints() + if self.input_schema is None: + from cloudg.mcp.schema import input_schema_for + + self.input_schema = input_schema_for(self.handler) + + def _derive_hints(self) -> None: + """Fill unset behaviour hints from the declared capabilities. + + Clients treat a missing hint pessimistically (destructive, + open-world), which makes them prompt for every read-only lookup. + """ + caps = self.capabilities + writes = caps & {Capability.WRITE_STATE, Capability.WRITE_FS, Capability.EXEC} + ann = self.annotations + if ann.read_only is None: + ann.read_only = not writes + if ann.destructive is None: + ann.destructive = False + # Tools that call cloud APIs or spawn processes reach outside the + # loaded data: open-world, and never idempotent by default (two + # collections a minute apart can differ), so caching and retry + # middleware leave them alone. + external = bool(caps & {Capability.CLOUD_ACCESS, Capability.EXEC}) + if ann.open_world is None: + ann.open_world = external + if ann.idempotent is None and ann.read_only and not external: + ann.idempotent = True + + def to_wire( + self, name: str | None = None, *, include_output_schema: bool = True + ) -> dict[str, Any]: + exposed = name or self.name + if not _NAME_RE.match(exposed): + raise ValueError(f"Invalid exposed tool name {exposed!r}") + out: dict[str, Any] = { + "name": exposed, + "description": self.description, + "inputSchema": self.input_schema, + } + if self.output_schema and include_output_schema: + out["outputSchema"] = self.output_schema + ann = self.annotations.to_wire() + if ann: + out["annotations"] = ann + return self._common_wire(out) + + +@dataclass +class ResourceSpec(_BaseSpec): + """A static resource at a fixed URI. ``handler(ctx)`` returns ``str``, + ``bytes``, a JSON-serialisable value, or a list of ResourceContents.""" + + uri: str = "" + mime_type: str = "application/json" + annotations: ContentAnnotations | None = None + size: int | None = None + + def to_wire(self) -> dict[str, Any]: + out: dict[str, Any] = {"uri": self.uri, "name": self.name, "mimeType": self.mime_type} + if self.description: + out["description"] = self.description + if self.size is not None: + out["size"] = self.size + if self.annotations: + ann = self.annotations.to_wire() + if ann: + out["annotations"] = ann + return self._common_wire(out) + + +@dataclass +class ResourceTemplateSpec(_BaseSpec): + """A parameterised resource (RFC 6570 level-1 ``{var}`` templates). + + ``handler(ctx, **variables)`` receives the URI variables. ``completions`` + maps a variable name to ``fn(ctx, partial) -> list[str]`` for + ``completion/complete``. + """ + + uri_template: str = "" + mime_type: str = "application/json" + annotations: ContentAnnotations | None = None + completions: dict[str, Callable[..., Any]] = field(default_factory=dict) + + def __post_init__(self) -> None: + super().__post_init__() + self._regex = _template_regex(self.uri_template) + + @property + def variables(self) -> list[str]: + return re.findall(r"{\+?([A-Za-z0-9_]+)\*?}", self.uri_template) + + def match(self, uri: str) -> dict[str, str] | None: + m = self._regex.match(uri) + if not m: + return None + from urllib.parse import unquote + + return {k: unquote(v) for k, v in m.groupdict().items()} + + def expand(self, **values: str) -> str: + from urllib.parse import quote + + out = self.uri_template + for k, v in values.items(): + multi = quote(str(v), safe="/:@!$&'()*+,;=") + out = ( + out.replace("{+" + k + "}", multi) + .replace("{" + k + "*}", multi) + .replace("{" + k + "}", quote(str(v), safe="")) + ) + return out + + def canonical(self, values: dict[str, str]) -> str: + """The URI with decoded ``values`` filled in verbatim (no + percent-encoding): ``cloudg://graph/%64%33`` and ``cloudg://graph/d3`` + share one canonical form, which is what access policies match.""" + out = self.uri_template + for k, v in values.items(): + for form in ("{+" + k + "}", "{" + k + "*}", "{" + k + "}"): + out = out.replace(form, str(v)) + return out + + def to_wire(self) -> dict[str, Any]: + out: dict[str, Any] = { + "uriTemplate": self.uri_template, + "name": self.name, + "mimeType": self.mime_type, + } + if self.description: + out["description"] = self.description + if self.annotations: + ann = self.annotations.to_wire() + if ann: + out["annotations"] = ann + return self._common_wire(out) + + +def _template_regex(template: str) -> re.Pattern[str]: + parts = re.split(r"({\+?[A-Za-z0-9_]+\*?})", template) + rx = "" + for part in parts: + m = re.fullmatch(r"{(\+?)([A-Za-z0-9_]+)(\*?)}", part) + if m: + # {+var} (RFC 6570 reserved expansion) and the legacy {var*} may + # span path segments; {var} stops at "/" and "?" + multi = bool(m.group(1) or m.group(3)) + rx += f"(?P<{m.group(2)}>.+)" if multi else f"(?P<{m.group(2)}>[^/?]+)" + else: + rx += re.escape(part) + return re.compile("^" + rx + "$") + + +@dataclass +class PromptArgument: + name: str + description: str = "" + required: bool = False + completion: Callable[..., Any] | None = None + title: str | None = None + + def to_wire(self) -> dict[str, Any]: + out: dict[str, Any] = {"name": self.name, "required": self.required} + if self.title: + out["title"] = self.title + if self.description: + out["description"] = self.description + return out + + +@dataclass +class PromptSpec(_BaseSpec): + """A prompt template. ``handler(ctx, **arguments)`` returns a + :class:`PromptResult`, a list of PromptMessage, or a plain string (sent + as one user message).""" + + arguments: list[PromptArgument] = field(default_factory=list) + # Prompts usually embed workspace summaries + sensitivity: Sensitivity = Sensitivity.INTERNAL + + def to_wire(self, name: str | None = None) -> dict[str, Any]: + exposed = name or self.name + if not _NAME_RE.match(exposed): + raise ValueError(f"Invalid exposed prompt name {exposed!r}") + out: dict[str, Any] = { + "name": exposed, + "description": self.description, + "arguments": [a.to_wire() for a in self.arguments], + } + return self._common_wire(out) + + +# --------------------------------------------------------------------------- +# Registry +# --------------------------------------------------------------------------- + + +#: Keyword options of :meth:`Registry.resource` / :meth:`Registry.resource_template` +#: and their defaults. +_RESOURCE_OPTIONS: dict[str, Any] = { + "name": None, + "title": None, + "description": None, + "mime_type": "application/json", + "category": "general", + "tags": (), + "sensitivity": Sensitivity.CONFIDENTIAL, + "annotations": None, + "capabilities": (Capability.READ_STATE,), + "icons": None, +} + + +class Registry: + """Collection of tool / resource / prompt specs. + + Catalog modules expose ``register(registry)`` functions; users can build + their own registry, merge registries, or register extra primitives with + the decorators:: + + reg = Registry() + + @reg.tool(category="custom", read_only=True) + def my_tool(ctx, asset_id: str) -> dict: ... + """ + + def __init__(self) -> None: + self.tools: dict[str, ToolSpec] = {} + self.resources: dict[str, ResourceSpec] = {} + self.templates: dict[str, ResourceTemplateSpec] = {} + self.prompts: dict[str, PromptSpec] = {} + + # -- adding --------------------------------------------------------- + + def add(self, spec: _BaseSpec, *, replace: bool = False) -> _BaseSpec: + if isinstance(spec, ToolSpec): + table: dict[str, Any] = self.tools + key = spec.name + elif isinstance(spec, ResourceSpec): + table, key = self.resources, spec.uri + elif isinstance(spec, ResourceTemplateSpec): + table, key = self.templates, spec.uri_template + elif isinstance(spec, PromptSpec): + table, key = self.prompts, spec.name + else: # pragma: no cover (defensive) + raise TypeError(f"Unsupported spec {type(spec).__name__}") + if key in table and not replace: + raise ValueError(f"Duplicate {type(spec).__name__} {key!r}") + table[key] = spec + return spec + + def merge(self, other: "Registry", *, replace: bool = False) -> "Registry": + for spec in [ + *other.tools.values(), + *other.resources.values(), + *other.templates.values(), + *other.prompts.values(), + ]: + self.add(spec, replace=replace) + return self + + def tool( + self, + name: str | None = None, + *, + title: str | None = None, + description: str | None = None, + category: str = "general", + tags: Iterable[str] = (), + sensitivity: Sensitivity = Sensitivity.CONFIDENTIAL, + capabilities: Iterable[Capability] = (Capability.READ_STATE,), + read_only: bool | None = None, + destructive: bool | None = None, + idempotent: bool | None = None, + open_world: bool | None = None, + output_schema: dict[str, Any] | None = None, + transform_hints: dict[str, Any] | None = None, + timeout_seconds: float | None = None, + icons: list[Icon] | None = None, + ) -> Callable[[Handler], Handler]: + def deco(fn: Handler) -> Handler: + self.add( + ToolSpec( + name=name or fn.__name__, + handler=fn, + title=title, + description=description or "", + category=category, + tags=set(tags), + sensitivity=sensitivity, + capabilities=set(capabilities), + annotations=ToolAnnotations( + title=title, + read_only=read_only, + destructive=destructive, + idempotent=idempotent, + open_world=open_world, + ), + output_schema=output_schema, + transform_hints=dict(transform_hints or {}), + timeout_seconds=timeout_seconds, + icons=icons, + ) + ) + return fn + + return deco + + def resource(self, uri: str, **options: Any) -> Callable[[Handler], Handler]: + """Decorator registering a static :class:`ResourceSpec` at ``uri``. + + Keyword options (all optional): ``name``, ``title``, ``description``, + ``mime_type`` (default ``application/json``), ``category``, ``tags``, + ``sensitivity`` (default confidential), ``annotations``, + ``capabilities`` and ``icons``. + """ + return self._resource_decorator(ResourceSpec, {"uri": uri}, options) + + def resource_template( + self, + uri_template: str, + *, + completions: dict[str, Callable[..., Any]] | None = None, + **options: Any, + ) -> Callable[[Handler], Handler]: + """Decorator registering a :class:`ResourceTemplateSpec`. + + Takes the same keyword options as :meth:`resource`, plus + ``completions`` (variable name to ``fn(ctx, partial)``). + """ + fixed = {"uri_template": uri_template, "completions": dict(completions or {})} + return self._resource_decorator(ResourceTemplateSpec, fixed, options) + + def _resource_decorator( + self, cls: type, fixed: dict[str, Any], options: dict[str, Any] + ) -> Callable[[Handler], Handler]: + unknown = sorted(set(options) - set(_RESOURCE_OPTIONS)) + if unknown: + raise TypeError(f"Unexpected resource option(s): {', '.join(unknown)}") + opts = {**_RESOURCE_OPTIONS, **options} + + def deco(fn: Handler) -> Handler: + self.add( + cls( + name=opts["name"] or fn.__name__, + handler=fn, + title=opts["title"], + description=opts["description"] or "", + mime_type=opts["mime_type"], + category=opts["category"], + tags=set(opts["tags"]), + sensitivity=opts["sensitivity"], + annotations=opts["annotations"], + capabilities=set(opts["capabilities"]), + icons=opts["icons"], + **fixed, + ) + ) + return fn + + return deco + + def prompt( + self, + name: str | None = None, + *, + title: str | None = None, + description: str | None = None, + arguments: list[PromptArgument] | None = None, + category: str = "general", + tags: Iterable[str] = (), + sensitivity: Sensitivity = Sensitivity.INTERNAL, + icons: list[Icon] | None = None, + ) -> Callable[[Handler], Handler]: + def deco(fn: Handler) -> Handler: + self.add( + PromptSpec( + name=name or fn.__name__, + handler=fn, + title=title, + description=description or "", + arguments=list(arguments or []), + category=category, + tags=set(tags), + sensitivity=sensitivity, + icons=icons, + ) + ) + return fn + + return deco + + # -- lookup --------------------------------------------------------- + + def find_template(self, uri: str) -> tuple[ResourceTemplateSpec, dict[str, str]] | None: + # Longest template first, so the most specific pattern wins + # (cloudg://assets/{+ref}/neighbors before cloudg://assets/{+ref}) + for spec in sorted(self.templates.values(), key=lambda s: -len(s.uri_template)): + values = spec.match(uri) + if values is not None: + return spec, values + return None + + def __len__(self) -> int: + return len(self.tools) + len(self.resources) + len(self.templates) + len(self.prompts) + + +def render_json(value: Any) -> str: + """How the layer renders structured data as text content and JSON + resources. Size budgets (the projection transform) measure with the + same function, so their numbers match what the client receives.""" + import json + + return json.dumps(value, indent=2, ensure_ascii=False, default=str) + + +async def maybe_await(value: Any) -> Any: + """Await ``value`` if it is awaitable.""" + if inspect.isawaitable(value): + return await value + return value diff --git a/cloudg/mcp/layer.py b/cloudg/mcp/layer.py new file mode 100644 index 0000000..860d0fe --- /dev/null +++ b/cloudg/mcp/layer.py @@ -0,0 +1,653 @@ +"""The cloudg MCP layer: one object that owns the registry, the workspace, +the policy and the middleware chain, and dispatches MCP requests. + +Adapters never call handlers directly; they go through +:meth:`CloudGMCPLayer.call_tool`, :meth:`read_resource`, +:meth:`get_prompt` and :meth:`complete`. Whichever server a request arrives +through, it gets the same validation, access control, transforms +(redaction / pseudonymisation / projection / annotation / substitution), +timeouts and audit trail. Output shaping lives in +:mod:`cloudg.mcp.layer_output`, construction options in +:mod:`cloudg.mcp.layer_options`. +""" + +from __future__ import annotations + +import asyncio +import inspect +import logging +import re +import time +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Awaitable, Callable + +from cloudg.mcp.context import Principal, ToolContext +from cloudg.mcp.core import ( + META_PREFIX, + BlobResourceContents, + InvalidArgumentsError, + MCPLayerError, + NotFoundError, + PromptResult, + PromptSpec, + ResourceSpec, + ResourceTemplateSpec, + TextResourceContents, + ToolResult, + ToolSpec, + maybe_await, + render_json, +) +from cloudg.mcp.layer_options import LayerOptions, filter_registry +from cloudg.mcp.layer_output import ( + OutputScope, + error_result, + jsonable, + restored_pairs, + shape_prompt, + shape_resource, + shape_tool_result, + transform_error, +) +from cloudg.mcp.schema import validate_arguments +from cloudg.mcp.transforms.base import TransformContext + +if TYPE_CHECKING: # pragma: no cover + from cloudg.config import CloudGConfig + from cloudg.mcp.policy import Policy + +__all__ = ["CallInfo", "CloudGMCPLayer", "LayerOptions", "Middleware", "Next"] + +logger = logging.getLogger("cloudg.mcp") + +DEFAULT_INSTRUCTIONS = """\ +cloudg exposes a multi-cloud (AWS / Azure / GCP) infrastructure inventory, +its relationship graph, security findings, compliance mappings and an RDF +ontology. Start with `workspace_status` (or read cloudg://workspace) to see +which datasets are loaded; load one with `load_dataset` or collect live data +with `map_inventory`. Identifiers in results may be pseudonymised by policy: +pass them back to tools unchanged and they resolve to the real resource. +Prefer the narrow query tools (find_assets, get_asset, neighbors, paths) +over dumping whole datasets; results are paginated.""" + + +@dataclass +class CallInfo: + """What a middleware sees about the request it wraps.""" + + kind: str # tool | resource | prompt | completion + name: str # tool/prompt name, resource URI, or template of a completion + spec: Any + arguments: dict[str, Any] + principal: Principal + context: ToolContext + started: float = field(default_factory=time.monotonic) + #: The call's :class:`~cloudg.mcp.layer_output.OutputScope` (set by the layer). + output: Any = field(default=None, repr=False) + + +_EMPTY_COMPLETION: dict[str, Any] = {"values": [], "total": 0, "hasMore": False} +_PREFIX_RE = re.compile(r"^[A-Za-z0-9_.\-]{0,64}$") +#: ``ToolContext.meta`` key holding the canonical URI of a resource read. +URI_META_KEY = f"{META_PREFIX}uri" + +Next = Callable[[CallInfo], Awaitable[Any]] +# ``async def middleware(info: CallInfo, call_next: Next) -> Any``; for tools +# the value is a ToolResult, for resources a list of contents, for prompts a +# PromptResult. Raise an MCPLayerError to reject the call. +Middleware = Callable[[CallInfo, Next], Awaitable[Any]] + +dumps = render_json +_URI_KEYWORD: dict[Any, bool] = {} + + +def _policy_call(fn: Callable[..., Any], *args: Any, uri: str | None = None) -> Any: + """Call a policy method, passing the concrete resource ``uri`` when the + policy accepts it (older policies only match the spec's own URI).""" + if uri is not None: + key = getattr(fn, "__func__", fn) + accepts = _URI_KEYWORD.get(key) + if accepts is None: + try: + params = inspect.signature(fn).parameters.values() + except (TypeError, ValueError): + params = () # type: ignore[assignment] + accepts = any(p.name == "uri" or p.kind is p.VAR_KEYWORD for p in params) + _URI_KEYWORD[key] = accepts + if accepts: + return fn(*args, uri=uri) + return fn(*args) + + +class CloudGMCPLayer: + """Plug-and-play MCP layer over cloudg. + + Args: + config: cloudg configuration (providers, credentials...). Defaults + to ``CloudGConfig()``. + options: All other settings as one :class:`LayerOptions`. + **kwargs: The same settings as keywords (each field of + :class:`LayerOptions`): ``workspace``, ``policy`` (a + :class:`~cloudg.mcp.policy.Policy`, profile name, policy file + path or dict), ``registry``, ``include_categories`` / + ``exclude_categories``, ``include_tools`` / ``exclude_tools``, + ``prefix`` (prepended to every tool and prompt name), + ``middleware`` (outermost first), ``default_timeout`` (seconds + before a tool call is cancelled), ``max_output_chars``, ``name``, + ``version`` and ``instructions``. Unknown keywords raise + :class:`TypeError`. + """ + + def __init__( + self, + config: "CloudGConfig | None" = None, + *, + options: LayerOptions | None = None, + **kwargs: Any, + ) -> None: + from cloudg import __version__ + from cloudg.mcp.policy import Policy as PolicyClass + + opts = LayerOptions.build(options, kwargs) + workspace = opts.workspace + if workspace is None: + from cloudg.config import CloudGConfig + from cloudg.mcp.state import Workspace + + workspace = Workspace(config or CloudGConfig()) + self.workspace = workspace + policy = opts.policy + self.policy: Policy = ( + policy if isinstance(policy, PolicyClass) else PolicyClass.load(policy) + ) + registry = opts.registry + if registry is None: + from cloudg.mcp.catalog import default_registry + + registry = default_registry() + self.registry = filter_registry(registry, opts) + self.prefix = "" + self.set_prefix(opts.prefix) + self.middleware: list[Middleware] = list(opts.middleware) + self.default_timeout = opts.default_timeout + self.max_output_chars = opts.max_output_chars + self.name = opts.name + self.version = opts.version or __version__ + self.instructions = ( + opts.instructions if opts.instructions is not None else DEFAULT_INSTRUCTIONS + ) + # Notified with (kind, uri) when a resource changes (dataset loaded, + # workspace mutated); adapters forward these as + # notifications/resources/updated and list_changed. + self._change_listeners: list[Callable[[str, str | None], Any]] = [] + self._loop: asyncio.AbstractEventLoop | None = None + if callable(getattr(workspace, "on_change", None)): + workspace.on_change(self.notify_change) + + # ------------------------------------------------------------------ + # Construction helpers + # ------------------------------------------------------------------ + + def set_prefix(self, prefix: str) -> None: + """Change the name prefix of every exposed tool and prompt. The + prefix belongs to the layer, so one layer has one prefix: mount two + layers to expose the catalog under two prefixes.""" + if not isinstance(prefix, str) or not _PREFIX_RE.match(prefix): + raise ValueError(f"Invalid prefix {prefix!r}: use up to 64 of A-Z a-z 0-9 _ . -") + self.prefix = prefix + + def use(self, middleware: Middleware) -> "CloudGMCPLayer": + """Append a middleware (innermost so far).""" + self.middleware.append(middleware) + return self + + def on_change(self, listener: Callable[[str, str | None], Any]) -> None: + """Register ``listener(kind, uri)``; kind is ``"resources"`` + (list changed), ``"resource"`` (one URI updated), ``"tools"`` or + ``"prompts"``.""" + self._change_listeners.append(listener) + + def notify_change(self, kind: str, uri: str | None = None) -> None: + """Fan a change out to the listeners, on the server's event loop: + sync handlers run in worker threads and adapters' notification + code is not thread-safe.""" + loop = self._loop + if loop is not None and loop.is_running(): + try: + running = asyncio.get_running_loop() + except RuntimeError: + running = None + if running is not loop: + loop.call_soon_threadsafe(self._dispatch_change, kind, uri) + return + self._dispatch_change(kind, uri) + + def _dispatch_change(self, kind: str, uri: str | None) -> None: + for listener in list(self._change_listeners): + try: + listener(kind, uri) + except Exception: + logger.debug("change listener failed", exc_info=True) + + # ------------------------------------------------------------------ + # Names + # ------------------------------------------------------------------ + + def exposed_name(self, name: str) -> str: + return f"{self.prefix}{name}" + + def internal_name(self, exposed: str) -> str: + """Registry name for an exposed tool or prompt name. With a prefix + set, only prefixed names resolve: a host server mounting the layer + may own a tool with the bare name.""" + if not self.prefix: + return exposed + if exposed.startswith(self.prefix): + return exposed[len(self.prefix) :] + return "\0unprefixed:" + exposed # never a registry key + + def context(self, principal: Principal | None = None, **kwargs: Any) -> ToolContext: + return ToolContext(layer=self, principal=principal or Principal.local(), **kwargs) + + def _remember_loop(self) -> None: + try: + self._loop = asyncio.get_running_loop() + except RuntimeError: + logger.debug("no running event loop to remember") + + # ------------------------------------------------------------------ + # Listing + # ------------------------------------------------------------------ + + def list_tools(self, principal: Principal | None = None) -> list[ToolSpec]: + p = principal or Principal.local() + return [s for s in self.registry.tools.values() if self.policy.is_allowed(s, p)] + + def list_resources(self, principal: Principal | None = None) -> list[ResourceSpec]: + p = principal or Principal.local() + return [s for s in self.registry.resources.values() if self.policy.is_allowed(s, p)] + + def list_resource_templates( + self, principal: Principal | None = None + ) -> list[ResourceTemplateSpec]: + p = principal or Principal.local() + return [s for s in self.registry.templates.values() if self.policy.is_allowed(s, p)] + + def list_prompts(self, principal: Principal | None = None) -> list[PromptSpec]: + p = principal or Principal.local() + return [s for s in self.registry.prompts.values() if self.policy.is_allowed(s, p)] + + def tools_wire(self, principal: Principal | None = None) -> list[dict[str, Any]]: + """``tools/list`` entries for ``principal``. + + ``outputSchema`` is withheld when the caller's output pipeline is + non-empty: masking or projection changes the shape of + ``structuredContent``, and clients reject results that do not + validate against an advertised schema. + """ + p = principal or Principal.local() + return [ + s.to_wire( + self.exposed_name(s.name), + include_output_schema=not self.policy.output_pipeline(s, p), + ) + for s in self.list_tools(p) + ] + + def resources_wire(self, principal: Principal | None = None) -> list[dict[str, Any]]: + return [s.to_wire() for s in self.list_resources(principal)] + + def resource_templates_wire(self, principal: Principal | None = None) -> list[dict[str, Any]]: + return [s.to_wire() for s in self.list_resource_templates(principal)] + + def prompts_wire(self, principal: Principal | None = None) -> list[dict[str, Any]]: + return [s.to_wire(self.exposed_name(s.name)) for s in self.list_prompts(principal)] + + # ------------------------------------------------------------------ + # Middleware plumbing, transforms + # ------------------------------------------------------------------ + + async def _run_chain(self, info: CallInfo, terminal: Next) -> Any: + chain: Next = terminal + for mw in reversed(self.middleware): + chain = _bind(mw, chain) + return await chain(info) + + def _transform( + self, value: Any, spec: Any, principal: Principal, kind: str + ) -> tuple[Any, dict]: + """``(transformed, report)`` for ``value`` under the caller's output + pipeline.""" + return OutputScope(self.policy, spec, principal, kind).apply(value) + + def _begin( + self, kind: str, name: str, spec: Any, args: dict[str, Any], ctx: ToolContext + ) -> CallInfo: + """Set up one call: its context and its output scope (also used for + the progress and log notifications the handler sends).""" + scope = OutputScope(self.policy, spec, ctx.principal, kind) + ctx.kind, ctx.name = kind, name + ctx.transform_output = lambda value: scope(jsonable(value)) + return CallInfo(kind, name, spec, args, ctx.principal, ctx, output=scope) + + def _scope(self, info: CallInfo) -> OutputScope: + if info.output is None: # a middleware built its own CallInfo + info.output = OutputScope(self.policy, info.spec, info.principal, info.kind) + return info.output + + def _untransform(self, info: CallInfo, values: dict[str, Any]) -> dict[str, Any]: + """Run the input pipeline (secret guards, pseudonym reversal) and + remember the real values it restored for the output side.""" + pipeline = self.policy.input_pipeline(info.spec, info.principal) + if not pipeline: + return values + tctx = TransformContext( + principal=info.principal, + spec=info.spec, + kind=info.kind, + direction="input", + vault=self.policy.vault, + ) + out = pipeline.apply(values, tctx) + self._scope(info).add_restored(restored_pairs(values, out, tctx)) + return out + + async def _invoke( + self, fn: Callable[..., Any], ctx: ToolContext, kwargs: dict[str, Any] + ) -> Any: + """Run a handler: coroutine functions on the loop, plain functions + in a worker thread so blocking graph/ontology work does not stall + the server.""" + self._remember_loop() + ctx.loop = self._loop + if asyncio.iscoroutinefunction(fn): + return await fn(ctx, **kwargs) + result = await asyncio.to_thread(fn, ctx, **kwargs) + return await maybe_await(result) + + async def _invoke_guarded(self, info: CallInfo, kwargs: dict[str, Any]) -> Any: + """:meth:`_invoke` for resources and prompts: an unexpected handler + exception becomes an internal :class:`MCPLayerError` (transformed by + the caller) instead of leaking raw exception text through an + adapter.""" + try: + return await self._invoke(info.spec.handler, info.context, kwargs) + except MCPLayerError: + raise + except Exception as exc: + logger.exception("%s %s failed", info.kind.capitalize(), info.name) + raise MCPLayerError(f"{type(exc).__name__}: {exc}") from None + + def _record(self, hook: str, *args: Any) -> None: + """Call an optional policy audit hook (``record_hidden`` / + ``record_rejection``); failures are logged, never raised.""" + record = getattr(self.policy, hook, None) + if record is not None: + try: + record(*args) + except Exception: + logger.debug("%s failed", hook, exc_info=True) + + # ------------------------------------------------------------------ + # Tools + # ------------------------------------------------------------------ + + def get_tool(self, name: str, principal: Principal | None = None) -> ToolSpec: + principal = principal or Principal.local() + spec = self.registry.tools.get(self.internal_name(name)) + if spec is None or not self.policy.is_allowed(spec, principal): + self._record("record_hidden", "tool", name, principal) + raise NotFoundError(f"Unknown tool: {name}") + return spec + + async def call_tool( + self, + name: str, + arguments: dict[str, Any] | None = None, + *, + principal: Principal | None = None, + context: ToolContext | None = None, + ) -> ToolResult: + """Call a tool. Unknown tools raise :class:`NotFoundError`; every + other failure (bad arguments, access denied, handler error, + timeout) comes back as an ``isError`` result the model can read.""" + principal = principal or (context.principal if context else Principal.local()) + spec = self.get_tool(name, principal) + ctx = context or self.context(principal) + ctx.principal = principal + info = self._begin("tool", spec.name, spec, dict(arguments or {}), ctx) + try: + return await self._run_chain(info, self._tool_terminal) + except MCPLayerError as exc: + # Includes NotFoundError raised by a handler ("asset not found"): + # the model can recover from an isError result, not from a + # protocol error. + return error_result(self._scope(info), exc.message, exc.code, exc.data) + + async def _tool_terminal(self, info: CallInfo) -> ToolResult: + spec: ToolSpec = info.spec + scope = self._scope(info) + try: + self.policy.check_call(spec, info.principal, info.arguments) + try: + arguments = self._untransform(info, info.arguments) + except MCPLayerError as exc: + # check_call already audited this call as allowed; the input + # pipeline (e.g. a secret in the arguments) refused it + self._record("record_rejection", spec, info.principal, info.arguments, exc) + raise + kwargs = validate_arguments(spec.handler, arguments) + timeout = spec.timeout_seconds or self.default_timeout + coro = self._invoke(spec.handler, info.context, kwargs) + raw = await (asyncio.wait_for(coro, timeout) if timeout else coro) + except (MCPLayerError, InvalidArgumentsError): + raise + except asyncio.TimeoutError: + return error_result(scope, f"Tool {spec.name} timed out", "timeout") + except Exception as exc: + logger.exception("Tool %s failed", spec.name) + # Exception text often quotes ARNs / account ids: it goes + # through the same output pipeline as a normal result. + return error_result(scope, f"{type(exc).__name__}: {exc}", "handler_error") + result, report = shape_tool_result(scope, raw, info.context.links, self.max_output_chars) + if report: + result.meta[f"{META_PREFIX}transforms"] = report + result.meta.setdefault( + f"{META_PREFIX}duration_ms", int((time.monotonic() - info.started) * 1000) + ) + return result + + # ------------------------------------------------------------------ + # Resources + # ------------------------------------------------------------------ + + def _resolve_resource(self, uri: str, principal: Principal) -> tuple[Any, dict[str, str], str]: + """``(spec, template values, canonical URI)`` for ``uri``. + + Template values arrive percent-decoded, so ``cloudg://graph/%64%33`` + names the same resource as ``cloudg://graph/d3``: the canonical URI + (values filled in verbatim) resolves to the static resource when one + exists, and is what the policy matches allow / deny patterns against. + """ + spec: Any = self.registry.resources.get(uri) + values: dict[str, str] = {} + canonical = uri + if spec is None: + found = self.registry.find_template(uri) + if found: + spec, values = found + canonical = spec.canonical(values) + static = self.registry.resources.get(canonical) + if static is not None: + spec, values = static, {} + if spec is None or not _policy_call(self.policy.is_allowed, spec, principal, uri=canonical): + self._record("record_hidden", "resource", uri, principal) + raise NotFoundError(f"Unknown resource: {uri}") + return spec, values, canonical + + async def read_resource( + self, + uri: str, + *, + principal: Principal | None = None, + context: ToolContext | None = None, + ) -> list[TextResourceContents | BlobResourceContents]: + principal = principal or (context.principal if context else Principal.local()) + spec, values, canonical = self._resolve_resource(uri, principal) + served = canonical if isinstance(spec, ResourceSpec) else uri + ctx = context or self.context(principal) + ctx.principal = principal + ctx.meta[URI_META_KEY] = canonical + info = self._begin("resource", served, spec, values, ctx) + try: + return await self._run_chain(info, self._resource_terminal) + except MCPLayerError as exc: + raise transform_error(self._scope(info), exc) from None + + async def _resource_terminal(self, info: CallInfo) -> list[Any]: + spec = info.spec + canonical = info.context.meta.get(URI_META_KEY, info.name) + _policy_call(self.policy.check_call, spec, info.principal, info.arguments, uri=canonical) + values = self._untransform(info, info.arguments) + raw = await self._invoke_guarded(info, values) + return shape_resource(self._scope(info), raw, info.name, spec.mime_type) + + # ------------------------------------------------------------------ + # Prompts + # ------------------------------------------------------------------ + + async def get_prompt( + self, + name: str, + arguments: dict[str, Any] | None = None, + *, + principal: Principal | None = None, + context: ToolContext | None = None, + ) -> PromptResult: + principal = principal or (context.principal if context else Principal.local()) + spec = self.registry.prompts.get(self.internal_name(name)) + if spec is None or not self.policy.is_allowed(spec, principal): + self._record("record_hidden", "prompt", name, principal) + raise NotFoundError(f"Unknown prompt: {name}") + args = dict(arguments or {}) + missing = [a.name for a in spec.arguments if a.required and not args.get(a.name)] + if missing: + raise InvalidArgumentsError(f"Missing required prompt arguments: {', '.join(missing)}") + ctx = context or self.context(principal) + ctx.principal = principal + info = self._begin("prompt", spec.name, spec, args, ctx) + try: + return await self._run_chain(info, self._prompt_terminal) + except MCPLayerError as exc: + raise transform_error(self._scope(info), exc) from None + + async def _prompt_terminal(self, info: CallInfo) -> PromptResult: + spec: PromptSpec = info.spec + self.policy.check_call(spec, info.principal, info.arguments) + args = self._untransform(info, info.arguments) + known = {a.name for a in spec.arguments} + raw = await self._invoke_guarded(info, {k: v for k, v in args.items() if k in known}) + return shape_prompt(self._scope(info), raw, spec.description) + + # ------------------------------------------------------------------ + # Completions + # ------------------------------------------------------------------ + + def _completion_target(self, ref: dict[str, Any], arg_name: str) -> tuple[Any, Any]: + """``(spec, completion fn)`` named by a completion ``ref``.""" + spec: Any = None + fn: Callable[..., Any] | None = None + if ref.get("type") == "ref/prompt": + spec = self.registry.prompts.get(self.internal_name(ref.get("name", ""))) + if spec is not None: + fn = next((a.completion for a in spec.arguments if a.name == arg_name), None) + elif ref.get("type") == "ref/resource": + spec = self.registry.templates.get(ref.get("uri", "")) + if spec is not None: + fn = spec.completions.get(arg_name) + return spec, fn + + async def complete( + self, + ref: dict[str, Any], + argument: dict[str, Any], + *, + principal: Principal | None = None, + context_arguments: dict[str, Any] | None = None, + ) -> dict[str, Any]: + """``completion/complete``. ``ref`` is ``{"type": "ref/prompt", + "name": ...}`` or ``{"type": "ref/resource", "uri":