""" NSX Policy API client supporting NSX 3.x, 4.x, and 9.x managers. Uses the Policy API (preferred over the deprecated Manager API). """ import logging import time from typing import Any, Dict, List, Optional, Tuple import requests import urllib3 logger = logging.getLogger(__name__) _POLICY_BASE = "/policy/api/v1" _PAGE_SIZE = 1000 _MAX_RETRIES = 3 _RETRY_BACKOFF = 2 # seconds class NSXAuthError(Exception): pass class NSXAPIError(Exception): def __init__(self, message: str, status_code: int = 0, body: str = ""): super().__init__(message) self.status_code = status_code self.body = body class NSXClient: """ Thin wrapper around the NSX Policy REST API. Supports basic auth (username + password) and handles: - Cursor-based pagination - Automatic retries on transient errors - Optional SSL verification bypass (for self-signed certs) """ def __init__( self, host: str, username: str, password: str, verify_ssl: bool = True, timeout: int = 30, label: str = "", ): self.base_url = f"https://{host.rstrip('/')}{_POLICY_BASE}" self.timeout = timeout self.label = label or host if not verify_ssl: urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) self._session = requests.Session() self._session.auth = (username, password) self._session.verify = verify_ssl self._session.headers.update( {"Content-Type": "application/json", "Accept": "application/json"} ) # ------------------------------------------------------------------ # Internal helpers # ------------------------------------------------------------------ def _url(self, path: str) -> str: return f"{self.base_url}/{path.lstrip('/')}" def _request( self, method: str, path: str, **kwargs, ) -> requests.Response: url = self._url(path) last_exc: Optional[Exception] = None for attempt in range(1, _MAX_RETRIES + 1): try: resp = self._session.request( method, url, timeout=self.timeout, **kwargs ) if resp.status_code == 401: raise NSXAuthError( f"[{self.label}] Authentication failed – check credentials." ) if resp.status_code == 403: raise NSXAuthError( f"[{self.label}] Access denied to {path} – check permissions." ) if resp.status_code >= 500 and attempt < _MAX_RETRIES: logger.warning( "[%s] HTTP %s on %s, retry %d/%d", self.label, resp.status_code, path, attempt, _MAX_RETRIES, ) time.sleep(_RETRY_BACKOFF * attempt) continue if not resp.ok: raise NSXAPIError( f"[{self.label}] HTTP {resp.status_code} on {method} {path}", status_code=resp.status_code, body=resp.text, ) return resp except (requests.ConnectionError, requests.Timeout) as exc: last_exc = exc if attempt < _MAX_RETRIES: logger.warning( "[%s] Connection error on %s, retry %d/%d: %s", self.label, path, attempt, _MAX_RETRIES, exc, ) time.sleep(_RETRY_BACKOFF * attempt) raise NSXAPIError( f"[{self.label}] Failed to reach {path} after {_MAX_RETRIES} attempts: {last_exc}" ) def _get_paged(self, path: str, params: Optional[Dict] = None) -> List[Dict]: """Collect all pages for a list endpoint using cursor-based pagination.""" params = dict(params or {}) params["page_size"] = _PAGE_SIZE results: List[Dict] = [] while True: resp = self._request("GET", path, params=params) data = resp.json() results.extend(data.get("results", [])) cursor = data.get("cursor") if not cursor: break params["cursor"] = cursor return results # ------------------------------------------------------------------ # Version / connectivity check # ------------------------------------------------------------------ def get_version(self) -> str: """Return the NSX manager version string.""" resp = self._request("GET", "/infra/sites") # Version is available via a different endpoint try: resp2 = self._session.get( f"https://{self.base_url.split('//')[1].split('/')[0]}/api/v1/node/version", timeout=self.timeout, ) if resp2.ok: return resp2.json().get("product_version", "unknown") except Exception: pass return "unknown" # ------------------------------------------------------------------ # Groups # ------------------------------------------------------------------ def list_groups(self, domain: str = "default") -> List[Dict]: """Return all security groups in the given policy domain.""" logger.debug("[%s] Listing groups in domain '%s'", self.label, domain) return self._get_paged(f"/infra/domains/{domain}/groups") def get_group(self, group_id: str, domain: str = "default") -> Dict: """Return a single group definition by ID or display_name.""" groups = self.list_groups(domain) # Try exact ID match first, then display_name for g in groups: if g.get("id") == group_id or g.get("display_name") == group_id: return g raise NSXAPIError( f"[{self.label}] Group '{group_id}' not found in domain '{domain}'" ) def get_group_effective_ips( self, group_id: str, domain: str = "default" ) -> List[str]: """ Return the list of effective (computed) IP addresses that are members of the group. This resolves VM-based, tag-based, etc. criteria. """ logger.debug( "[%s] Fetching effective IPs for group '%s'", self.label, group_id ) path = f"/infra/domains/{domain}/groups/{group_id}/members/ip-addresses" raw = self._get_paged(path) # Each item is a plain string (the IP / CIDR) ips = [item if isinstance(item, str) else item.get("ip_address", "") for item in raw] return sorted(ip for ip in ips if ip) # ------------------------------------------------------------------ # IP Address Expression helpers # ------------------------------------------------------------------ @staticmethod def _extract_ip_expressions( expressions: List[Dict], ) -> Tuple[List[Dict], List[str]]: """ Split expressions into (non-ip-entries, existing-static-ips). Removes IPAddressExpression items and their preceding ConjunctionOperators. Returns a clean expression list ready for re-injection and the existing IPs. """ existing_ips: List[str] = [] clean: List[Dict] = [] skip_next_conjunction = False for expr in expressions: rtype = expr.get("resource_type", "") if rtype == "IPAddressExpression": existing_ips.extend(expr.get("ip_addresses", [])) # Remove the conjunction that precedes this entry (if we added it) if clean and clean[-1].get("resource_type") == "ConjunctionOperator": clean.pop() continue clean.append(expr) _ = skip_next_conjunction # unused, kept for clarity return clean, existing_ips def patch_group_ips( self, group_id: str, ips_to_add: List[str], ips_to_remove: List[str], domain: str = "default", whatif: bool = False, ) -> Dict: """ Modify the static IP address entries in a target group: - Adds ips_to_add as static entries - Removes ips_to_remove from static entries - Leaves all non-IP expressions (Conditions, Tags, etc.) untouched Returns a result dict describing what was (or would have been) done. """ group = self.get_group(group_id, domain) gid = group["id"] display = group.get("display_name", gid) expressions = list(group.get("expression", [])) base_exprs, current_static_ips = self._extract_ip_expressions(expressions) current_set = set(current_static_ips) desired_set = (current_set | set(ips_to_add)) - set(ips_to_remove) actually_added = sorted(set(ips_to_add) - current_set) actually_removed = sorted(current_set & set(ips_to_remove)) result = { "group_id": gid, "display_name": display, "ips_added": actually_added, "ips_removed": actually_removed, "whatif": whatif, } if not actually_added and not actually_removed: logger.info("[%s] Group '%s': no changes needed.", self.label, display) result["status"] = "no_change" return result if whatif: logger.info( "[%s] [WHATIF] Group '%s': would add %d, remove %d IPs.", self.label, display, len(actually_added), len(actually_removed), ) result["status"] = "whatif" return result # Build the new expression list new_exprs = list(base_exprs) if desired_set: if new_exprs: new_exprs.append( {"resource_type": "ConjunctionOperator", "conjunction_operator": "OR"} ) new_exprs.append( { "resource_type": "IPAddressExpression", "ip_addresses": sorted(desired_set), } ) patch_body = dict(group) patch_body["expression"] = new_exprs self._request( "PATCH", f"/infra/domains/{domain}/groups/{gid}", json=patch_body, ) logger.info( "[%s] Group '%s': added %d, removed %d IPs.", self.label, display, len(actually_added), len(actually_removed), ) result["status"] = "updated" return result