From b37f37e486e57936b6565b74338969c4e2e1d829 Mon Sep 17 00:00:00 2001
From: Cole <cole@testifysec.com>
Date: Sun, 10 Nov 2024 12:38:16 -0600
Subject: [PATCH] add user defined options for sign

---
 cmd/run.go      |  3 ++-
 cmd/sign.go     | 26 +++++++++++++++++++++++++-
 options/sign.go |  3 ++-
 3 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/cmd/run.go b/cmd/run.go
index 61910cff..5174fd26 100644
--- a/cmd/run.go
+++ b/cmd/run.go
@@ -117,7 +117,8 @@ func runRun(ctx context.Context, ro options.RunOptions, args []string, signers .
 	runOptions := []witness.RunOption{
 		witness.RunWithSigners(signers...),
 		witness.RunWithAttestors(attestors),
-		witness.RunWithAttestationOpts(attestation.WithWorkingDir(ro.WorkingDir), attestation.WithHashes(roHashes)),
+		witness.RunWithAttestationOpts(attestation.WithWorkingDir(ro.WorkingDir),
+			attestation.WithHashes(roHashes)),
 		witness.RunWithTimestampers(timestampers...),
 	}
 
diff --git a/cmd/sign.go b/cmd/sign.go
index 3aad6913..2d317173 100644
--- a/cmd/sign.go
+++ b/cmd/sign.go
@@ -16,12 +16,14 @@ package cmd
 
 import (
 	"context"
+	"crypto"
 	"fmt"
 	"os"
 
 	witness "github.com/in-toto/go-witness"
 	"github.com/in-toto/go-witness/cryptoutil"
 	"github.com/in-toto/go-witness/dsse"
+	"github.com/in-toto/go-witness/log"
 	"github.com/in-toto/go-witness/timestamp"
 	"github.com/in-toto/witness/options"
 	"github.com/spf13/cobra"
@@ -81,5 +83,27 @@ func runSign(ctx context.Context, so options.SignOptions, signers ...cryptoutil.
 	}
 
 	defer outFile.Close()
-	return witness.Sign(inFile, so.DataType, outFile, dsse.SignWithSigners(signers[0]), dsse.SignWithTimestampers(timestampers...))
+
+	// Aggregate all user-defined subjects into a single map
+	allSubjects := make(map[string]cryptoutil.DigestSet)
+
+	// Iterate over user-defined subjects and add them to the aggregated map
+	for _, userDefinedSubject := range so.UserDefinedSubjects {
+		fmt.Printf("User-defined subject: %v\n", userDefinedSubject)
+		ds, err := cryptoutil.CalculateDigestSetFromBytes([]byte(userDefinedSubject),
+			[]cryptoutil.DigestValue{{
+				Hash:   crypto.SHA256,
+				GitOID: false,
+			}})
+
+		if err != nil {
+			log.Debugf("(witness) failed to record user-defined subject %v: %v", userDefinedSubject, err)
+			continue
+		}
+		// Add the user-defined subject to the aggregated map
+		allSubjects["https://witness.dev/internal/user:"+userDefinedSubject] = ds
+	}
+
+	return witness.Sign(inFile, so.DataType, outFile, dsse.SignWithSigners(signers[0]),
+		dsse.SignWithTimestampers(timestampers...), dsse.SignWithUserDefinedSubject(allSubjects))
 }
diff --git a/options/sign.go b/options/sign.go
index 112993c3..0f73f1e2 100644
--- a/options/sign.go
+++ b/options/sign.go
@@ -23,6 +23,7 @@ type SignOptions struct {
 	OutFilePath              string
 	InFilePath               string
 	TimestampServers         []string
+	UserDefinedSubjects      []string
 }
 
 var RequiredSignFlags = []string{
@@ -37,6 +38,6 @@ func (so *SignOptions) AddFlags(cmd *cobra.Command) {
 	cmd.Flags().StringVarP(&so.OutFilePath, "outfile", "o", "", "File to write signed data. Defaults to stdout")
 	cmd.Flags().StringVarP(&so.InFilePath, "infile", "f", "", "Witness policy file to sign")
 	cmd.Flags().StringSliceVar(&so.TimestampServers, "timestamp-servers", []string{}, "Timestamp Authority Servers to use when signing envelope")
-
+	cmd.Flags().StringSliceVarP(&so.UserDefinedSubjects, "user-defined-subject", "u", []string{}, "User-defined linked subjects to include in the attestation")
 	cmd.MarkFlagsRequiredTogether(RequiredSignFlags...)
 }