agent-governance/preflight/inventory_check.py

#!/usr/bin/env python3
"""
Inventory Validator
===================
Validates that target resources exist and are in allowed pools.
Part of Phase 3: Execution Pipeline - Preflight System.
"""

import json
import subprocess
import sys
from dataclasses import dataclass
from datetime import datetime, timezone
from enum import Enum
from typing import Optional


class CheckStatus(str, Enum):
    PASS = "PASS"
    FAIL = "FAIL"
    WARN = "WARN"
    SKIP = "SKIP"


@dataclass
class CheckResult:
    check_name: str
    status: CheckStatus
    message: str
    details: dict
    timestamp: str

    def to_dict(self) -> dict:
        return {
            "check_name": self.check_name,
            "status": self.status.value,
            "message": self.message,
            "details": self.details,
            "timestamp": self.timestamp
        }


class InventoryChecker:
    """
    Validates inventory targets against allowed pools.

    Checks:
    - Target exists in inventory
    - Target is in an allowed pool for the agent's tier
    - Target is not in a forbidden pool
    """

    POOL_TIERS = {
        "pve-sandbox": 1,    # Tier 1+ can access
        "pve-staging": 2,    # Tier 2+ can access
        "pve-prod": 3,       # Tier 3+ with approval
    }

    FORBIDDEN_POOLS = {
        0: ["pve-sandbox", "pve-staging", "pve-prod"],  # Tier 0: no pools
        1: ["pve-staging", "pve-prod"],                  # Tier 1: only sandbox
        2: ["pve-prod"],                                  # Tier 2: sandbox + staging
        3: [],                                            # Tier 3: all with approval
        4: [],                                            # Tier 4: all
    }

    def __init__(self):
        self.vault_token = self._get_vault_token()

    def _get_vault_token(self) -> str:
        with open("/opt/vault/init-keys.json") as f:
            return json.load(f)["root_token"]

    def _now(self) -> str:
        return datetime.now(timezone.utc).isoformat()

    def _vault_read(self, path: str) -> Optional[dict]:
        """Read from Vault KV"""
        result = subprocess.run([
            "curl", "-sk",
            "-H", f"X-Vault-Token: {self.vault_token}",
            f"https://127.0.0.1:8200/v1/secret/data/{path}"
        ], capture_output=True, text=True)

        try:
            data = json.loads(result.stdout)
            if "data" in data and "data" in data["data"]:
                return data["data"]["data"]
        except:
            pass
        return None

    def get_inventory(self) -> dict:
        """Get infrastructure inventory from Vault"""
        inventory = self._vault_read("inventory/proxmox")
        if not inventory:
            # Return mock inventory for testing
            return {
                "cluster": "pve-cluster-01",
                "pools": {
                    "pve-sandbox": {
                        "nodes": ["sandbox-vm-01", "sandbox-vm-02", "sandbox-vm-03"],
                        "tier_required": 1
                    },
                    "pve-staging": {
                        "nodes": ["staging-vm-01", "staging-vm-02"],
                        "tier_required": 2
                    },
                    "pve-prod": {
                        "nodes": ["prod-vm-01", "prod-vm-02", "prod-db-01"],
                        "tier_required": 3
                    }
                }
            }
        return inventory

    def check_target_exists(self, target: str) -> CheckResult:
        """Check if target exists in any pool"""
        inventory = self.get_inventory()

        for pool_name, pool_data in inventory.get("pools", {}).items():
            nodes = pool_data.get("nodes", [])
            if target in nodes:
                return CheckResult(
                    check_name="target_exists",
                    status=CheckStatus.PASS,
                    message=f"Target '{target}' found in pool '{pool_name}'",
                    details={"target": target, "pool": pool_name, "nodes": nodes},
                    timestamp=self._now()
                )

        return CheckResult(
            check_name="target_exists",
            status=CheckStatus.FAIL,
            message=f"Target '{target}' not found in any inventory pool",
            details={"target": target, "searched_pools": list(inventory.get("pools", {}).keys())},
            timestamp=self._now()
        )

    def check_pool_access(self, target: str, agent_tier: int) -> CheckResult:
        """Check if agent tier has access to target's pool"""
        inventory = self.get_inventory()

        # Find which pool contains the target
        target_pool = None
        for pool_name, pool_data in inventory.get("pools", {}).items():
            if target in pool_data.get("nodes", []):
                target_pool = pool_name
                break

        if not target_pool:
            return CheckResult(
                check_name="pool_access",
                status=CheckStatus.FAIL,
                message=f"Target '{target}' not found in inventory",
                details={"target": target},
                timestamp=self._now()
            )

        # Check if pool is forbidden for this tier
        forbidden = self.FORBIDDEN_POOLS.get(agent_tier, [])
        if target_pool in forbidden:
            return CheckResult(
                check_name="pool_access",
                status=CheckStatus.FAIL,
                message=f"Tier {agent_tier} agents cannot access pool '{target_pool}'",
                details={
                    "target": target,
                    "pool": target_pool,
                    "agent_tier": agent_tier,
                    "forbidden_pools": forbidden
                },
                timestamp=self._now()
            )

        return CheckResult(
            check_name="pool_access",
            status=CheckStatus.PASS,
            message=f"Tier {agent_tier} has access to pool '{target_pool}'",
            details={
                "target": target,
                "pool": target_pool,
                "agent_tier": agent_tier
            },
            timestamp=self._now()
        )

    def check_not_production(self, target: str) -> CheckResult:
        """Warn if target is in production pool"""
        inventory = self.get_inventory()

        for pool_name, pool_data in inventory.get("pools", {}).items():
            if target in pool_data.get("nodes", []):
                if "prod" in pool_name.lower():
                    return CheckResult(
                        check_name="not_production",
                        status=CheckStatus.WARN,
                        message=f"Target '{target}' is in PRODUCTION pool '{pool_name}'",
                        details={
                            "target": target,
                            "pool": pool_name,
                            "warning": "Production access requires explicit approval"
                        },
                        timestamp=self._now()
                    )
                return CheckResult(
                    check_name="not_production",
                    status=CheckStatus.PASS,
                    message=f"Target '{target}' is not in production",
                    details={"target": target, "pool": pool_name},
                    timestamp=self._now()
                )

        return CheckResult(
            check_name="not_production",
            status=CheckStatus.SKIP,
            message=f"Target '{target}' not found in inventory",
            details={"target": target},
            timestamp=self._now()
        )

    def run_all_checks(self, target: str, agent_tier: int) -> list[CheckResult]:
        """Run all inventory checks"""
        results = []

        # Check 1: Target exists
        results.append(self.check_target_exists(target))

        # Check 2: Pool access
        results.append(self.check_pool_access(target, agent_tier))

        # Check 3: Production warning
        results.append(self.check_not_production(target))

        return results

    def preflight_report(self, targets: list[str], agent_tier: int) -> dict:
        """Generate a full preflight report for multiple targets"""
        report = {
            "report_type": "inventory_preflight",
            "agent_tier": agent_tier,
            "timestamp": self._now(),
            "targets": {},
            "summary": {
                "total_checks": 0,
                "passed": 0,
                "failed": 0,
                "warnings": 0,
                "skipped": 0
            },
            "can_proceed": True
        }

        for target in targets:
            results = self.run_all_checks(target, agent_tier)
            report["targets"][target] = [r.to_dict() for r in results]

            for r in results:
                report["summary"]["total_checks"] += 1
                if r.status == CheckStatus.PASS:
                    report["summary"]["passed"] += 1
                elif r.status == CheckStatus.FAIL:
                    report["summary"]["failed"] += 1
                    report["can_proceed"] = False
                elif r.status == CheckStatus.WARN:
                    report["summary"]["warnings"] += 1
                elif r.status == CheckStatus.SKIP:
                    report["summary"]["skipped"] += 1

        return report


