Object Storage Integrity and Streaming Replica Parity Audit

id: object-storage-integrity-replica-parity-audit
namespace: company.storage

inputs:
  - id: primary_endpoint
    type: STRING
    defaults: ""
    description: "Custom endpoint URL for the primary storage (e.g.
      'http://minio:9000' for MinIO or LocalStack). Leave empty for standard AWS
      S3."

  - id: primary_bucket
    type: STRING
    defaults: "production-data-lake"
    description: "Name of the source S3-compatible bucket containing authoritative data."

  - id: primary_region
    type: STRING
    defaults: "us-east-1"
    description: "AWS region or deployment region for the primary bucket."

  - id: replica_endpoint
    type: STRING
    defaults: ""
    description: "Custom endpoint URL for the replica storage (e.g.
      'http://minio-replica:9000'). Leave empty for standard AWS S3."

  - id: replica_bucket
    type: STRING
    defaults: "disaster-recovery-replica"
    description: "Name of the replica S3-compatible bucket being verified for parity."

  - id: replica_region
    type: STRING
    defaults: "us-east-1"
    description: "AWS region or deployment region for the replica bucket."

  - id: prefix
    type: STRING
    defaults: "raw-events/"
    description: "Object key prefix to audit. Leave empty to audit the entire bucket."

  - id: max_keys
    type: INT
    defaults: 1000
    description: "Maximum number of object keys to retrieve and compare per audit run."

  - id: chunk_size_bytes
    type: INT
    defaults: 65536
    description: "Buffer size in bytes for streaming SHA-256 reads (default: 64 KB)."

  - id: enable_byte_level_diff
    type: BOOL
    defaults: true
    description: "When a hash mismatch occurs, perform byte-level stream inspection
      to identify the exact differing offset."

  - id: enable_quarantine
    type: BOOL
    defaults: false
    description: "When true, isolates corrupted replica objects to the quarantine
      prefix. Safe report-only mode is default."

  - id: quarantine_prefix
    type: STRING
    defaults: "quarantine-isolated/"
    description: "Destination prefix within the replica bucket where corrupted
      objects are moved."

  - id: fail_on_discrepancy
    type: BOOL
    defaults: true
    description: "Whether to fail the workflow execution when any mismatch, missing
      file, or corruption is found."

triggers:
  - id: daily_scheduled_audit
    type: io.kestra.plugin.core.trigger.Schedule
    cron: "0 2 * * *"

  - id: replication_sync_webhook
    type: io.kestra.plugin.core.trigger.Webhook
    key: "storage-parity-sync-hook"

