#!/bin/bash
set -euo pipefail

# verify-package.sh: Verify a Python package against pip-witness attestations
#
# Usage:
#   ./verify-package.sh requests                    # verify latest
#   ./verify-package.sh requests==2.33.1             # verify specific version
#   ./verify-package.sh --scan-first litellm         # scan then verify
#
# This script:
# 1. Looks up the package attestation in Archivista (archivista.testifysec.io)
# 2. Verifies the attestation against our signed Rego security policies
# 3. Produces a policyverify attestation as evidence of the evaluation
# 4. Stores the verification result back in Archivista

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
ARCHIVISTA_SERVER="${ARCHIVISTA_SERVER:-https://archivista.testifysec.io}"
POLICY_FILE="${SCRIPT_DIR}/policies/pip-witness-policy-signed.json"
POLICY_KEY="${SCRIPT_DIR}/policies/policy-key.pub"
SIGNING_KEY="${SCRIPT_DIR}/policies/policy-key.pem"

SCAN_FIRST=false
PACKAGE=""

usage() {
    echo "verify-package.sh: Verify a Python package against pip-witness attestations"
    echo ""
    echo "Usage: ./verify-package.sh [options] <package-spec>"
    echo ""
    echo "Options:"
    echo "  --scan-first    Scan the package first, then verify"
    echo "  --archivista    Archivista server URL (default: archivista.testifysec.io)"
    echo ""
    echo "The script looks up the scan attestation in Archivista by package name,"
    echo "then evaluates it against 9 Rego security policies covering:"
    echo "  - Network exfiltration (unknown IPs contacted during install)"
    echo "  - Credential harvesting (.ssh, .aws, .kube file access)"
    echo "  - Code execution (shells, curl, wget spawned)"
    echo "  - Persistence (cron, systemd, .pth injection)"
    echo "  - TeamPCP IOCs (LiteLLM attack patterns)"
    echo "  - OpenClaw IOCs (browser/wallet theft patterns)"
    echo "  - Supply chain (setup.py analysis)"
    echo "  - Package signing (PEP 740 provenance)"
    echo ""
    echo "Results are stored as a policyverify attestation in Archivista."
    exit 1
}

while [[ $# -gt 0 ]]; do
    case "$1" in
        --scan-first) SCAN_FIRST=true; shift ;;
        --archivista) ARCHIVISTA_SERVER="$2"; shift 2 ;;
        -h|--help) usage ;;
        *) PACKAGE="$1"; shift ;;
    esac
done

[ -z "${PACKAGE}" ] && usage

echo "=========================================="
echo "  pip-witness verify: ${PACKAGE}"
echo "=========================================="
echo "  Archivista: ${ARCHIVISTA_SERVER}"

# Step 1: Optionally scan first
if [ "${SCAN_FIRST}" = true ]; then
    echo ""
    echo "[1/3] Scanning ${PACKAGE} in sandbox..."
    ./pip-witness "${PACKAGE}"
fi

# Step 2: Find the attestation in Archivista
echo ""
echo "[2/3] Looking up attestation in Archivista..."

# Search for attestations with this collection name (step name = pip-install-<pkg>)
STEP_NAME="pip-install-$(echo "${PACKAGE}" | sed 's/[^a-zA-Z0-9]/-/g' | tr '[:upper:]' '[:lower:]')"

# Query Archivista GraphQL for matching attestations
SEARCH_RESULT=$(curl -s -X POST "${ARCHIVISTA_SERVER}/query" \
    -H 'Content-Type: application/json' \
    -d "{
        \"query\": \"query { dsses(where: { hasStatementWith: { hasAttestationCollectionsWith: { name: \\\"${STEP_NAME}\\\" } } }, orderBy: { direction: DESC, field: CREATED_AT }, first: 1) { edges { node { gitoidSha256 } } } }\"
    }" 2>/dev/null)

