From 8bc40f353dcfcde2cd701088bd7a83e8a212bc98 Mon Sep 17 00:00:00 2001
From: Chris Hager <chris@linuxuser.at>
Date: Fri, 15 Nov 2024 21:52:21 +0100
Subject: [PATCH] CLI tool to print verified measurements from an aTLS server

---
 cmd/get-measurements/main.go | 121 +++++++++++++++++++++++++++++++++++
 measurements-empty.json      |   1 +
 2 files changed, 122 insertions(+)
 create mode 100644 cmd/get-measurements/main.go
 create mode 100644 measurements-empty.json

diff --git a/cmd/get-measurements/main.go b/cmd/get-measurements/main.go
new file mode 100644
index 0000000..32060c5
--- /dev/null
+++ b/cmd/get-measurements/main.go
@@ -0,0 +1,121 @@
+package main
+
+//
+// CLI tool to get and print verified measurements from an aTLS server.
+//
+// Currently only works for Azure TDX but should be easy to expand.
+//
+// Usage:
+//
+//	 go run cmd/get-measurements/main.go instance_ip:port
+//
+
+import (
+	"crypto/tls"
+	"encoding/asn1"
+	"encoding/hex"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"log"
+	"os"
+
+	"github.com/flashbots/cvm-reverse-proxy/common"
+	"github.com/flashbots/cvm-reverse-proxy/internal/atls"
+	"github.com/flashbots/cvm-reverse-proxy/internal/attestation/variant"
+	"github.com/flashbots/cvm-reverse-proxy/proxy"
+	"github.com/urfave/cli/v2" // imports as package "cli"
+)
+
+var flags []cli.Flag = []cli.Flag{
+	&cli.BoolFlag{
+		Name:  "log-debug",
+		Value: false,
+		Usage: "log debug messages",
+	},
+}
+
+func main() {
+	app := &cli.App{
+		Name:   "get-measurements",
+		Usage:  "Get verified measurements",
+		Flags:  flags,
+		Action: runClient,
+	}
+
+	if err := app.Run(os.Args); err != nil {
+		log.Fatal(err)
+	}
+}
+
+func runClient(cCtx *cli.Context) error {
+	logDebug := cCtx.Bool("log-debug")
+
+	// Setup logging
+	log := common.SetupLogger(&common.LoggingOpts{
+		Debug:   logDebug,
+		JSON:    false,
+		Service: "get-measurements",
+		Version: common.Version,
+	})
+
+	addr := cCtx.Args().Get(0)
+	if addr == "" {
+		log.Error("Please provide an address as cli argument")
+		return errors.New("provide an address as argument")
+	}
+
+	log.Info("Getting verified measurements from " + addr + " ...")
+
+	// Prepare aTLS stuff
+	serverAttestationType := proxy.AttestationAzureTDX
+	issuer, err := proxy.CreateAttestationIssuer(log, serverAttestationType)
+	if err != nil {
+		log.Error("could not create attestation issuer", "err", err)
+		return err
+	}
+
+	validators, err := proxy.CreateAttestationValidators(log, serverAttestationType, "measurements-empty.json")
+	if err != nil {
+		log.Error("could not create attestation validators", "err", err)
+		return err
+	}
+
+	tlsConfig, err := atls.CreateAttestationClientTLSConfig(issuer, validators)
+	if err != nil {
+		log.Error("could not create atls config", "err", err)
+		return err
+	}
+
+	// Open connection to the TDX server and verify the aTLS attestation
+	conn, err := tls.Dial("tcp", addr, tlsConfig)
+	if err != nil {
+		log.Error("Error in Dial", "err", err)
+		return err
+	}
+	defer conn.Close()
+
+	// Extract the aTLS variant and measurements from the TLS connection
+	certs := conn.ConnectionState().PeerCertificates
+	atlsVariant, extractedMeasurements, err := proxy.GetMeasurementsFromTLS(certs, []asn1.ObjectIdentifier{variant.AzureTDX{}.OID()})
+	if err != nil {
+		log.Error("Error in getMeasurementsFromTLS", "err", err)
+		return err
+	}
+
+	measurementsInHeaderFormat := make(map[uint32]string, len(extractedMeasurements))
+	for pcr, value := range extractedMeasurements {
+		measurementsInHeaderFormat[pcr] = hex.EncodeToString(value)
+	}
+
+	marshaledPcrs, err := json.MarshalIndent(measurementsInHeaderFormat, "", "    ")
+	if err != nil {
+		return errors.New("could not marshal measurement extracted from tls extension")
+	}
+
+	log.Info("Variant: " + atlsVariant.String())
+	// log.Info("Measurements", "measurements", string(marshaledPcrs))
+	fmt.Println(string(marshaledPcrs))
+
+	return nil
+}
diff --git a/measurements-empty.json b/measurements-empty.json
new file mode 100644
index 0000000..f9eea9d
--- /dev/null
+++ b/measurements-empty.json
@@ -0,0 +1 @@
+{"azure-tdx": {}}