diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md index 94b57548f5..449c5fb538 100644 --- a/RELEASE-NOTES.md +++ b/RELEASE-NOTES.md @@ -19,6 +19,7 @@ ## Deprecations ## New additions +* `snow connection test --enable-diag` now runs SnowCD-style per-endpoint connectivity checks on the live connection. It merges `SYSTEM$ALLOWLIST()` with `SYSTEM$ALLOWLIST_PRIVATELINK()` (deduped by `(type, host, port)`; the privatelink call is best-effort and skipped on deployments where it is unavailable), probes each resolvable endpoint (DNS, TCP, plus a TLS handshake on port 443) with measured latency, and prints a streaming `Checking : ✅` line per endpoint, followed by a per-endpoint table (`url | type | status | latency_ms | issuer | cert_expires`), an effective-network-policy summary (account/user policy, current IP, allowed/blocked IP and rule lists), and a `Results: X Healthy, Y Unhealthy out of Z. N skipped (non-resolvable patterns)` line. Wildcard hostnames are counted as `Skipped`. The TLS probe respects the `REQUESTS_CA_BUNDLE` and `SSL_CERT_FILE` env vars (matching `snowflake-connector-python`) for custom CA bundles. JSON output (`--format json`) carries the same data under a `Diagnostic` key. The legacy `SnowflakeConnectionTestReport.txt` written by `snowflake-connector-python` is still produced. Replaces the deprecated SnowCD tool. * Added a `--protocol` option (plus matching `SNOWFLAKE_PROTOCOL` env var and `protocol` config key) to `snow connection add` and the global connection overrides. This allows selecting `http` or `https` as the connection protocol without editing `config.toml`, which is primarily useful for local development against `http` deployments. * Added `snow helpers generate-project-schema` command to emit a JSON Schema for the `snowflake.yml` project definition file. The output follows [JSON Schema Draft 2020-12](https://json-schema.org/draft/2020-12/schema) and can be plugged into supported editors (e.g. VS Code via the YAML extension) or CI pipelines to get completion and to catch structural mistakes (unknown keys, wrong types, missing required fields) before a deploy. The schema is generated from the CLI's own models, so some cross-field and semantic checks are still only applied at load/deploy time. Use `--definition-version` to select the project definition version (`1`, `1.1`, or `2`; defaults to `2`) and `--output-file`/`-o` to write the schema to a file. diff --git a/src/snowflake/cli/_plugins/connection/commands.py b/src/snowflake/cli/_plugins/connection/commands.py index 0218ee9af3..71b645efa2 100644 --- a/src/snowflake/cli/_plugins/connection/commands.py +++ b/src/snowflake/cli/_plugins/connection/commands.py @@ -16,6 +16,7 @@ import logging import os.path +from dataclasses import asdict from pathlib import Path from typing import Optional @@ -28,6 +29,13 @@ ) from click.core import ParameterSource # type: ignore from snowflake import connector +from snowflake.cli._plugins.connection.diagnostic import ( + DiagnosticReport, + NetworkPolicySnapshot, + collect_network_policy, + run_diagnostic, + status_line, +) from snowflake.cli._plugins.connection.util import ( strip_if_value_present, ) @@ -67,12 +75,15 @@ from snowflake.cli.api.config_ng.masking import mask_sensitive_value from snowflake.cli.api.console import cli_console from snowflake.cli.api.constants import DEFAULT_SIZE_LIMIT_MB, ObjectType +from snowflake.cli.api.output.formats import OutputFormat from snowflake.cli.api.output.types import ( CollectionResult, CommandResult, MessageResult, + MultipleResults, ObjectResult, ) +from snowflake.cli.api.sanitizers import sanitize_for_terminal from snowflake.cli.api.secure_path import SecurePath from snowflake.connector import ProgrammingError from snowflake.connector.constants import CONNECTIONS_FILE @@ -380,6 +391,15 @@ def test( # Test connection cli_context = get_cli_context() + conn_ctx = cli_context.connection_context + + # The connector's `enable_diag=True` runs ConnectionDiagnostic during + # connect() — a silent multi-second probe of every endpoint. Surface a + # heads-up so the wait isn't unexplained; our own per-endpoint stream + # comes later via `_connection_test_with_diag`. + if conn_ctx.enable_diag and cli_context.output_format == OutputFormat.TABLE: + cli_console.message("Running connector connectivity probe...") + conn = cli_context.connection # Test session attributes @@ -401,7 +421,6 @@ def test( except ProgrammingError as err: raise ClickException(str(err)) - conn_ctx = cli_context.connection_context result = { "Connection name": conn_ctx.connection_name, "Status": "OK", @@ -417,10 +436,146 @@ def test( result["Diag Report Location"] = os.path.join( conn_ctx.diag_log_path, "SnowflakeConnectionTestReport.txt" ) + return _connection_test_with_diag(conn, conn_ctx, result) return ObjectResult(result) +def _connection_test_with_diag( + conn, + conn_ctx, + connection_summary: dict, +) -> CommandResult: + """Run the SnowCD-style per-endpoint diagnostic and assemble the result. + + For TABLE output: streams `Checking : ` lines as the + run progresses, then appends a per-endpoint table, a network-policy block, + and a `Results: ...` summary line. + + For JSON / CSV output: returns the structured diagnostic nested under a + `Diagnostic` key so consumers can read everything off one object. + """ + is_table_output = get_cli_context().output_format == OutputFormat.TABLE + + if is_table_output: + cli_console.message("Fetching allowlist from Snowflake...") + report = run_diagnostic( + conn=conn, + allowlist_path=conn_ctx.diag_allowlist_path, + on_check=lambda c: ( + cli_console.message(status_line(c)) if is_table_output else None + ), + ) + + if is_table_output: + cli_console.message("Inspecting network policies...") + policy = collect_network_policy(conn, user=conn.user) + + if not is_table_output: + connection_summary["Diagnostic"] = { + "checks": [asdict(c) for c in report.checks], + "healthy": report.healthy, + "unhealthy": report.unhealthy, + "skipped": report.skipped, + "tested": report.tested, + "network_policy": asdict(policy), + } + return ObjectResult(connection_summary) + + return _diagnostic_table_results(connection_summary, report, policy) + + +def _safe(value: Optional[str], fallback: str = "") -> str: + """Return `value` with ANSI escapes stripped, or `fallback` when missing.""" + if value is None or value == "": + return fallback + return sanitize_for_terminal(value) or fallback + + +def _diagnostic_table_results( + connection_summary: dict, + report: DiagnosticReport, + policy: NetworkPolicySnapshot, +) -> CommandResult: + """Assemble the TABLE-format output: summary, per-endpoint rows, policy block, totals. + + Every server-derived string (host, IP, policy/rule names, allow/block lists, + cert metadata) is run through `sanitize_for_terminal` before display so a + crafted response from `SYSTEM$ALLOWLIST()` or `DESC NETWORK POLICY` cannot + inject ANSI escape sequences into the user's terminal. + """ + results = MultipleResults() + results.add(ObjectResult(connection_summary)) + + tested_rows = [ + { + "url": _safe(c.host), + "type": _safe(c.type), + "status": c.status, + "latency_ms": c.latency_ms if c.latency_ms is not None else "", + "issuer": _safe(c.cert_issuer), + "cert_expires": _safe(c.cert_expires), + } + for c in report.checks + if c.status != "Skipped" + ] + if tested_rows: + results.add(CollectionResult(tested_rows)) + + if policy.has_policy(): + policy_summary = { + "Effective network policy": _safe(policy.effective_policy), + "Source": "user" if policy.user_policy else "account", + "Account-level": _safe(policy.account_policy, "(none)"), + "User-level": _safe(policy.user_policy, "(none)"), + "Current IP": _safe(policy.current_ip, "(unknown)"), + "Allowed IPs": ", ".join( + sanitize_for_terminal(v) or "" for v in policy.allowed_ip_list + ) + or "(none)", + "Blocked IPs": ", ".join( + sanitize_for_terminal(v) or "" for v in policy.blocked_ip_list + ) + or "(none)", + "Allowed network rules": ", ".join( + sanitize_for_terminal(v) or "" for v in policy.allowed_network_rule_list + ) + or "(none)", + "Blocked network rules": ", ".join( + sanitize_for_terminal(v) or "" for v in policy.blocked_network_rule_list + ) + or "(none)", + } + if policy.error: + # `error` interpolates the server-supplied policy name. + policy_summary["Note"] = _safe(policy.error) + results.add(ObjectResult(policy_summary)) + if policy.rules: + results.add( + CollectionResult( + [ + { + "rule": _safe(r.name), + "mode": _safe(r.mode), + "type": _safe(r.type), + "values": ", ".join( + sanitize_for_terminal(v) or "" for v in r.values + ), + } + for r in policy.rules + ] + ) + ) + elif policy.current_ip: + results.add( + MessageResult( + f"No network policy in effect (current IP: {_safe(policy.current_ip)})." + ) + ) + results.add(MessageResult(report.summary_line())) + return results + + @app.command(requires_connection=False) def set_default( name: str = typer.Argument( diff --git a/src/snowflake/cli/_plugins/connection/diagnostic.py b/src/snowflake/cli/_plugins/connection/diagnostic.py new file mode 100644 index 0000000000..2d965d2790 --- /dev/null +++ b/src/snowflake/cli/_plugins/connection/diagnostic.py @@ -0,0 +1,513 @@ +# Copyright (c) 2024 Snowflake Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Connectivity diagnostics for `snow connection test --enable-diag`. + +Replicates the per-endpoint checks SnowCD used to provide. Endpoints come from +`SYSTEM$ALLOWLIST()` (or a JSON file via `--diag-allowlist-path`); each one is +probed with TCP connect, plus a TLS handshake on port 443 to capture the cert +issuer and expiry. +""" + +from __future__ import annotations + +import json +import logging +import os +import socket +import ssl +import time +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Callable, Iterable, Literal, Optional + +from snowflake.cli._plugins.connection.util import ALLOWLIST_QUERY +from snowflake.cli.api.exceptions import CliError +from snowflake.cli.api.identifiers import FQN +from snowflake.cli.api.project.util import identifier_to_show_like_pattern +from snowflake.cli.api.sanitizers import sanitize_for_terminal +from snowflake.cli.api.secure_path import SecurePath +from snowflake.connector import SnowflakeConnection +from snowflake.connector.cursor import DictCursor + +log = logging.getLogger(__name__) + +DEFAULT_TIMEOUT_SECONDS: float = 5.0 +TLS_PORT: int = 443 +ALLOWLIST_PRIVATELINK_QUERY = "SELECT SYSTEM$ALLOWLIST_PRIVATELINK()" +ALLOWLIST_FILE_SIZE_LIMIT_MB: float = 1.0 +NETWORK_POLICY_PARAM_NAME = "NETWORK_POLICY" + +Status = Literal["Healthy", "Unhealthy", "Skipped"] + + +@dataclass +class EndpointCheck: + host: str + port: int + type: str # noqa: A003 — public field name; matches `SYSTEM$ALLOWLIST()` JSON + status: Status + error: Optional[str] = None + cert_issuer: Optional[str] = None + cert_expires: Optional[str] = None + latency_ms: Optional[float] = None + + +@dataclass +class DiagnosticReport: + checks: list[EndpointCheck] = field(default_factory=list) + + def _count(self, status: Status) -> int: + return sum(1 for c in self.checks if c.status == status) + + @property + def healthy(self) -> int: + return self._count("Healthy") + + @property + def unhealthy(self) -> int: + return self._count("Unhealthy") + + @property + def skipped(self) -> int: + return self._count("Skipped") + + @property + def tested(self) -> int: + return self.healthy + self.unhealthy + + def summary_line(self) -> str: + return ( + f"Results: {self.healthy} Healthy, {self.unhealthy} Unhealthy " + f"out of {self.tested} endpoints. " + f"{self.skipped} skipped (non-resolvable patterns)." + ) + + +# --------------------------------------------------------------------------- # +# Allowlist loading +# --------------------------------------------------------------------------- # + + +def _query_json_list(conn: SnowflakeConnection, sql: str, key: str) -> list[dict]: + """Run an allowlist-returning system function; return its JSON list, or [] on any failure. + + Failures are silent because either query may be unavailable: low-priv roles + can't call `SYSTEM$ALLOWLIST()`, and `SYSTEM$ALLOWLIST_PRIVATELINK()` isn't + enabled in every deployment. The caller decides what to do with an empty + result (typically: fall back to the connection host). + """ + try: + *_, cursor = conn.execute_string(sql, cursor_class=DictCursor) + row = cursor.fetchone() + raw = row.get(key) if row else None + parsed = json.loads(raw) if raw else [] + except Exception: + log.debug("Allowlist query failed: %s", sql, exc_info=True) + return [] + return parsed if isinstance(parsed, list) else [] + + +def _dedupe_entries(entries: list[dict[str, Any]]) -> list[dict[str, Any]]: + """Drop duplicate `(type, host, port)` triples, preserving first-seen order.""" + seen: set[tuple[str, str, Any]] = set() + out: list[dict[str, Any]] = [] + for e in entries: + key = (str(e.get("type", "")), str(e.get("host", "")), e.get("port")) + if key not in seen: + seen.add(key) + out.append(e) + return out + + +def load_allowlist( + conn: SnowflakeConnection, allowlist_path: Optional[Path] +) -> list[dict[str, Any]]: + """Return the raw allowlist as a list of `{type, host, port}` dicts. + + If `allowlist_path` is given, parse it as JSON. Otherwise merge + `SYSTEM$ALLOWLIST()` and `SYSTEM$ALLOWLIST_PRIVATELINK()` from the open + connection, deduping `(type, host, port)`. Either query can fail silently + (low-priv role, function unavailable); the caller treats an empty result + as a signal to fall back to the connection host. + """ + if allowlist_path is not None: + try: + payload = json.loads( + SecurePath(allowlist_path).read_text( + file_size_limit_mb=ALLOWLIST_FILE_SIZE_LIMIT_MB + ) + ) + except (OSError, json.JSONDecodeError) as exc: + raise CliError( + f"Could not read allowlist file {allowlist_path}: {exc}. " + "Check the path and that the file is valid JSON." + ) + if not isinstance(payload, list): + raise CliError( + f"Allowlist file {allowlist_path} must contain a JSON array " + "of {type, host, port} entries." + ) + return payload + + public = _query_json_list(conn, ALLOWLIST_QUERY, "SYSTEM$ALLOWLIST()") + privatelink = _query_json_list( + conn, ALLOWLIST_PRIVATELINK_QUERY, "SYSTEM$ALLOWLIST_PRIVATELINK()" + ) + return _dedupe_entries([*public, *privatelink]) + + +# --------------------------------------------------------------------------- # +# Endpoint probing +# --------------------------------------------------------------------------- # + + +def is_resolvable(host: str) -> bool: + """False for hostnames the OS resolver can never satisfy (wildcards, empty).""" + return bool(host) and "*" not in host + + +def _resolve_cafile() -> Optional[str]: + """Honour `REQUESTS_CA_BUNDLE` then `SSL_CERT_FILE`, matching the connector. + + Returns `None` if neither points at a readable file, in which case + `ssl.create_default_context()` uses the system trust store. + """ + for var in ("REQUESTS_CA_BUNDLE", "SSL_CERT_FILE"): + path = os.environ.get(var) + if path and os.path.isfile(path): + return path + return None + + +def _issuer_org_name(cert: dict) -> Optional[str]: + """Extract organizationName from a peer cert's issuer RDN sequence.""" + for rdn in cert.get("issuer") or (): + for entry in rdn: + if ( + isinstance(entry, tuple) + and len(entry) == 2 + and entry[0] == "organizationName" + ): + return entry[1] + return None + + +def _probe_tls(sock: socket.socket, host: str) -> tuple[Optional[str], Optional[str]]: + """Wrap `sock` in TLS; return `(issuer_org, not_after)` from the peer cert.""" + context = ssl.create_default_context(cafile=_resolve_cafile()) + context.minimum_version = ssl.TLSVersion.TLSv1_2 + with context.wrap_socket(sock, server_hostname=host) as ssock: + cert: Any = ssock.getpeercert() or {} + not_after = cert.get("notAfter") + return _issuer_org_name(cert), not_after if isinstance(not_after, str) else None + + +def check_endpoint( + host: str, + port: int, + type_: str, + timeout: float = DEFAULT_TIMEOUT_SECONDS, +) -> EndpointCheck: + """Probe a single endpoint. Never raises; failures map to Unhealthy.""" + if not is_resolvable(host): + return EndpointCheck( + host, port, type_, "Skipped", error="non-resolvable pattern" + ) + + try: + start = time.perf_counter() + sock = socket.create_connection((host, port), timeout=timeout) + latency_ms = round((time.perf_counter() - start) * 1000, 1) + except OSError as exc: + return EndpointCheck(host, port, type_, "Unhealthy", error=str(exc)) + + try: + if port == TLS_PORT: + try: + issuer, expires = _probe_tls(sock, host) + except (ssl.SSLError, OSError) as exc: + return EndpointCheck( + host, + port, + type_, + "Unhealthy", + error=f"TLS handshake failed: {exc}", + ) + else: + issuer = expires = None + finally: + try: + sock.close() + except OSError: + pass + + return EndpointCheck( + host, + port, + type_, + "Healthy", + cert_issuer=issuer, + cert_expires=expires, + latency_ms=latency_ms, + ) + + +def _normalise_entries( + entries: Iterable[dict[str, Any]], +) -> list[tuple[str, int, str]]: + """Pull `(host, port, type)` triples from allowlist entries, dropping malformed rows.""" + out: list[tuple[str, int, str]] = [] + for entry in entries: + host = entry.get("host") + if not isinstance(host, str): + continue + try: + port = int(entry.get("port", TLS_PORT)) + except (TypeError, ValueError): + continue + out.append((host, port, str(entry.get("type", "UNKNOWN")))) + return out + + +def run_diagnostic( + conn: SnowflakeConnection, + allowlist_path: Optional[Path], + on_check: Callable[[EndpointCheck], None] = lambda _check: None, + timeout: float = DEFAULT_TIMEOUT_SECONDS, +) -> DiagnosticReport: + """Load the allowlist, probe each entry, return a full report. + + Calls `on_check(...)` once per endpoint in input order so the CLI can + stream `Checking : ✅` lines as the run progresses. + + If the queryable allowlist comes back empty (typically permission denied + on `SYSTEM$ALLOWLIST()` for low-privilege roles) and no `allowlist_path` + was supplied, fall back to checking just `(conn.host, 443)`. + """ + allowlist = load_allowlist(conn, allowlist_path) + if not allowlist and allowlist_path is None: + log.warning( + "Allowlist empty (likely permission denied on SYSTEM$ALLOWLIST()); " + "falling back to the connection host." + ) + allowlist = [ + {"type": "SNOWFLAKE_DEPLOYMENT", "host": conn.host, "port": TLS_PORT} + ] + + report = DiagnosticReport() + for host, port, type_ in _normalise_entries(allowlist): + check = check_endpoint(host, port, type_, timeout=timeout) + report.checks.append(check) + on_check(check) + return report + + +def status_line(check: EndpointCheck) -> str: + """Format a streaming line: `Checking : [ (detail)]`. + + Server-derived fields (`host`, `type`) are sanitised before display so + crafted ANSI sequences in an allowlist response cannot rewrite the + terminal. + """ + icon = {"Healthy": "✅", "Unhealthy": "❌", "Skipped": "⏭"}[check.status] + host = sanitize_for_terminal(check.host) or "" + type_ = sanitize_for_terminal(check.type) or "" + if check.status == "Healthy" and check.latency_ms is not None: + suffix = f" ({check.latency_ms} ms)" + elif check.status == "Unhealthy" and check.error: + suffix = f" ({sanitize_for_terminal(check.error)})" + else: + suffix = "" + return f"Checking {type_}: {host} {icon}{suffix}" + + +# --------------------------------------------------------------------------- # +# Network policy snapshot +# --------------------------------------------------------------------------- # + + +@dataclass +class NetworkRule: + name: str + mode: str # INGRESS | EGRESS | INTERNAL_STAGE + type: str # noqa: A003 — IPV4 | HOST_PORT | AWSVPCEID | AZURELINKID | PRIVATE_HOST_PORT + values: list[str] = field(default_factory=list) + + +@dataclass +class NetworkPolicySnapshot: + """Snapshot of network-policy state at the time of `snow connection test`. + + `account_policy` and `user_policy` are the raw values from + `SHOW PARAMETERS LIKE 'NETWORK_POLICY'` at each scope; `effective_policy` + is `user_policy or account_policy` per Snowflake precedence. + """ + + current_ip: Optional[str] = None + account_policy: Optional[str] = None + user_policy: Optional[str] = None + effective_policy: Optional[str] = None + allowed_ip_list: list[str] = field(default_factory=list) + blocked_ip_list: list[str] = field(default_factory=list) + allowed_network_rule_list: list[str] = field(default_factory=list) + blocked_network_rule_list: list[str] = field(default_factory=list) + rules: list[NetworkRule] = field(default_factory=list) + error: Optional[str] = None + + def has_policy(self) -> bool: + return self.effective_policy is not None + + +def _query_rows(conn: SnowflakeConnection, sql: str) -> list[dict]: + """Run a query and return rows as case-folded dicts; `[]` on any failure. + + Network-policy diagnostics are best-effort — a low-priv role may not be + able to `SHOW PARAMETERS` at account scope or `DESC NETWORK POLICY`. Fail + closed and let the caller surface the gap. + """ + try: + *_, cursor = conn.execute_string(sql, cursor_class=DictCursor) + return [{str(k).lower(): v for k, v in row.items()} for row in cursor] + except Exception as exc: + log.debug("Network-policy query failed: %s\n%s", sql, exc) + return [] + + +def _scalar_value(rows: list[dict], column: str) -> Optional[str]: + if not rows: + return None + val = rows[0].get(column.lower()) + return str(val) if val not in (None, "") else None + + +def _split_csv(value: Optional[str]) -> list[str]: + return [v.strip() for v in (value or "").split(",") if v.strip()] + + +def _safe_fqn(name: str) -> Optional[FQN]: + """Parse an identifier, returning `None` if it fails the FQN regex.""" + try: + return FQN.from_string(name) + except Exception: + log.debug("Could not parse identifier as FQN: %r", name) + return None + + +def _describe_network_policy( + conn: SnowflakeConnection, name: str +) -> Optional[dict[str, list[str]]]: + """Return parsed allow/block lists from `DESC NETWORK POLICY`, or None on failure.""" + fqn = _safe_fqn(name) + if fqn is None: + return None + rows = _query_rows(conn, f"DESC NETWORK POLICY {fqn.sql_identifier}") + if not rows: + return None + by_name = { + str(r["name"]).upper(): str(r["value"]) + for r in rows + if r.get("name") and r.get("value") is not None + } + return { + "allowed_ip_list": _split_csv(by_name.get("ALLOWED_IP_LIST")), + "blocked_ip_list": _split_csv(by_name.get("BLOCKED_IP_LIST")), + "allowed_network_rule_list": _split_csv( + by_name.get("ALLOWED_NETWORK_RULE_LIST") + ), + "blocked_network_rule_list": _split_csv( + by_name.get("BLOCKED_NETWORK_RULE_LIST") + ), + } + + +def _describe_network_rule( + conn: SnowflakeConnection, qualified_name: str +) -> Optional[NetworkRule]: + fqn = _safe_fqn(qualified_name) + if fqn is None: + return None + rows = _query_rows(conn, f"DESC NETWORK RULE {fqn.sql_identifier}") + if not rows: + return None + head = rows[0] + return NetworkRule( + name=qualified_name, + mode=str(head.get("mode") or ""), + type=str(head.get("type") or ""), + values=_split_csv(str(head.get("value_list") or "")), + ) + + +def collect_network_policy( + conn: SnowflakeConnection, user: Optional[str] = None +) -> NetworkPolicySnapshot: + """Best-effort snapshot of the active network policy and its referenced rules. + + Order of precedence: + 1. `SHOW PARAMETERS LIKE 'NETWORK_POLICY' FOR USER ` (user wins) + 2. `SHOW PARAMETERS LIKE 'NETWORK_POLICY' IN ACCOUNT` + 3. `DESC NETWORK POLICY ` for inline allow/block lists + 4. `DESC NETWORK RULE ` for referenced rule values + + Always returns a snapshot; failures populate `snapshot.error` instead. + """ + snapshot = NetworkPolicySnapshot() + + snapshot.current_ip = _scalar_value( + _query_rows(conn, "SELECT CURRENT_IP_ADDRESS() AS IP"), "IP" + ) + policy_pattern = identifier_to_show_like_pattern(NETWORK_POLICY_PARAM_NAME) + snapshot.account_policy = _scalar_value( + _query_rows(conn, f"SHOW PARAMETERS LIKE {policy_pattern} IN ACCOUNT"), + "value", + ) + if isinstance(user, str) and user.strip(): + user_fqn = _safe_fqn(user) + if user_fqn is not None: + snapshot.user_policy = _scalar_value( + _query_rows( + conn, + f"SHOW PARAMETERS LIKE {policy_pattern} " + f"FOR USER {user_fqn.sql_identifier}", + ), + "value", + ) + snapshot.effective_policy = snapshot.user_policy or snapshot.account_policy + + if not snapshot.effective_policy: + return snapshot + + described = _describe_network_policy(conn, snapshot.effective_policy) + if described is None: + snapshot.error = ( + f"Could not DESC NETWORK POLICY {snapshot.effective_policy} " + "(role lacks privilege?). Allowed/blocked lists not shown." + ) + return snapshot + + snapshot.allowed_ip_list = described["allowed_ip_list"] + snapshot.blocked_ip_list = described["blocked_ip_list"] + snapshot.allowed_network_rule_list = described["allowed_network_rule_list"] + snapshot.blocked_network_rule_list = described["blocked_network_rule_list"] + + for rule_name in ( + *snapshot.allowed_network_rule_list, + *snapshot.blocked_network_rule_list, + ): + rule = _describe_network_rule(conn, rule_name) + if rule is not None: + snapshot.rules.append(rule) + + return snapshot diff --git a/tests/test_connection_diagnostic.py b/tests/test_connection_diagnostic.py new file mode 100644 index 0000000000..3ee7c9acf1 --- /dev/null +++ b/tests/test_connection_diagnostic.py @@ -0,0 +1,493 @@ +# Copyright (c) 2024 Snowflake Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for `snow connection test --enable-diag` SnowCD-style diagnostics.""" + +from __future__ import annotations + +import json +import socket +import ssl +from pathlib import Path +from tempfile import NamedTemporaryFile +from unittest import mock + +import pytest +from snowflake.cli._plugins.connection.diagnostic import ( + DiagnosticReport, + EndpointCheck, + NetworkPolicySnapshot, + check_endpoint, + collect_network_policy, + is_resolvable, + load_allowlist, + run_diagnostic, + status_line, +) +from snowflake.cli.api.exceptions import CliError + +ALLOWLIST_FIXTURE = [ + { + "type": "SNOWFLAKE_DEPLOYMENT", + "host": "acct.snowflakecomputing.com", + "port": 443, + }, + {"type": "OCSP_RESPONDER", "host": "ocsp.x.com", "port": 80}, + {"type": "STAGE", "host": "*.region.snowflakecomputing.com", "port": 443}, +] + + +@pytest.mark.parametrize( + "host, expected", + [ + ("acct.snowflakecomputing.com", True), + ("sfc-stage.s3.amazonaws.com", True), + ("*.region.snowflakecomputing.com", False), + ("", False), + ], +) +def test_is_resolvable(host, expected): + assert is_resolvable(host) is expected + + +def test_status_line_includes_latency_for_healthy(): + line = status_line( + EndpointCheck("x.com", 443, "SNOWFLAKE_DEPLOYMENT", "Healthy", latency_ms=12.3) + ) + assert "✅" in line + assert "x.com" in line + assert "12.3 ms" in line + + +def test_status_line_includes_error_for_unhealthy(): + line = status_line( + EndpointCheck( + "x.com", 443, "SNOWFLAKE_DEPLOYMENT", "Unhealthy", error="DNS fail" + ) + ) + assert "❌" in line + assert "DNS fail" in line + + +def test_status_line_strips_ansi_from_server_host(): + """Server-supplied host must not be able to inject ANSI escapes into the terminal.""" + malicious = "\x1b[31mevil.com\x1b[0m" + line = status_line( + EndpointCheck(malicious, 443, "SNOWFLAKE_DEPLOYMENT", "Healthy", latency_ms=1.0) + ) + assert "\x1b" not in line + assert "evil.com" in line + + +def test_load_allowlist_from_file(): + with NamedTemporaryFile("w+", suffix=".json", delete=False) as f: + f.write(json.dumps(ALLOWLIST_FIXTURE)) + path = Path(f.name) + conn = mock.MagicMock() + out = load_allowlist(conn, path) + assert out == ALLOWLIST_FIXTURE + conn.execute_string.assert_not_called() + + +def test_load_allowlist_file_invalid_json_raises(): + with NamedTemporaryFile("w+", suffix=".json", delete=False) as f: + f.write("not json") + path = Path(f.name) + with pytest.raises(CliError): + load_allowlist(mock.MagicMock(), path) + + +def test_load_allowlist_from_query_when_no_path(): + cursor = mock.MagicMock() + cursor.fetchone.return_value = {"SYSTEM$ALLOWLIST()": json.dumps(ALLOWLIST_FIXTURE)} + conn = mock.MagicMock() + conn.execute_string.return_value = (cursor,) + assert load_allowlist(conn, None) == ALLOWLIST_FIXTURE + + +def test_load_allowlist_merges_privatelink_entries(): + """SYSTEM$ALLOWLIST_PRIVATELINK() entries are appended after the public set.""" + pl_extra = [ + { + "type": "SNOWFLAKE_DEPLOYMENT", + "host": "acct.privatelink.snowflakecomputing.com", + "port": 443, + } + ] + + def execute_string(sql, cursor_class=None): + cur = mock.MagicMock() + if "PRIVATELINK" in sql: + cur.fetchone.return_value = { + "SYSTEM$ALLOWLIST_PRIVATELINK()": json.dumps(pl_extra) + } + else: + cur.fetchone.return_value = { + "SYSTEM$ALLOWLIST()": json.dumps(ALLOWLIST_FIXTURE) + } + return (cur,) + + conn = mock.MagicMock() + conn.execute_string.side_effect = execute_string + out = load_allowlist(conn, None) + assert out == ALLOWLIST_FIXTURE + pl_extra + + +def test_load_allowlist_dedupes_overlap_with_privatelink(): + """A (type, host, port) appearing in both lists collapses to one entry.""" + overlap = ALLOWLIST_FIXTURE[0] # SNOWFLAKE_DEPLOYMENT acct... + + def execute_string(sql, cursor_class=None): + cur = mock.MagicMock() + if "PRIVATELINK" in sql: + cur.fetchone.return_value = { + "SYSTEM$ALLOWLIST_PRIVATELINK()": json.dumps([overlap]) + } + else: + cur.fetchone.return_value = { + "SYSTEM$ALLOWLIST()": json.dumps(ALLOWLIST_FIXTURE) + } + return (cur,) + + conn = mock.MagicMock() + conn.execute_string.side_effect = execute_string + out = load_allowlist(conn, None) + # Same length — the duplicate from PL was dropped. + assert len(out) == len(ALLOWLIST_FIXTURE) + + +def test_load_allowlist_silently_skips_failing_privatelink_query(): + """If SYSTEM$ALLOWLIST_PRIVATELINK is unavailable the public list is still returned.""" + + def execute_string(sql, cursor_class=None): + if "PRIVATELINK" in sql: + raise RuntimeError("Function not enabled in this deployment") + cur = mock.MagicMock() + cur.fetchone.return_value = { + "SYSTEM$ALLOWLIST()": json.dumps(ALLOWLIST_FIXTURE) + } + return (cur,) + + conn = mock.MagicMock() + conn.execute_string.side_effect = execute_string + assert load_allowlist(conn, None) == ALLOWLIST_FIXTURE + + +def test_check_endpoint_skips_wildcards(): + result = check_endpoint("*.x.com", 443, "STAGE") + assert result.status == "Skipped" + assert result.error == "non-resolvable pattern" + + +def test_check_endpoint_marks_unhealthy_on_dns_failure(): + with mock.patch( + "snowflake.cli._plugins.connection.diagnostic.socket.create_connection", + side_effect=socket.gaierror("Name or service not known"), + ): + result = check_endpoint("nonexistent.invalid", 443, "SNOWFLAKE_DEPLOYMENT") + assert result.status == "Unhealthy" + assert "Name or service" in (result.error or "") + assert result.latency_ms is None + + +def test_check_endpoint_marks_unhealthy_on_tls_failure(): + fake_sock = mock.MagicMock() + with ( + mock.patch( + "snowflake.cli._plugins.connection.diagnostic.socket.create_connection", + return_value=fake_sock, + ), + mock.patch( + "snowflake.cli._plugins.connection.diagnostic._probe_tls", + side_effect=ssl.SSLError("cert verify failed"), + ), + ): + result = check_endpoint("x.com", 443, "SNOWFLAKE_DEPLOYMENT") + assert result.status == "Unhealthy" + assert "TLS handshake failed" in (result.error or "") + + +def test_check_endpoint_records_latency_and_cert(): + fake_sock = mock.MagicMock() + with ( + mock.patch( + "snowflake.cli._plugins.connection.diagnostic.socket.create_connection", + return_value=fake_sock, + ), + mock.patch( + "snowflake.cli._plugins.connection.diagnostic._probe_tls", + return_value=("DigiCert Inc", "Jan 19 23:59:59 2027 GMT"), + ), + ): + result = check_endpoint("x.com", 443, "SNOWFLAKE_DEPLOYMENT") + assert result.status == "Healthy" + assert result.cert_issuer == "DigiCert Inc" + assert result.cert_expires == "Jan 19 23:59:59 2027 GMT" + assert result.latency_ms is not None + assert result.latency_ms >= 0 + + +def test_check_endpoint_skips_tls_for_port_80(): + fake_sock = mock.MagicMock() + with ( + mock.patch( + "snowflake.cli._plugins.connection.diagnostic.socket.create_connection", + return_value=fake_sock, + ), + mock.patch( + "snowflake.cli._plugins.connection.diagnostic._probe_tls" + ) as probe_tls, + ): + result = check_endpoint("crl.example.com", 80, "CRL_DISTRIBUTION_POINT") + probe_tls.assert_not_called() + assert result.status == "Healthy" + assert result.cert_issuer is None + assert result.cert_expires is None + + +def test_diagnostic_report_counts(): + report = DiagnosticReport( + checks=[ + EndpointCheck("a", 443, "X", "Healthy"), + EndpointCheck("b", 443, "X", "Healthy"), + EndpointCheck("c", 443, "X", "Unhealthy"), + EndpointCheck("d", 443, "X", "Skipped"), + ] + ) + assert report.healthy == 2 + assert report.unhealthy == 1 + assert report.skipped == 1 + assert report.tested == 3 + assert ( + report.summary_line() + == "Results: 2 Healthy, 1 Unhealthy out of 3 endpoints. 1 skipped (non-resolvable patterns)." + ) + + +def test_run_diagnostic_streams_per_endpoint(monkeypatch): + cursor = mock.MagicMock() + cursor.fetchone.return_value = {"SYSTEM$ALLOWLIST()": json.dumps(ALLOWLIST_FIXTURE)} + conn = mock.MagicMock() + conn.execute_string.return_value = (cursor,) + + fake_sock = mock.MagicMock() + monkeypatch.setattr( + "snowflake.cli._plugins.connection.diagnostic.socket.create_connection", + lambda *a, **kw: fake_sock, + ) + monkeypatch.setattr( + "snowflake.cli._plugins.connection.diagnostic._probe_tls", + lambda sock, host: ("DigiCert Inc", "Jan 19 23:59:59 2027 GMT"), + ) + + streamed = [] + report = run_diagnostic(conn, allowlist_path=None, on_check=streamed.append) + + assert len(streamed) == 3 + assert report.healthy == 2 + assert report.skipped == 1 + assert [c.type for c in streamed] == [ + "SNOWFLAKE_DEPLOYMENT", + "OCSP_RESPONDER", + "STAGE", + ] + + +def test_run_diagnostic_falls_back_when_allowlist_query_fails(monkeypatch): + conn = mock.MagicMock() + conn.host = "fallback.snowflakecomputing.com" + conn.execute_string.side_effect = RuntimeError( + "permission denied on SYSTEM$ALLOWLIST" + ) + + fake_sock = mock.MagicMock() + monkeypatch.setattr( + "snowflake.cli._plugins.connection.diagnostic.socket.create_connection", + lambda *a, **kw: fake_sock, + ) + monkeypatch.setattr( + "snowflake.cli._plugins.connection.diagnostic._probe_tls", + lambda sock, host: (None, None), + ) + + report = run_diagnostic(conn, allowlist_path=None) + assert report.tested == 1 + assert report.checks[0].host == "fallback.snowflakecomputing.com" + assert report.checks[0].type == "SNOWFLAKE_DEPLOYMENT" + + +def test_run_diagnostic_normalises_malformed_entries(monkeypatch): + cursor = mock.MagicMock() + cursor.fetchone.return_value = { + "SYSTEM$ALLOWLIST()": json.dumps( + [ + {"type": "X", "host": "good.com", "port": 443}, + {"type": "X", "host": None, "port": 443}, + {"type": "X", "host": "bad.com", "port": "not-a-number"}, + ] + ) + } + conn = mock.MagicMock() + conn.execute_string.return_value = (cursor,) + monkeypatch.setattr( + "snowflake.cli._plugins.connection.diagnostic.socket.create_connection", + lambda *a, **kw: mock.MagicMock(), + ) + monkeypatch.setattr( + "snowflake.cli._plugins.connection.diagnostic._probe_tls", + lambda *a, **kw: (None, None), + ) + report = run_diagnostic(conn, allowlist_path=None) + assert [c.host for c in report.checks] == ["good.com"] + + +def _scripted_cursor(rows): + cur = mock.MagicMock() + cur.__iter__ = lambda self: iter(rows) + return cur + + +def test_collect_network_policy_user_overrides_account(): + conn = mock.MagicMock() + seen_sql: list[str] = [] + + def execute_string(sql, cursor_class=None): + seen_sql.append(sql) + if "CURRENT_IP_ADDRESS" in sql: + return (_scripted_cursor([{"IP": "1.2.3.4"}]),) + if "SHOW PARAMETERS" in sql and "IN ACCOUNT" in sql: + return (_scripted_cursor([{"value": "ACCT_POLICY"}]),) + if "SHOW PARAMETERS" in sql and "FOR USER" in sql: + return (_scripted_cursor([{"value": "USER_POLICY"}]),) + if "DESC NETWORK POLICY" in sql: + return ( + _scripted_cursor( + [ + {"name": "ALLOWED_IP_LIST", "value": "0.0.0.0/0"}, + {"name": "BLOCKED_IP_LIST", "value": ""}, + {"name": "ALLOWED_NETWORK_RULE_LIST", "value": "DB.SCHEMA.NR1"}, + {"name": "BLOCKED_NETWORK_RULE_LIST", "value": ""}, + ] + ), + ) + if "DESC NETWORK RULE" in sql: + return ( + _scripted_cursor( + [ + { + "mode": "INGRESS", + "type": "IPV4", + "value_list": "1.0.0.0/8,2.0.0.0/8", + } + ] + ), + ) + return (_scripted_cursor([]),) + + conn.execute_string.side_effect = execute_string + snap = collect_network_policy(conn, user="alice") + assert snap.current_ip == "1.2.3.4" + assert snap.account_policy == "ACCT_POLICY" + assert snap.user_policy == "USER_POLICY" + assert snap.effective_policy == "USER_POLICY" + assert snap.allowed_ip_list == ["0.0.0.0/0"] + assert snap.allowed_network_rule_list == ["DB.SCHEMA.NR1"] + assert len(snap.rules) == 1 + assert snap.rules[0].values == ["1.0.0.0/8", "2.0.0.0/8"] + assert snap.rules[0].mode == "INGRESS" + assert snap.rules[0].type == "IPV4" + # Identifiers must be wrapped in IDENTIFIER('...') — never bare-interpolated. + user_sql = next(s for s in seen_sql if "FOR USER" in s) + assert "IDENTIFIER('alice')" in user_sql + desc_sql = next(s for s in seen_sql if "DESC NETWORK POLICY" in s) + assert "IDENTIFIER('USER_POLICY')" in desc_sql + rule_sql = next(s for s in seen_sql if "DESC NETWORK RULE" in s) + assert "IDENTIFIER('DB.SCHEMA.NR1')" in rule_sql + + +def test_collect_network_policy_handles_no_policy(): + conn = mock.MagicMock() + conn.execute_string.return_value = (_scripted_cursor([]),) + snap = collect_network_policy(conn, user="alice") + assert snap.has_policy() is False + assert snap.effective_policy is None + + +def test_collect_network_policy_survives_query_errors(): + conn = mock.MagicMock() + conn.execute_string.side_effect = RuntimeError("insufficient privileges") + snap = collect_network_policy(conn, user="alice") + assert isinstance(snap, NetworkPolicySnapshot) + assert snap.has_policy() is False + + +def test_resolve_cafile_prefers_requests_ca_bundle(tmp_path, monkeypatch): + from snowflake.cli._plugins.connection.diagnostic import _resolve_cafile + + primary = tmp_path / "primary.pem" + primary.write_text("primary") + secondary = tmp_path / "secondary.pem" + secondary.write_text("secondary") + monkeypatch.setenv("REQUESTS_CA_BUNDLE", str(primary)) + monkeypatch.setenv("SSL_CERT_FILE", str(secondary)) + assert _resolve_cafile() == str(primary) + + +def test_resolve_cafile_falls_back_to_ssl_cert_file(tmp_path, monkeypatch): + from snowflake.cli._plugins.connection.diagnostic import _resolve_cafile + + bundle = tmp_path / "bundle.pem" + bundle.write_text("bundle") + monkeypatch.delenv("REQUESTS_CA_BUNDLE", raising=False) + monkeypatch.setenv("SSL_CERT_FILE", str(bundle)) + assert _resolve_cafile() == str(bundle) + + +def test_resolve_cafile_returns_none_when_env_points_at_missing_file(monkeypatch): + from snowflake.cli._plugins.connection.diagnostic import _resolve_cafile + + monkeypatch.setenv("REQUESTS_CA_BUNDLE", "/nonexistent/ca.pem") + monkeypatch.delenv("SSL_CERT_FILE", raising=False) + assert _resolve_cafile() is None + + +def test_resolve_cafile_returns_none_when_unset(monkeypatch): + from snowflake.cli._plugins.connection.diagnostic import _resolve_cafile + + monkeypatch.delenv("REQUESTS_CA_BUNDLE", raising=False) + monkeypatch.delenv("SSL_CERT_FILE", raising=False) + assert _resolve_cafile() is None + + +def test_probe_tls_passes_cafile_to_ssl_context(tmp_path, monkeypatch): + """_probe_tls plumbs the resolved CA bundle into ssl.create_default_context.""" + bundle = tmp_path / "bundle.pem" + bundle.write_text("bundle") + monkeypatch.setenv("REQUESTS_CA_BUNDLE", str(bundle)) + + fake_context = mock.MagicMock() + fake_ssock = mock.MagicMock() + fake_ssock.__enter__ = lambda self: fake_ssock + fake_ssock.__exit__ = lambda self, *a: False + fake_ssock.getpeercert.return_value = {} + fake_context.wrap_socket.return_value = fake_ssock + + with mock.patch( + "snowflake.cli._plugins.connection.diagnostic.ssl.create_default_context", + return_value=fake_context, + ) as create_default_context: + from snowflake.cli._plugins.connection.diagnostic import _probe_tls + + _probe_tls(mock.MagicMock(), "x.com") + create_default_context.assert_called_once_with(cafile=str(bundle))