tasks:
  - id: init_audit_context
    type: io.kestra.plugin.core.log.Log
    level: INFO
    message: |
      [STORAGE PARITY AUDIT] Starting integrity verification
      Primary Target: [{{ inputs.primary_bucket }}] @ {{ inputs.primary_endpoint != "" ? inputs.primary_endpoint : "AWS Default" }}
      Replica Target: [{{ inputs.replica_bucket }}] @ {{ inputs.replica_endpoint != "" ? inputs.replica_endpoint : "AWS Default" }}
      Auditing Prefix: "{{ inputs.prefix }}" | Max Keys: {{ inputs.max_keys }} | Quarantine Active: {{ inputs.enable_quarantine }}

  - id: parallel_inventory_extraction
    type: io.kestra.plugin.core.flow.Parallel
    description: "Concurrently extract object inventories from both storage locations"
    tasks:
      - id: list_primary_inventory
        type: io.kestra.plugin.aws.s3.List
        description: "Retrieve object listing from primary bucket"
        endpointOverride: "{{ inputs.primary_endpoint != '' ? inputs.primary_endpoint : null }}"
        bucket: "{{ inputs.primary_bucket }}"
        prefix: "{{ inputs.prefix }}"
        region: "{{ inputs.primary_region }}"
        compatibilityMode: true
        forcePathStyle: true
        accessKeyId: "{{ secret('PRIMARY_S3_ACCESS_KEY_ID') }}"
        secretKeyId: "{{ secret('PRIMARY_S3_SECRET_ACCESS_KEY') }}"
        filter: FILES
        maxKeys: "{{ inputs.max_keys }}"

      - id: list_replica_inventory
        type: io.kestra.plugin.aws.s3.List
        description: "Retrieve object listing from replica bucket"
        endpointOverride: "{{ inputs.replica_endpoint != '' ? inputs.replica_endpoint : null }}"
        bucket: "{{ inputs.replica_bucket }}"
        prefix: "{{ inputs.prefix }}"
        region: "{{ inputs.replica_region }}"
        compatibilityMode: true
        forcePathStyle: true
        accessKeyId: "{{ secret('REPLICA_S3_ACCESS_KEY_ID') }}"
        secretKeyId: "{{ secret('REPLICA_S3_SECRET_ACCESS_KEY') }}"
        filter: FILES
        maxKeys: "{{ inputs.max_keys }}"

  - id: reconcile_and_verify_parity
    type: io.kestra.plugin.scripts.python.Script
    description: "Reconcile object sets, perform streaming SHA-256 validation, and
      pinpoint byte offsets"
    taskRunner:
      type: io.kestra.plugin.core.runner.Process
    beforeCommands:
      - pip install boto3 botocore kestra
    outputFiles:
      - "integrity_audit_report.json"
    env:
      PRIMARY_ENDPOINT: "{{ inputs.primary_endpoint }}"
      PRIMARY_BUCKET: "{{ inputs.primary_bucket }}"
      PRIMARY_REGION: "{{ inputs.primary_region }}"
      PRIMARY_KEY: "{{ secret('PRIMARY_S3_ACCESS_KEY_ID') }}"
      PRIMARY_SECRET: "{{ secret('PRIMARY_S3_SECRET_ACCESS_KEY') }}"
      REPLICA_ENDPOINT: "{{ inputs.replica_endpoint }}"
      REPLICA_BUCKET: "{{ inputs.replica_bucket }}"
      REPLICA_REGION: "{{ inputs.replica_region }}"
      REPLICA_KEY: "{{ secret('REPLICA_S3_ACCESS_KEY_ID') }}"
      REPLICA_SECRET: "{{ secret('REPLICA_S3_SECRET_ACCESS_KEY') }}"
      CHUNK_SIZE: "{{ inputs.chunk_size_bytes }}"
      ENABLE_DIFF: "{{ inputs.enable_byte_level_diff }}"
      ENABLE_QUARANTINE: "{{ inputs.enable_quarantine }}"
      QUARANTINE_PREFIX: "{{ inputs.quarantine_prefix }}"
    script: |
      import os
      import json
      import hashlib
      import datetime
      import boto3
      from botocore.config import Config
      from botocore.exceptions import ClientError, BotoCoreError
      from kestra import Kestra

      # ---------------------------------------------------------------------------
      # 1. Configuration & Client Initialization
      # ---------------------------------------------------------------------------
      chunk_size = int(os.environ.get("CHUNK_SIZE", 65536))
      enable_diff = os.environ.get("ENABLE_DIFF", "true").lower() == "true"
      enable_quarantine = os.environ.get("ENABLE_QUARANTINE", "false").lower() == "true"
      quarantine_prefix = os.environ.get("QUARANTINE_PREFIX", "quarantine-isolated/")
      if not quarantine_prefix.endswith("/"):
          quarantine_prefix += "/"

      def get_s3_client(endpoint, region, key, secret):
          kwargs = {
              "region_name": region or "us-east-1",
              "aws_access_key_id": key,
              "aws_secret_access_key": secret,
              "config": Config(
                  signature_version="s3v4",
                  s3={"addressing_style": "path"},
                  retries={"max_attempts": 3, "mode": "standard"},
                  connect_timeout=10,
                  read_timeout=30
              )
          }
          if endpoint and endpoint.strip():
              kwargs["endpoint_url"] = endpoint.strip()
          return boto3.client("s3", **kwargs)

      primary_client = get_s3_client(
          os.environ.get("PRIMARY_ENDPOINT"),
          os.environ.get("PRIMARY_REGION"),
          os.environ.get("PRIMARY_KEY"),
          os.environ.get("PRIMARY_SECRET")
      )
      replica_client = get_s3_client(
          os.environ.get("REPLICA_ENDPOINT"),
          os.environ.get("REPLICA_REGION"),
          os.environ.get("REPLICA_KEY"),
          os.environ.get("REPLICA_SECRET")
      )

      primary_bucket = os.environ.get("PRIMARY_BUCKET")
      replica_bucket = os.environ.get("REPLICA_BUCKET")

      # ---------------------------------------------------------------------------
      # 2. Ingest Inventory from Upstream Tasks
      # ---------------------------------------------------------------------------
      primary_objects_raw = {{ (outputs.list_primary_inventory.objects ?? []) | toJson }}
      replica_objects_raw = {{ (outputs.list_replica_inventory.objects ?? []) | toJson }}

      primary_map = {item["key"]: item for item in primary_objects_raw if "key" in item}
      replica_map = {item["key"]: item for item in replica_objects_raw if "key" in item}

      all_primary_keys = set(primary_map.keys())
      all_replica_keys = set(replica_map.keys())

      missing_in_replica = sorted(list(all_primary_keys - all_replica_keys))
      unexpected_in_replica = sorted(list(all_replica_keys - all_primary_keys))
      common_keys = sorted(list(all_primary_keys & all_replica_keys))

      size_mismatches = []
      content_mismatches = []
      verified_matches = []
      read_errors = []
      quarantined_items = []
      total_bytes_streamed = 0

      # ---------------------------------------------------------------------------
      # 3. Streaming Content Verification Loop
      # ---------------------------------------------------------------------------
      for key in common_keys:
          p_meta = primary_map[key]
          r_meta = replica_map[key]
          p_size = p_meta.get("size", 0)
          r_size = r_meta.get("size", 0)

          if p_size != r_size:
              size_mismatches.append({
                  "key": key,
                  "primary_size_bytes": p_size,
                  "replica_size_bytes": r_size,
                  "size_delta_bytes": abs(p_size - r_size),
                  "primary_etag": p_meta.get("eTag", ""),
                  "replica_etag": r_meta.get("eTag", "")
              })
              continue

          # Bounded streaming SHA-256 calculation
          p_stream = None
          r_stream = None
          try:
              p_resp = primary_client.get_object(Bucket=primary_bucket, Key=key)
              r_resp = replica_client.get_object(Bucket=replica_bucket, Key=key)
              p_stream = p_resp["Body"]
              r_stream = r_resp["Body"]

              hasher_p = hashlib.sha256()
              hasher_r = hashlib.sha256()

              first_differing_offset = None
              primary_differing_byte = None
              replica_differing_byte = None
              current_offset = 0

              while True:
                  chunk_p = p_stream.read(chunk_size)
                  chunk_r = r_stream.read(chunk_size)

                  if not chunk_p and not chunk_r:
                      break

                  total_bytes_streamed += len(chunk_p) + len(chunk_r)
                  hasher_p.update(chunk_p)
                  hasher_r.update(chunk_r)

                  # Byte-level diffing when chunks diverge
                  if enable_diff and first_differing_offset is None and chunk_p != chunk_r:
                      min_len = min(len(chunk_p), len(chunk_r))
                      for idx in range(min_len):
                          if chunk_p[idx] != chunk_r[idx]:
                              first_differing_offset = current_offset + idx
                              primary_differing_byte = f"0x{chunk_p[idx]:02x}"
                              replica_differing_byte = f"0x{chunk_r[idx]:02x}"
                              break
                      if first_differing_offset is None:
                          first_differing_offset = current_offset + min_len
                          primary_differing_byte = f"0x{chunk_p[min_len]:02x}" if len(chunk_p) > min_len else "EOF"
                          replica_differing_byte = f"0x{chunk_r[min_len]:02x}" if len(chunk_r) > min_len else "EOF"

                  current_offset += max(len(chunk_p), len(chunk_r))

              digest_p = hasher_p.hexdigest()
              digest_r = hasher_r.hexdigest()

              if digest_p == digest_r:
                  verified_matches.append({
                      "key": key,
                      "size_bytes": p_size,
                      "sha256": digest_p
                  })
              else:
                  mismatch_entry = {
                      "key": key,
                      "size_bytes": p_size,
                      "primary_sha256": digest_p,
                      "replica_sha256": digest_r,
                      "primary_etag": p_meta.get("eTag", ""),
                      "replica_etag": r_meta.get("eTag", ""),
                      "first_differing_offset": first_differing_offset,
                      "primary_byte": primary_differing_byte,
                      "replica_byte": replica_differing_byte
                  }
                  content_mismatches.append(mismatch_entry)

                  # Safe opt-in quarantine for corrupted replica
                  if enable_quarantine:
                      timestamp_tag = datetime.datetime.now(datetime.timezone.utc).strftime("%Y%m%dT%H%M%SZ")
                      clean_key_name = key.replace("/", "_")
                      quarantine_key = f"{quarantine_prefix}corrupted_{timestamp_tag}_{clean_key_name}"
                      try:
                          replica_client.copy_object(
                              Bucket=replica_bucket,
                              CopySource={"Bucket": replica_bucket, "Key": key},
                              Key=quarantine_key
                          )
                          replica_client.delete_object(Bucket=replica_bucket, Key=key)
                          quarantined_items.append({
                              "original_key": key,
                              "quarantine_key": quarantine_key,
                              "reason": "SHA256_MISMATCH",
                              "timestamp": timestamp_tag
                          })
                      except Exception as q_err:
                          read_errors.append({
                              "key": key,
                              "target": "replica_quarantine",
                              "error": str(q_err)
                          })

          except (ClientError, BotoCoreError) as err:
              read_errors.append({
                  "key": key,
                  "error_type": type(err).__name__,
                  "error_message": str(err)
              })
          finally:
              if p_stream:
                  p_stream.close()
              if r_stream:
                  r_stream.close()

      # ---------------------------------------------------------------------------
      # 4. Synthesize Audit Status and Export Manifest
      # ---------------------------------------------------------------------------
      total_discrepancies = (
          len(missing_in_replica) +
          len(unexpected_in_replica) +
          len(size_mismatches) +
          len(content_mismatches) +
          len(read_errors)
      )
      audit_passed = (total_discrepancies == 0)

      report_manifest = {
          "audit_timestamp": datetime.datetime.now(datetime.timezone.utc).isoformat(),
          "primary_bucket": primary_bucket,
          "replica_bucket": replica_bucket,
          "audit_passed": audit_passed,
          "summary": {
              "total_primary_objects": len(all_primary_keys),
              "total_replica_objects": len(all_replica_keys),
              "verified_matches_count": len(verified_matches),
              "missing_in_replica_count": len(missing_in_replica),
              "unexpected_in_replica_count": len(unexpected_in_replica),
              "size_mismatches_count": len(size_mismatches),
              "content_mismatches_count": len(content_mismatches),
              "read_errors_count": len(read_errors),
              "quarantined_count": len(quarantined_items),
              "total_bytes_streamed": total_bytes_streamed
          },
          "discrepancies": {
              "missing_in_replica": missing_in_replica,
              "unexpected_in_replica": unexpected_in_replica,
              "size_mismatches": size_mismatches,
              "content_mismatches": content_mismatches,
              "read_errors": read_errors
          },
          "quarantine_log": quarantined_items,
          "verified_matches": verified_matches
      }

      with open("integrity_audit_report.json", "w", encoding="utf-8") as f:
          json.dump(report_manifest, f, indent=2)

      print(f"[AUDIT COMPLETE] Passed: {audit_passed} | Verified: {len(verified_matches)} | Discrepancies: {total_discrepancies}")

      Kestra.outputs({
          "audit_passed": audit_passed,
          "discrepancy_count": total_discrepancies,
          "verified_keys_count": len(verified_matches),
          "missing_keys_count": len(missing_in_replica),
          "unexpected_keys_count": len(unexpected_in_replica),
          "size_mismatch_count": len(size_mismatches),
          "content_mismatch_count": len(content_mismatches),
          "read_error_count": len(read_errors),
          "quarantined_keys_count": len(quarantined_items),
          "total_bytes_streamed": total_bytes_streamed
      })

  - id: route_audit_outcome
    type: io.kestra.plugin.core.flow.If
    description: "Route workflow execution based on parity audit results"
    condition: "{{ outputs.reconcile_and_verify_parity.vars.audit_passed == true }}"
    then:
      - id: log_audit_attestation
        type: io.kestra.plugin.core.log.Log
        level: INFO
        message: |
          [PARITY ATTESTATION PASSED] Replica [{{ inputs.replica_bucket }}] matches primary [{{ inputs.primary_bucket }}].
          Verified Objects: {{ outputs.reconcile_and_verify_parity.vars.verified_keys_count }}
          Total Data Streamed: {{ outputs.reconcile_and_verify_parity.vars.total_bytes_streamed }} bytes.
          Integrity Report generated: integrity_audit_report.json

    else:
      - id: log_discrepancy_summary
        type: io.kestra.plugin.core.log.Log
        level: WARN
        message: |
          [PARITY DISCREPANCY DETECTED] Total Discrepancies: {{ outputs.reconcile_and_verify_parity.vars.discrepancy_count }}
          - Missing in Replica: {{ outputs.reconcile_and_verify_parity.vars.missing_keys_count }}
          - Unexpected in Replica: {{ outputs.reconcile_and_verify_parity.vars.unexpected_keys_count }}
          - Size Mismatches: {{ outputs.reconcile_and_verify_parity.vars.size_mismatch_count }}
          - Content Bitrot / SHA256 Mismatches: {{ outputs.reconcile_and_verify_parity.vars.content_mismatch_count }}
          - Read/Permission Errors: {{ outputs.reconcile_and_verify_parity.vars.read_error_count }}
          - Quarantined Objects: {{ outputs.reconcile_and_verify_parity.vars.quarantined_keys_count }}

      - id: enforce_pipeline_gate
        type: io.kestra.plugin.core.execution.Fail
        condition: "{{ inputs.fail_on_discrepancy == true }}"
        errorMessage: |
          Object storage replica parity audit detected {{ outputs.reconcile_and_verify_parity.vars.discrepancy_count }} discrepancies between [{{ inputs.primary_bucket }}] and [{{ inputs.replica_bucket }}]. Downstream pipeline blocked. Check integrity_audit_report.json.