#!/bin/bash
set -euo pipefail

# Build and sign a pip-witness policy from Rego modules.
# Creates a signed DSSE envelope that cilock verify can evaluate.
#
# Usage: ./build_policy.sh [signing-key-path] [output-path]

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_DIR="${SCRIPT_DIR}/.."
REGO_DIR="${SCRIPT_DIR}/rego"
SIGNING_KEY="${1:-${PROJECT_DIR}/policies/policy-key.pem}"
OUTPUT="${2:-${PROJECT_DIR}/policies/pip-witness-policy-signed.json}"

# Generate policy signing key if it doesn't exist
if [ ! -f "${SIGNING_KEY}" ]; then
    echo "Generating policy signing key..."
    openssl ecparam -genkey -name prime256v1 -noout -out "${SIGNING_KEY}"
    openssl ec -in "${SIGNING_KEY}" -pubout -out "${SIGNING_KEY%.pem}.pub"
    echo "  Key: ${SIGNING_KEY}"
    echo "  Pub: ${SIGNING_KEY%.pem}.pub"
fi

# Get public key ID (SHA256 of DER-encoded public key)
PUBKEY_FILE="${SIGNING_KEY%.pem}.pub"
if [ ! -f "${PUBKEY_FILE}" ]; then
    openssl ec -in "${SIGNING_KEY}" -pubout -out "${PUBKEY_FILE}" 2>/dev/null
fi
KEYID=$(openssl ec -pubin -in "${PUBKEY_FILE}" -outform DER 2>/dev/null | shasum -a 256 | cut -d' ' -f1)
PUBKEY_B64=$(base64 < "${PUBKEY_FILE}" | tr -d '\n')

echo "Building pip-witness policy..."
echo "  Key ID: ${KEYID}"

# Read each Rego module and base64-encode it
REGO_MODULES=""
for rego_file in "${REGO_DIR}"/*.rego; do
    name=$(basename "${rego_file}" .rego)
    module_b64=$(base64 < "${rego_file}" | tr -d '\n')
    echo "  Loaded: ${name}.rego"

    if [ -n "${REGO_MODULES}" ]; then
        REGO_MODULES="${REGO_MODULES},"
    fi
    REGO_MODULES="${REGO_MODULES}{\"name\": \"${name}\", \"module\": \"${module_b64}\"}"
done

# Build the policy JSON
# The policy has one step "pip-install" that expects:
# - command-run attestation (with Rego policies for network, execution, persistence)
# - pip-install attestation (with Rego policies for supply chain, signing)
POLICY_RAW="${SCRIPT_DIR}/pip-witness-policy-raw.json"

cat > "${POLICY_RAW}" << POLICY_EOF
{
    "expires": "2030-12-01T00:00:00Z",
    "steps": {
        "pip-install": {
            "name": "pip-install",
            "functionaries": [
                {
                    "type": "publickey",
                    "publickeyid": "${KEYID}"
                }
            ],
            "attestations": [
                {
                    "type": "https://aflock.ai/attestations/command-run/v0.1",
                    "regopolicies": [
$(for rego_file in "${REGO_DIR}"/network_exfiltration.rego "${REGO_DIR}"/code_execution.rego "${REGO_DIR}"/persistence.rego "${REGO_DIR}"/credential_harvesting.rego; do
    if [ -f "${rego_file}" ]; then
        name=$(basename "${rego_file}" .rego)
        module_b64=$(base64 < "${rego_file}" | tr -d '\n')
        echo "                        {\"name\": \"${name}\", \"module\": \"${module_b64}\"},"
    fi
done | sed '$ s/,$//')
                    ]
                },
                {
                    "type": "https://aflock.ai/attestations/pip-install/v0.1",
                    "regopolicies": [
$(for rego_file in "${REGO_DIR}"/supply_chain.rego "${REGO_DIR}"/package_signing.rego; do
    if [ -f "${rego_file}" ]; then
        name=$(basename "${rego_file}" .rego)
        module_b64=$(base64 < "${rego_file}" | tr -d '\n')
        echo "                        {\"name\": \"${name}\", \"module\": \"${module_b64}\"},"
    fi
done | sed '$ s/,$//')
                    ]
                },
                {
                    "type": "https://aflock.ai/attestations/environment/v0.1"
                },
                {
                    "type": "https://aflock.ai/attestations/material/v0.1"
                },
                {
                    "type": "https://aflock.ai/attestations/product/v0.1"
                }
            ]
        }
    },
    "publickeys": {
        "${KEYID}": {
            "keyid": "${KEYID}",
            "key": "${PUBKEY_B64}"
        }
    }
}
POLICY_EOF

echo "  Raw policy: ${POLICY_RAW}"

# Sign the policy with cilock
if command -v cilock &>/dev/null; then
    CILOCK=cilock
elif [ -f "/usr/local/bin/cilock" ]; then
    CILOCK=/usr/local/bin/cilock
else
    # Try to build it or use Docker
    echo "  cilock not found locally, using Docker..."
    docker run --rm \
        -v "${SCRIPT_DIR}:/policies" \
        --entrypoint cilock \
        pip-witness:latest \
        sign \
            --signer-file-key-path /policies/policy-key.pem \
            --infile "/policies/pip-witness-policy-raw.json" \
            --outfile "/policies/pip-witness-policy-signed.json" \
            --datatype "https://witness.testifysec.com/policy/v0.1"
    cp "${SCRIPT_DIR}/pip-witness-policy-signed.json" "${OUTPUT}" 2>/dev/null || true
    echo "  Signed policy: ${OUTPUT}"
    exit 0
fi

"${CILOCK}" sign \
    --signer-file-key-path "${SIGNING_KEY}" \
    --infile "${POLICY_RAW}" \
    --outfile "${OUTPUT}" \
    --datatype "https://witness.testifysec.com/policy/v0.1"

echo "  Signed policy: ${OUTPUT}"
echo ""
echo "Policy built successfully with ${#REGO_DIR}/*.rego modules."
echo "Use with: cilock verify --policy ${OUTPUT} --attestation <scan.json> --key ${PUBKEY_FILE}"