#!/usr/bin/env bash
set -euo pipefail

readonly statement_type=https://in-toto.io/Statement/v0.1
readonly provenance_type=https://slsa.dev/provenance/v1
readonly spdx_type=https://spdx.dev/Document

image_subject() {
  local image_ref=$1
  [[ $image_ref =~ ^(.+)@sha256:([0-9a-f]{64})$ ]] || {
    echo "Expected immutable sha256 image reference: $image_ref" >&2
    return 1
  }
  printf '%s\n%s\n' "${BASH_REMATCH[1]}" "${BASH_REMATCH[2]}"
}

create_provenance() {
  local output=$1 source_uri=$2 source_ref=$3 commit=$4 builder=$5 invocation=$6
  jq -n \
    --arg source_uri "$source_uri" \
    --arg source_ref "$source_ref" \
    --arg commit "$commit" \
    --arg builder "$builder" \
    --arg invocation "$invocation" \
    '{
      buildDefinition: {
        buildType: "https://github.com/Attestations/GitHubActionsWorkflow@v1",
        externalParameters: {source: {uri: $source_uri, ref: $source_ref}},
        internalParameters: {},
        resolvedDependencies: [{uri: $source_uri, digest: {gitCommit: $commit}}]
      },
      runDetails: {
        builder: {id: $builder},
        metadata: {invocationId: $invocation}
      }
    }' > "$output"
}

extract_sbom() (
  local output=$1 image_ref=$2 wrapper
  wrapper=$(mktemp)
  trap 'rm -f "$wrapper"' EXIT
  docker buildx imagetools inspect "$image_ref" --format '{{ json .SBOM }}' > "$wrapper"
  jq -e '
    .SPDX
    | .SPDXID == "SPDXRef-DOCUMENT"
      and (.spdxVersion | startswith("SPDX-"))
      and (.packages | type == "array")
  ' "$wrapper" > /dev/null
  jq '.SPDX' "$wrapper" > "$output"
)

verify_subject() {
  local statement=$1 subject_ref=$2 expected_name expected_digest
  mapfile -t subject < <(image_subject "$subject_ref")
  expected_name=${subject[0]}
  expected_digest=${subject[1]}
  jq -e \
    --arg statement_type "$statement_type" \
    --arg expected_name "$expected_name" \
    --arg expected_digest "$expected_digest" '
      ._type == $statement_type
      and (.subject | length == 1)
      and .subject[0].name == $expected_name
      and .subject[0].digest.sha256 == $expected_digest
    ' "$statement" > /dev/null
}

verify_statement() {
  local kind=$1 statement=$2 subject_ref=$3
  verify_subject "$statement" "$subject_ref"
  shift 3

  case $kind in
    provenance)
      local source_uri=$1 source_ref=$2 commit=$3 builder=$4 invocation=$5
      jq -e \
        --arg predicate_type "$provenance_type" \
        --arg source_uri "$source_uri" \
        --arg source_ref "$source_ref" \
        --arg commit "$commit" \
        --arg builder "$builder" \
        --arg invocation "$invocation" '
          .predicateType == $predicate_type
          and .predicate.buildDefinition.buildType == "https://github.com/Attestations/GitHubActionsWorkflow@v1"
          and .predicate.buildDefinition.externalParameters.source.uri == $source_uri
          and .predicate.buildDefinition.externalParameters.source.ref == $source_ref
          and (.predicate.buildDefinition.resolvedDependencies | length == 1)
          and .predicate.buildDefinition.resolvedDependencies[0].uri == $source_uri
          and .predicate.buildDefinition.resolvedDependencies[0].digest.gitCommit == $commit
          and .predicate.runDetails.builder.id == $builder
          and .predicate.runDetails.metadata.invocationId == $invocation
        ' "$statement" > /dev/null
      ;;
    sbom)
      jq -e \
        --arg predicate_type "$spdx_type" '
          .predicateType == $predicate_type
          and .predicate.SPDXID == "SPDXRef-DOCUMENT"
          and (.predicate.spdxVersion | startswith("SPDX-"))
          and (.predicate.packages | type == "array")
        ' "$statement" > /dev/null
      ;;
    *)
      echo "Unknown attestation kind: $kind" >&2
      return 1
      ;;
  esac
}

