llm-intelligence/scripts/vertex_pricing_signature_guard_lib.go

//go:build llm_script

package main

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"time"
)

type vertexPricingSignatureGuardConfig struct {
	URL            string
	Fixture        string
	SnapshotDir    string
	BaselinePath   string
	Timeout        time.Duration
	AllowBootstrap bool
}

type vertexPricingSignatureGuardResult struct {
	SnapshotPath         string
	SignaturePath        string
	BaselinePath         string
	DriftDetected        bool
	BaselineInitialized  bool
	PreviousBaselineHash string
	CurrentSignature     vertexPricingStructureSignature
}

func runVertexPricingSignatureGuard(cfg vertexPricingSignatureGuardConfig, now time.Time) (vertexPricingSignatureGuardResult, error) {
	snapshotDir := cfg.SnapshotDir
	if snapshotDir == "" {
		snapshotDir = filepath.Join("logs", "vertex-pricing-snapshots")
	}
	if err := os.MkdirAll(snapshotDir, 0o755); err != nil {
		return vertexPricingSignatureGuardResult{}, fmt.Errorf("mkdir snapshot dir: %w", err)
	}

	baseName := fmt.Sprintf("vertex-pricing-%s", now.Format("20060102-150405"))
	snapshotPath := filepath.Join(snapshotDir, baseName+".html")
	signaturePath := filepath.Join(snapshotDir, baseName+".signature.json")
	baselinePath := cfg.BaselinePath
	if baselinePath == "" {
		baselinePath = filepath.Join(snapshotDir, "baseline.signature.json")
	}

	clientCfg := vertexPricingImportConfig{
		URL:          cfg.URL,
		Fixture:      cfg.Fixture,
		DryRun:       true,
		Timeout:      cfg.Timeout,
		SnapshotOnly: true,
		SnapshotOut:  snapshotPath,
		SignatureOut: signaturePath,
	}
	if err := runVertexPricingImport(clientCfg, nil, ioDiscard{}); err != nil {
		return vertexPricingSignatureGuardResult{}, err
	}

	current, err := readVertexPricingStructureSignature(signaturePath)
	if err != nil {
		return vertexPricingSignatureGuardResult{}, err
	}
	result := vertexPricingSignatureGuardResult{
		SnapshotPath:     snapshotPath,
		SignaturePath:    signaturePath,
		BaselinePath:     baselinePath,
		CurrentSignature: current,
	}

	previous, err := readVertexPricingStructureSignature(baselinePath)
	if err != nil {
		if os.IsNotExist(err) {
			if !cfg.AllowBootstrap {
				return result, fmt.Errorf("vertex pricing baseline missing: %s", baselinePath)
			}
			if err := copyFileCommon(signaturePath, baselinePath); err != nil {
				return result, fmt.Errorf("initialize baseline: %w", err)
			}
			result.BaselineInitialized = true
			return result, nil
		}
		return result, err
	}

	result.PreviousBaselineHash = previous.StructureSHA256
	if previous.StructureSHA256 != current.StructureSHA256 {
		result.DriftDetected = true
		return result, fmt.Errorf(
			"vertex pricing structure drift detected: baseline=%s current=%s baseline_path=%s signature_path=%s snapshot_path=%s",
			previous.StructureSHA256, current.StructureSHA256, baselinePath, signaturePath, snapshotPath,
		)
	}
	return result, nil
}

func formatVertexPricingSignatureGuardSummary(result vertexPricingSignatureGuardResult) string {
	return fmt.Sprintf(
		"source=vertex-pricing-signature-guard drift=%t baseline_initialized=%t structure_sha256=%s previous_baseline_sha256=%s snapshot_out=%s signature_out=%s baseline_path=%s",
		result.DriftDetected,
		result.BaselineInitialized,
		result.CurrentSignature.StructureSHA256,
		emptyIfBlank(result.PreviousBaselineHash),
		result.SnapshotPath,
		result.SignaturePath,
		result.BaselinePath,
	)
}

func buildVertexPricingSignatureAuditRecord(cfg vertexPricingSignatureGuardConfig, result vertexPricingSignatureGuardResult, checkedAt time.Time, runErr error) officialImportSignatureAuditRecord {
	record := officialImportSignatureAuditRecord{
		SourceKey:               "vertex_pricing_signature",
		CheckedAt:               checkedAt,
		Status:                  officialImportSignatureAuditStatus(result.DriftDetected, result.BaselineInitialized, runErr),
		DriftDetected:           result.DriftDetected,
		BaselineInitialized:     result.BaselineInitialized,
		SourceURL:               strings.TrimSpace(cfg.URL),
		FixturePath:             strings.TrimSpace(cfg.Fixture),
		SnapshotPath:            strings.TrimSpace(result.SnapshotPath),
		SignaturePath:           strings.TrimSpace(result.SignaturePath),
		BaselinePath:            strings.TrimSpace(result.BaselinePath),
		StructureSHA256:         strings.TrimSpace(result.CurrentSignature.StructureSHA256),
		PreviousStructureSHA256: strings.TrimSpace(result.PreviousBaselineHash),
		ByteSize:                result.CurrentSignature.ByteSize,
		ErrorMessage:            errorMessageText(runErr),
	}
	if hasVertexPricingStructureSignature(result.CurrentSignature) {
		signatureCopy := result.CurrentSignature
		record.SignaturePayload = &signatureCopy
	}
	return record
}

func persistVertexPricingSignatureAuditIfConfigured(cfg vertexPricingSignatureGuardConfig, result vertexPricingSignatureGuardResult, checkedAt time.Time, runErr error) error {
	return persistOfficialImportSignatureAuditIfConfigured(buildVertexPricingSignatureAuditRecord(cfg, result, checkedAt, runErr))
}

func readVertexPricingStructureSignature(path string) (vertexPricingStructureSignature, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return vertexPricingStructureSignature{}, err
	}
	var signature vertexPricingStructureSignature
	if err := json.Unmarshal(data, &signature); err != nil {
		return vertexPricingStructureSignature{}, fmt.Errorf("unmarshal signature %s: %w", path, err)
	}
	return signature, nil
}

func hasVertexPricingStructureSignature(signature vertexPricingStructureSignature) bool {
	return signature.ByteSize > 0 ||
		strings.TrimSpace(signature.StructureSHA256) != "" ||
		strings.TrimSpace(signature.SHA256) != "" ||
		len(signature.TagCounts) > 0 ||
		len(signature.Headings) > 0
}