From 53194fccb32bb2d24f2524b853f9f6a30c4f0409 Mon Sep 17 00:00:00 2001
From: Dave Grijalva <grijalva@gmail.com>
Date: Tue, 31 Jan 2017 11:36:56 -0800
Subject: [PATCH] allow claims and headers to be specified at command line

---
 cmd/jwt/app.go  | 24 +++++++++++++++++++++++-
 cmd/jwt/args.go | 23 +++++++++++++++++++++++
 2 files changed, 46 insertions(+), 1 deletion(-)
 create mode 100644 cmd/jwt/args.go

diff --git a/cmd/jwt/app.go b/cmd/jwt/app.go
index c037114..c6f1a38 100644
--- a/cmd/jwt/app.go
+++ b/cmd/jwt/app.go
@@ -25,14 +25,20 @@ var (
 	flagKey     = flag.String("key", "", "path to key file or '-' to read from stdin")
 	flagCompact = flag.Bool("compact", false, "output compact JSON")
 	flagDebug   = flag.Bool("debug", false, "print out all kinds of debug data")
+	flagClaims  = make(ArgList)
+	flagHead    = make(ArgList)
 
 	// Modes - exactly one of these is required
-	flagSign   = flag.String("sign", "", "path to claims object to sign or '-' to read from stdin")
+	flagSign   = flag.String("sign", "", "path to claims object to sign, '-' to read from stdin, or '+' to use only -claim args")
 	flagVerify = flag.String("verify", "", "path to JWT token to verify or '-' to read from stdin")
 	flagShow   = flag.String("show", "", "path to JWT file or '-' to read from stdin")
 )
 
 func main() {
+	// Plug in Var flags
+	flag.Var(flagClaims, "claim", "add additional claims. may be used more than once")
+	flag.Var(flagHead, "head", "add additional header params. may be used more than once")
+
 	// Usage message if you ask for -help or if you mess up inputs.
 	flag.Usage = func() {
 		fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
@@ -74,6 +80,8 @@ func loadData(p string) ([]byte, error) {
 	var rdr io.Reader
 	if p == "-" {
 		rdr = os.Stdin
+	} else if p == "+" {
+		return []byte("{}"), nil
 	} else {
 		if f, err := os.Open(p); err == nil {
 			rdr = f
@@ -171,6 +179,13 @@ func signToken() error {
 		return fmt.Errorf("Couldn't parse claims JSON: %v", err)
 	}
 
+	// add command line claims
+	if len(flagClaims) > 0 {
+		for k, v := range flagClaims {
+			claims[k] = v
+		}
+	}
+
 	// get the key
 	var key interface{}
 	key, err = loadData(*flagKey)
@@ -187,6 +202,13 @@ func signToken() error {
 	// create a new token
 	token := jwt.NewWithClaims(alg, claims)
 
+	// add command line headers
+	if len(flagHead) > 0 {
+		for k, v := range flagHead {
+			token.Header[k] = v
+		}
+	}
+
 	if isEs() {
 		if k, ok := key.([]byte); !ok {
 			return fmt.Errorf("Couldn't convert key data to key")
diff --git a/cmd/jwt/args.go b/cmd/jwt/args.go
new file mode 100644
index 0000000..a5bba5b
--- /dev/null
+++ b/cmd/jwt/args.go
@@ -0,0 +1,23 @@
+package main
+
+import (
+	"encoding/json"
+	"fmt"
+	"strings"
+)
+
+type ArgList map[string]string
+
+func (l ArgList) String() string {
+	data, _ := json.Marshal(l)
+	return string(data)
+}
+
+func (l ArgList) Set(arg string) error {
+	parts := strings.SplitN(arg, "=", 2)
+	if len(parts) != 2 {
+		return fmt.Errorf("Invalid argument '%v'.  Must use format 'key=value'. %v", arg, parts)
+	}
+	l[parts[0]] = parts[1]
+	return nil
+}