jwtcurl/main.go

package main

import (
	"bytes"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strings"
	"time"

	"gitlab-app.eng.qops.net/golang/jwt"
)

type LagReader struct {
	lag time.Duration
	src io.Reader
}

func main() {
	err := Main()
	if err != nil {
		log.Println(err)
		os.Exit(1)
	}
}

func Main() error {
	var bodyRepeat int
	var path, host, method, body, headers, brandID, userID, issuer, basicAuth, claims string
	var ca, cert, key, secret string
	var needJWT, verbose, jsonPP bool
	var timeout, lag time.Duration
	flag.StringVar(&method, "method", "get", "method for request")
	flag.StringVar(&path, "path", "fieldsetdefinitions/v1/index/surveys/SV_031sm3MMOPSa8Tz/fieldsets?assumeHasPermission=true", "path for request")
	flag.StringVar(&host, "host", "data-platform.service.b1-prv.consul:8080", "host and port for request")
	flag.StringVar(&body, "body", "", "body for request")
	flag.IntVar(&bodyRepeat, "bodyrepeat", 1, "repeat body for request")
	flag.StringVar(&brandID, "brand", "testencresponse", "brandID for request JWT")
	flag.StringVar(&userID, "user", "breel", "userid for request JWT")
	flag.StringVar(&basicAuth, "auth", "", "comma separated user,password for basic auth")
	flag.StringVar(&headers, "headers", "", "headers as k=v,k=v for request")
	flag.StringVar(&issuer, "issuer", "dataprocessing,responseengine,fieldset-definitions,qualtrics,objectstore,svs,monolith,ex,blixt,null,responseengine", "issuer for jwt")
	flag.BoolVar(&needJWT, "jwt", true, "need jwt boolean")
	flag.BoolVar(&jsonPP, "jpp", true, "try json pretty print")
	flag.BoolVar(&verbose, "v", false, "is verbose")
	flag.DurationVar(&timeout, "t", time.Second*10, "request timeout")
	flag.DurationVar(&lag, "lag", time.Second*0, "writing request lag after connecting")
	flag.StringVar(&ca, "ca", "", "ca for server")
	flag.StringVar(&cert, "cert", "", "cert for client")
	flag.StringVar(&key, "key", "", "key for client")
	flag.StringVar(&secret, "secret", "", "secret for jwt")
	flag.StringVar(&claims, "claims", "", "extra claims as k=v,k=v")
	flag.Parse()

	if !strings.HasPrefix(host, "http") {
		host = "http://" + host
	}

	c := makeClient(timeout, ca, cert, key)
	var reqBody io.Reader
	if bodyRepeat >= 1 {
		reqBody = strings.NewReader(strings.Repeat(body, bodyRepeat))
	} else {
		reqBody = os.Stdin
	}
	reqBody = NewLagReader(lag, reqBody)
	req, err := http.NewRequest(
		strings.ToUpper(method),
		host+"/"+strings.Trim(path, "/"),
		reqBody,
	)
	if err != nil {
		return err
	}
	req.Header.Add("brandId", brandID)
	req.Header.Add("Content-Type", "application/json")
	if len(headers) > 0 {
		for _, pair := range strings.Split(headers, ",") {
			kv := strings.Split(pair, "=")
			req.Header.Add(kv[0], kv[1])
		}
	}
	if needJWT {
		setJWT(verbose, req, brandID, userID, issuer, secret, claims)
	}
	if basicAuth != "" {
		splits := strings.Split(basicAuth, ",")
		req.SetBasicAuth(splits[0], splits[1])
	}

	if verbose {
		fmt.Fprintf(os.Stderr, "%v\n", req)
	}
	start := time.Now()
	resp, err := c.Do(req)
	elapsed := time.Since(start)
	if err != nil {
		return fmt.Errorf("DO failed: %v", err)
	}
	if verbose {
		fmt.Fprintf(os.Stderr, "%v\n", resp.Header)
	}
	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return fmt.Errorf("READ BODY failed: %v", err)
	}
	defer resp.Body.Close()

	fmt.Fprintf(os.Stderr, "(%d / %v) ", resp.StatusCode, elapsed)
	if jsonPP {
		var v interface{}
		if err := json.Unmarshal(b, &v); err == nil {
			if c, err := json.MarshalIndent(v, "", "   "); err == nil {
				b = c
			}
		}
	}
	fmt.Printf("%s\n", bytes.TrimSpace(b))
	if resp.StatusCode >= http.StatusBadRequest && resp.StatusCode != http.StatusNotFound {
		return fmt.Errorf("Status %v", resp.StatusCode)
	}
	return nil
}

func makeClient(timeout time.Duration, ca, cert, key string) *http.Client {
	transport := &http.Transport{
		TLSClientConfig: &tls.Config{},
	}
	if ca == "" {
		transport.TLSClientConfig.InsecureSkipVerify = true
	} else {
		caBytes, err := ioutil.ReadFile(ca)
		if err != nil {
			panic(err)
		}
		rootCAs := x509.NewCertPool()
		rootCAs.AppendCertsFromPEM(caBytes)
		transport.TLSClientConfig.RootCAs = rootCAs
	}
	if cert != "" && key != "" {
		clientCert, err := tls.LoadX509KeyPair(cert, key)
		if err != nil {
			panic(err)
		}
		transport.TLSClientConfig.Certificates = []tls.Certificate{clientCert}
		transport.TLSClientConfig.BuildNameToCertificate()
	}
	return &http.Client{
		Timeout:   timeout,
		Transport: transport,
	}
}

func setJWT(verbose bool, r *http.Request, brandID, userID string, issuer, secret, claims string) {
	signer := &jwt.Signer{
		Key: []byte(secret),
		DefaultClaims: jwt.Claims{
			Audience: "qualtrics",
			Issuer:   issuer,
			UserID:   userID,
			BrandID:  brandID,
			Custom: map[string]interface{}{
				"IsolationPartitionID": brandID,
				"userType":             "UT_SERVERADMIN",
			},
		},
		IncludeBodyHash: true,
	}
	for _, claim := range strings.Split(claims, ",") {
		c := strings.Split(claim, "=")
		if len(c) < 2 {
			continue
		}
		signer.DefaultClaims.Custom[c[0]] = c[1]
	}

	if err := signer.Sign(r, jwt.Claims{}); err != nil {
		panic(err)
	}

	if verbose {
		log.Println(*signer)
	}
}

func NewLagReader(lag time.Duration, src io.Reader) *LagReader {
	return &LagReader{
		lag: lag,
		src: src,
	}
}

func (lr *LagReader) Read(p []byte) (n int, err error) {
	if lr.lag > 0 {
		<-time.After(lr.lag)
		lr.lag = 0
	}
	return lr.src.Read(p)
}