# =============================================================================
# CLI
# =============================================================================

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Inventory Preflight Checker")
    parser.add_argument("targets", nargs="+", help="Target nodes to check")
    parser.add_argument("--tier", type=int, default=1, help="Agent tier (0-4)")
    parser.add_argument("--json", action="store_true", help="Output JSON")

    args = parser.parse_args()

    checker = InventoryChecker()
    report = checker.preflight_report(args.targets, args.tier)

    if args.json:
        print(json.dumps(report, indent=2))
    else:
        print("\n" + "=" * 60)
        print("INVENTORY PREFLIGHT REPORT")
        print("=" * 60)
        print(f"Agent Tier: {report['agent_tier']}")
        print(f"Timestamp: {report['timestamp']}")
        print()

        for target, checks in report["targets"].items():
            print(f"\nTarget: {target}")
            print("-" * 40)
            for check in checks:
                status_icon = {
                    "PASS": "[OK]",
                    "FAIL": "[FAIL]",
                    "WARN": "[WARN]",
                    "SKIP": "[SKIP]"
                }.get(check["status"], "[?]")
                print(f"  {status_icon} {check['check_name']}: {check['message']}")

        print("\n" + "-" * 60)
        print("SUMMARY")
        print("-" * 60)
        s = report["summary"]
        print(f"  Total Checks: {s['total_checks']}")
        print(f"  Passed: {s['passed']}")
        print(f"  Failed: {s['failed']}")
        print(f"  Warnings: {s['warnings']}")
        print(f"  Skipped: {s['skipped']}")
        print()

        if report["can_proceed"]:
            print("[OK] PREFLIGHT PASSED - Can proceed with execution")
        else:
            print("[FAIL] PREFLIGHT FAILED - Cannot proceed")

        print("=" * 60)

    sys.exit(0 if report["can_proceed"] else 1)