GITOID=$(echo "${SEARCH_RESULT}" | python3 -c "
import json, sys
try:
    data = json.load(sys.stdin)
    edges = data.get('data', {}).get('dsses', {}).get('edges', [])
    if edges:
        print(edges[0]['node']['gitoidSha256'])
except:
    pass
" 2>/dev/null)

if [ -z "${GITOID}" ]; then
    echo "  No attestation found for '${STEP_NAME}' in Archivista"
    echo "  Run with --scan-first to create one, or check the step name"

    # Try broader search
    echo "  Trying broader search..."
    SEARCH_RESULT2=$(curl -s -X POST "${ARCHIVISTA_SERVER}/query" \
        -H 'Content-Type: application/json' \
        -d "{
            \"query\": \"query { dsses(where: { hasStatementWith: { hasAttestationCollectionsWith: { nameContains: \\\"${PACKAGE%%=*}\\\" } } }, first: 5) { edges { node { gitoidSha256 } } } }\"
        }" 2>/dev/null)

    echo "  Search result: ${SEARCH_RESULT2}" | head -2
    exit 1
fi

echo "  Found attestation: ${GITOID}"

# Download the attestation
ATTESTATION_FILE="/tmp/pip-witness-verify-${GITOID}.json"
curl -s "${ARCHIVISTA_SERVER}/download/${GITOID}" -o "${ATTESTATION_FILE}"
echo "  Downloaded to: ${ATTESTATION_FILE}"

# Step 3: Evaluate policies
echo ""
echo "[3/3] Evaluating against security policies..."

# Run OPA policy evaluation
POLICY_RESULT=$(python3 "${SCRIPT_DIR}/viewer/policy_eval.py" "${ATTESTATION_FILE}" 2>&1)

echo "${POLICY_RESULT}" | python3 -c "
import json, sys
r = json.load(sys.stdin)
print(f'  Risk Level: {r[\"risk_level\"]}')
print(f'  Risk Score: {r[\"risk_score\"]}')
print(f'  Violations: {r[\"total_violations\"]}')
print()
for name, pol in r.get('policies', {}).items():
    v = pol.get('violations', [])
    status = 'PASS' if not v else f'FAIL ({len(v)})'
    icon = '✓' if not v else '✗'
    print(f'  {icon} {name:30s} {status}')
    for violation in v[:3]:
        print(f'      {violation[:100]}')
    if len(v) > 3:
        print(f'      ... +{len(v)-3} more')
"

echo ""
echo "=========================================="

# If we have a signing key, produce a policyverify attestation and upload to Archivista
if [ -f "${SIGNING_KEY}" ]; then
    echo "  Signing verification result and uploading to Archivista..."

    VERIFY_OUT="/tmp/pip-witness-policyverify-${GITOID}.json"

    # Store the policy evaluation result as a signed envelope
    echo "${POLICY_RESULT}" | python3 -c "
import json, sys, base64
result = json.load(sys.stdin)
# Wrap in in-toto statement format
statement = {
    '_type': 'https://in-toto.io/Statement/v0.1',
    'predicateType': 'https://aflock.ai/attestations/pip-witness-policy-verify/v0.1',
    'subject': [{'name': '${PACKAGE}', 'digest': {'gitoid': '${GITOID}'}}],
    'predicate': result
}
print(json.dumps(statement))
" > "/tmp/pip-witness-statement-${GITOID}.json"

    # Sign it with cilock
    docker run --rm \
        -v "/tmp:/data" \
        -v "${SIGNING_KEY}:/key.pem" \
        --entrypoint cilock \
        pip-witness:latest \
        sign \
        --signer-file-key-path /key.pem \
        --infile "/data/pip-witness-statement-${GITOID}.json" \
        --outfile "/data/pip-witness-policyverify-${GITOID}.json" \
        --datatype "https://aflock.ai/attestations/pip-witness-policy-verify/v0.1" \
        2>/dev/null

    if [ -f "${VERIFY_OUT}" ]; then
        # Upload to Archivista
        UPLOAD_RESULT=$(curl -s -X POST "${ARCHIVISTA_SERVER}/upload" \
            -H 'Content-Type: application/json' \
            -d @"${VERIFY_OUT}" 2>/dev/null)
        VERIFY_GITOID=$(echo "${UPLOAD_RESULT}" | python3 -c "import json,sys; print(json.load(sys.stdin).get('gitoid',''))" 2>/dev/null)

        if [ -n "${VERIFY_GITOID}" ]; then
            echo "  Verification attestation stored: ${VERIFY_GITOID}"
            echo "  View: ${ARCHIVISTA_SERVER}/download/${VERIFY_GITOID}"
        else
            echo "  Upload result: ${UPLOAD_RESULT}"
        fi
    fi
fi

echo ""
echo "  Scan attestation:   ${ARCHIVISTA_SERVER}/download/${GITOID}"
echo "  Policy file:        ${POLICY_FILE}"
echo "=========================================="