verify_attestation() (
  local kind=$1 lookup_ref=$2 subject_ref=$3 identity=$4 predicate_alias envelope statement payload
  shift 4
  case $kind in
    provenance) predicate_alias=slsaprovenance1 ;;
    sbom) predicate_alias=spdxjson ;;
    *) echo "Unknown attestation kind: $kind" >&2; return 1 ;;
  esac

  envelope=$(mktemp)
  statement=$(mktemp)
  trap 'rm -f "$envelope" "$statement"' EXIT
  cosign verify-attestation \
    --type "$predicate_alias" \
    --certificate-identity "$identity" \
    --certificate-oidc-issuer https://token.actions.githubusercontent.com \
    "$lookup_ref" > "$envelope"
  payload=$(jq -sr -e 'map(select(.payload?))[0].payload' "$envelope")
  printf '%s' "$payload" | base64 --decode > "$statement"
  verify_statement "$kind" "$statement" "$subject_ref" "$@"
)

attest_images() (
  local source_uri=$1 source_ref=$2 commit=$3 builder=$4 invocation=$5
  local image_ref sbom provenance temp_dir
  shift 5
  (($# > 0))
  temp_dir=$(mktemp -d)
  trap 'rm -rf "$temp_dir"' EXIT

  for image_ref in "$@"; do
    sbom="$temp_dir/sbom.json"
    provenance="$temp_dir/provenance.json"
    extract_sbom "$sbom" "$image_ref"
    create_provenance "$provenance" "$source_uri" "$source_ref" "$commit" "$builder" "$invocation"
    cosign attest --yes --type spdxjson --predicate "$sbom" "$image_ref"
    cosign attest --yes --type slsaprovenance1 --predicate "$provenance" "$image_ref"
  done
)

verify_image() {
  local lookup_ref=$1 subject_ref=$2 identity=$3 source_uri=$4 source_ref=$5 commit=$6 builder=$7 invocation=$8
  cosign verify \
    --certificate-identity "$identity" \
    --certificate-oidc-issuer https://token.actions.githubusercontent.com \
    "$lookup_ref"
  verify_attestation sbom "$lookup_ref" "$subject_ref" "$identity"
  verify_attestation provenance "$lookup_ref" "$subject_ref" "$identity" \
    "$source_uri" "$source_ref" "$commit" "$builder" "$invocation"
}

verify_images() {
  local identity=$1 source_uri=$2 source_ref=$3 commit=$4 builder=$5 invocation=$6 image_ref
  shift 6
  (($# > 0))
  for image_ref in "$@"; do
    verify_image "$image_ref" "$image_ref" "$identity" \
      "$source_uri" "$source_ref" "$commit" "$builder" "$invocation"
  done
}

verify_promoted() {
  local identity=$1 source_uri=$2 source_ref=$3 commit=$4 builder=$5 invocation=$6
  local entry subject_ref target expected actual
  shift 6
  (($# > 0))
  for entry in "$@"; do
    subject_ref=${entry%%|*}
    target=${entry#*|}
    expected=${subject_ref##*@}
    actual=$(crane digest "$target")
    [[ $actual == "$expected" ]] || {
      echo "Digest mismatch: $target resolved to $actual, expected $expected" >&2
      return 1
    }
    echo "$target -> $actual"
    verify_image "$target" "$subject_ref" "$identity" \
      "$source_uri" "$source_ref" "$commit" "$builder" "$invocation"
  done
}

case ${1:-} in
  create-provenance) shift; create_provenance "$@" ;;
  extract-sbom) shift; extract_sbom "$@" ;;
  attest-images) shift; attest_images "$@" ;;
  verify-statement) shift; verify_statement "$@" ;;
  verify) shift; verify_attestation "$@" ;;
  verify-images) shift; verify_images "$@" ;;
  verify-promoted) shift; verify_promoted "$@" ;;
  *)
    echo "Usage: $0 {create-provenance|extract-sbom|attest-images|verify-statement|verify|verify-images|verify-promoted} ..." >&2
    exit 2
    ;;
esac