Merkle Audit Chain

registerFunction(
  { id: "security::audit" },
  async ({ type, agentId, detail }) => {
    // Get the latest entry to retrieve previous hash
    const prev = await trigger("state::get", {
      scope: "audit",
      key: "__latest",
    }).catch(() => ({ hash: "0".repeat(64) }));

    const entry: AuditEntry = {
      id: crypto.randomUUID(),
      timestamp: Date.now(),
      type,
      agentId,
      detail: detail || {},
      prevHash: prev.hash,
      hash: "",  // Computed below
    };

    // Compute HMAC-SHA256 hash linking to previous entry
    entry.hash = createHash("sha256")
      .update(JSON.stringify({ ...entry, hash: undefined }) + prev.hash)
      .digest("hex");

    // Store the entry
    await trigger("state::set", {
      scope: "audit",
      key: entry.id,
      value: entry,
    });

    // Update latest pointer
    await trigger("state::set", {
      scope: "audit",
      key: "__latest",
      value: { hash: entry.hash, id: entry.id, timestamp: entry.timestamp },
    });

    return { id: entry.id, hash: entry.hash };
  }
);

use hmac::{Hmac, Mac};
use sha2::Sha256;
type HmacSha256 = Hmac<Sha256>;

fn audit_hmac_key() -> &'static [u8] {
    static KEY: OnceLock<Vec<u8>> = OnceLock::new();
    KEY.get_or_init(|| {
        std::env::var("AUDIT_HMAC_KEY")
            .unwrap_or_else(|_| "dev-default-hmac-key-change-in-prod".to_string())
            .into_bytes()
    })
}

async fn append_audit(iii: &III, input: Value) -> Result<Value, IIIError> {
    let prev: Value = iii
        .trigger("state::get", json!({ "scope": "audit", "key": "__latest" }))
        .await
        .unwrap_or(json!({ "hash": "0".repeat(64) }));

    let prev_hash = prev["hash"].as_str().unwrap_or(&"0".repeat(64)).to_string();
    let id = uuid::Uuid::new_v4().to_string();
    let timestamp = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or_default()
        .as_millis() as u64;

    let entry_data = json!({
        "id": &id,
        "timestamp": timestamp,
        "type": input.get("type"),
        "agentId": input.get("agentId"),
        "detail": input.get("detail").unwrap_or(&json!({})),
        "prevHash": &prev_hash,
    });

    // Compute HMAC-SHA256
    let mut mac = HmacSha256::new_from_slice(audit_hmac_key())
        .map_err(|e| IIIError::Handler(format!("HMAC key error: {}", e)))?;
    mac.update(entry_data.to_string().as_bytes());
    mac.update(prev_hash.as_bytes());
    let hash = hex::encode(mac.finalize().into_bytes());

    let full_entry = json!({
        "id": &id,
        "timestamp": timestamp,
        "type": input.get("type"),
        "agentId": input.get("agentId"),
        "detail": input.get("detail").unwrap_or(&json!({})),
        "hash": &hash,
        "prevHash": &prev_hash,
    });

    iii.trigger("state::set", json!({
        "scope": "audit",
        "key": &id,
        "value": &full_entry,
    })).await?;

    iii.trigger("state::set", json!({
        "scope": "audit",
        "key": "__latest",
        "value": { "hash": &hash, "id": &id, "timestamp": timestamp },
    })).await?;

    Ok(json!({ "id": id, "hash": hash }))
}

registerFunction(
  { id: "security::verify_audit" },
  async (req) => {
    const entries = await trigger("state::list", { scope: "audit" });
    const chain: AuditEntry[] = entries
      .filter((e) => e.key !== "__latest" && e.value?.hash)
      .map((e) => e.value)
      .sort((a, b) => a.timestamp - b.timestamp);

    let prevHash = "0".repeat(64);
    const violations: string[] = [];

    for (const entry of chain) {
      // Check chain linkage
      if (entry.prevHash !== prevHash) {
        violations.push(
          `Chain break at ${entry.id}: expected ${prevHash}, got ${entry.prevHash}`
        );
      }

      // Recompute hash
      const computed = createHash("sha256")
        .update(JSON.stringify({ ...entry, hash: undefined }) + entry.prevHash)
        .digest("hex");

      // Verify integrity
      if (computed !== entry.hash) {
        violations.push(`Tampered entry ${entry.id}: hash mismatch`);
      }

      prevHash = entry.hash;
    }

    return {
      valid: violations.length === 0,
      entries: chain.length,
      violations,
    };
  }
);

async fn verify_audit(iii: &III) -> Result<Value, IIIError> {
    let entries: Value = iii
        .trigger("state::list", json!({ "scope": "audit" }))
        .await?;

    let mut chain: Vec<AuditEntry> = entries
        .as_array()
        .unwrap_or(&vec![])
        .iter()
        .filter_map(|e| {
            let val = e.get("value")?;
            if e["key"].as_str() == Some("__latest") {
                return None;
            }
            serde_json::from_value(val.clone()).ok()
        })
        .collect();

    chain.sort_by_key(|e| e.timestamp);

    let zeros = "0".repeat(64);
    let mut prev_hash = zeros.as_str();
    let mut violations = Vec::new();

    for entry in &chain {
        // Check linkage
        if entry.prev_hash != prev_hash {
            violations.push(format!(
                "Chain break at {}: expected {}, got {}",
                entry.id, prev_hash, entry.prev_hash
            ));
        }

        // Recompute HMAC
        let check_data = json!({
            "id": &entry.id,
            "timestamp": entry.timestamp,
            "type": &entry.entry_type,
            "agentId": &entry.agent_id,
            "detail": &entry.detail,
            "prevHash": &entry.prev_hash,
        });

        let mut mac = HmacSha256::new_from_slice(audit_hmac_key())
            .map_err(|_| IIIError::Handler("HMAC key error".into()))?;
        mac.update(check_data.to_string().as_bytes());
        mac.update(entry.prev_hash.as_bytes());
        let computed = hex::encode(mac.finalize().into_bytes());

        // Detect tampering
        if computed != entry.hash {
            violations.push(format!("Tampered entry {}: hash mismatch", entry.id));
        }

        prev_hash = &entry.hash;
    }

    Ok(json!({
        "valid": violations.is_empty(),
        "entries": chain.len(),
        "violations": violations,
    }))
}