vault_totp/cmd/user/funcs.go

package main

import (
	"bytes"
	`encoding/base64`
	"encoding/json"
	`errors`
	"fmt"
	"hash/crc32"
	`maps`
	"net/url"
	"os"
	`path/filepath`
	"strings"
	`text/template`

	"github.com/davecgh/go-spew/spew"
	"github.com/hashicorp/vault-client-go"
	"github.com/hashicorp/vault-client-go/schema"
	vAPI "github.com/hashicorp/vault/api"
	"github.com/mdp/qrterminal/v3"
	"golang.org/x/term"
	"r00t2.io/goutils/stringsx"
	`r00t2.io/sysutils/paths`
)

func configure() (err error) {

	var buf []byte
	var apiSecret *vAPI.Secret
	var apiHealth *vAPI.HealthResponse
	var clientMap *vault.Response[map[string]interface{}]
	var clientHealth *vault.Response[schema.SealStatusResponse]

	logger.Debug("vault_totp.configure: Configuring client(s).")

	if apiUrl, err = url.Parse(args.VaultConn.URI); err != nil {
		logger.Err("vault_totp.configure: Received error parsing Vault URI '%s': %v", args.VaultConn.URI, err)
		return
	}

	if err = getTok(&args.VaultConn.Token); err != nil {
		logger.Err("vault_totp.configure: Received error getting token: %v", err)
		return
	}
	args.Mounts.Auth = strings.TrimRight(args.Mounts.Auth, "/")
	if len(args.Gen.LookupFmt) == 0 {
		args.Gen.LookupFmt = []string{
			"json",
			"tpl_ref",
		}
	}

	if !args.NoConfirm {
		buf = make([]byte, 1)
		fmt.Printf(
			"Do the following settings look correct? (Y/n)\n"+
				"\tURL:            %s\n"+
				"\tToken (masked): %s\n"+
				"Y/n: ",
			apiUrl.String(),
			stringsx.Redact(apiTok, "*", 5, 5, false),
		)
		if _, err = fmt.Scanln(&buf); err != nil {
			if err.Error() == "unexpected newline" {
				err = nil
			} else {
				logger.Err("vault_totp.configure: Received error receiving configuration confirmation: %v", err)
				return
			}
		}
		if buf[0] != 0x00 && strings.ToLower(string(buf[0])) != "y" {
			logger.Warning("vault_totp.configure: User deferred execution.")
			fmt.Println("Please re-run with any reconfiguration necessary. Quitting.")
			err = fmt.Errorf("user deferred")
			return
		}
	}

	apiCfg = vAPI.DefaultConfig()
	apiCfg.Address = apiUrl.String()

	clientOpts = []vault.ClientOption{
		vault.WithEnvironment(),
		vault.WithAddress(apiUrl.String()),
	}

	if strings.ToLower(apiUrl.Scheme) == "https" {
		// I don't know why Vault/Hashicorp insists on using their own TLS configuration structs instead of crypto/tls.Config.
		// It's super lame and dumb and non-idiomatic.
		apiTls = &vAPI.TLSConfig{
			Insecure: args.VaultConn.Insecure,
		}
		if err = apiCfg.ConfigureTLS(apiTls); err != nil {
			logger.Err("vault_totp.configure: Received error setting API client TLS: %v", err)
			return
		}
		// And this is even MORE dumb. You can only set this during the client initialization.
		clientOpts = append(
			clientOpts,
			vault.WithTLS(
				vault.TLSConfiguration{
					InsecureSkipVerify: args.VaultConn.Insecure,
				},
			),
		)
	}
	if api, err = vAPI.NewClient(apiCfg); err != nil {
		logger.Err("vault_totp.configure: Received error initializing API client: %v", err)
		return
	}
	api.SetToken(apiTok)

	if client, err = vault.New(clientOpts...); err != nil {
		logger.Err("vault_totp.configure: Received error initializing client: %v", err)
		return
	}
	if err = client.SetToken(apiTok); err != nil {
		logger.Err("vault_totp.configure: Received error setting token on client: %v", err)
		return
	}

	// Health checks.
	if apiHealth, err = api.Sys().Health(); err != nil {
		logger.Err("vault_totp.configure: Received error getting health from API client: %v", err)
		return
	}
	if clientHealth, err = client.System.SealStatus(ctx); err != nil {
		logger.Err("vault_totp.configure: Received error getting health from client: %v", err)
		return
	}
	logger.Debug("vault_totp.configure: API health:\n%s", spew.Sdump(apiHealth))
	logger.Debug("vault_totp.configure: Client health:\n%s", spew.Sdump(clientHealth))

	// Token checks.
	if apiSecret, err = api.Auth().Token().LookupSelf(); err != nil {
		logger.Err("vault_totp.configure: Received error getting self token from API: %v", err)
		return
	}
	if clientMap, err = client.Auth.TokenLookUpSelf(ctx); err != nil {
		logger.Err("vault_totp.configure: Received error getting self token from client: %v", err)
		return
	}
	logger.Debug("vault_totp.configure: API self token:\n%s", spew.Sdump(apiSecret))
	logger.Debug("vault_totp.configure: Client self token:\n%s", spew.Sdump(clientMap))

	logger.Debug("vault_totp.configure: Finished configuring client(s).")

	return
}

func genTotp() (err error) {

	var maxKeyLen int
	var resp *vault.Response[map[string]interface{}]
	var req schema.MfaAdminGenerateTotpSecretRequest = schema.MfaAdminGenerateTotpSecretRequest{
		EntityId: entity["id"].(string),
		MethodId: totpMethodId,
	}

	logger.Debug(
		"vault_totp.genTotp: Setting/resetting TOTP for entity ID '%s' on '%s' ('%s' on '%s')...",
		req.EntityId, req.MethodId,
		entity["aliases"].([]interface{})[0].(map[string]interface{})["name"], args.Mounts.Auth,
	)

	if resp, err = client.Identity.MfaAdminGenerateTotpSecret(ctx, req); err != nil {
		logger.Err("vault_totp.genTotp: Received error (re-)generating TOTP secret: %v", err)
		return
	}

	totpSeed = resp.Data

	if totpSeed == nil || len(totpSeed) == 0 {
		// If there wasn't an error but the return's empty/nil, the TOTP exists.
		if args.Gen.Force {
			logger.Debug("vault_totp.genTotp: Clearing existing TOTP secret.")
			if resp, err = client.Identity.MfaAdminDestroyTotpSecret(
				ctx,
				schema.MfaAdminDestroyTotpSecretRequest{
					EntityId: entity["id"].(string),
					MethodId: totpMethodId,
				},
			); err != nil {
				logger.Err("vault_totp.genTotp: Received error clearing existing TOTP: %v", err)
				return
			}
			logger.Debug("vault_totp.genTotp: Successfully cleared TOTP secret; received response %#v", resp)
			if err = genTotp(); err != nil {
				logger.Err("vault_totp.genTotp: Received error re-trying generating TOTP secret: %v", err)
				return
			}
		} else {
			if !args.Gen.Silent {
				fmt.Println("TOTP secret already exists and force was not enabled.")
				totpSeed = nil
			}
			logger.Debug("vault_totp.genTotp: TOTP secret existed and force was not enabled.")
		}
	} else {
		if totpUrl, err = url.Parse(totpSeed["url"].(string)); err != nil {
			logger.Err("vault_totp.genTotp: Received error parsing TOTP URL: %v", err)
			return
		}
		if !args.Gen.Silent {
			// Print it to STDOUT
			if args.Gen.PrintQr {
				// qrterminal.Generate(totpUrl.String(), qrterminal.H, os.Stdout)
				qrterminal.GenerateWithConfig(
					totpUrl.String(),
					qrterminal.Config{
						Level:          qrterminal.L,
						Writer:         os.Stdout,
						HalfBlocks:     false,
						BlackChar:      qrterminal.BLACK,
						BlackWhiteChar: qrterminal.BLACK_WHITE,
						WhiteChar:      qrterminal.WHITE,
						WhiteBlackChar: qrterminal.WHITE_BLACK,
						QuietZone:      1,
						WithSixel:      false,
					},
				)
				fmt.Println()
			}
			fmt.Printf("\n----\nTOTP:\t%s\n\n", strings.TrimPrefix(totpUrl.Path, "/"))
			// I hate that there's not an easier way to do this besides iterating twice.
			for k := range maps.Keys(totpUrl.Query()) {
				if len(k) > maxKeyLen {
					maxKeyLen = len(k)
				}
			}
			for k, v := range totpUrl.Query() {
				for _, s := range v {
					fmt.Printf("%-*s\t%s\n", maxKeyLen, k+":", s)
				}
			}
			fmt.Print("----\n\n")
		}
		if err = writeTotp(); err != nil {
			logger.Err("vault_totp.genTotp: Received error writing TOTP QR to persistent storage: %v", err)
			return
		}

		logger.Debug("vault_totp.genTotp: Generated TOTP.")
	}

	return
}

func getAuthMnt() (err error) {

	var ok bool
	var i interface{}
	var isl []interface{}
	var auths map[string]*vAPI.AuthMount
	var enforcements *vault.Response[schema.StandardListResponse]
	var enforcement *vault.Response[map[string]interface{}]
	// var totpList *vault.Response[schema.StandardListResponse]
	var totpResp *vault.Response[map[string]interface{}]
	var enforcedMfa map[string]struct{} = make(map[string]struct{})

	logger.Debug("vault_totp.getAuthMnt: Fetching auth accessor for '%s'.", args.Mounts.Auth)

	if auths, err = api.Sys().ListAuth(); err != nil {
		logger.Err("vault_totp.getAuthMnt: Received error getting list of enabled auths: %v", err)
		return
	}

	for k, v := range auths {
		if strings.TrimRight(k, "/") != args.Mounts.Auth {
			continue
		}
		// It's not actually the accessor (v.Accessor), despite what some docs say, that the MFA uses.
		// But we need it for the entity lookup.
		authMntAccessor = v.Accessor
		// TOTP uses the UUID.
		authMntId = v.UUID
		logger.Debug("vault_totp.getAuthMnt: Found auth mount '%s':\n%s", args.Mounts.Auth, spew.Sdump(v))
		break
	}

	// I think they intentionally make this a pain in the ass so you pay for Enterprise.
	if enforcements, err = client.Identity.MfaListLoginEnforcements(ctx); err != nil {
		logger.Err("vault_totp.getAuthMnt: Received error getting list of enabled TOTP enforcements: %v", err)
		return
	}
	for _, e := range enforcements.Data.Keys {
		if enforcement, err = client.Identity.MfaReadLoginEnforcement(ctx, e); err != nil {
			logger.Err("vault_totp.getAuthMnt: Received error reading MFA login enforcement '%s': %v", e, err)
			return
		}
		if i, ok = enforcement.Data["auth_method_accessors"]; !ok {
			continue
		}
		if isl, ok = i.([]interface{}); !ok {
			continue
		}
		for _, am := range isl {
			if authMntAccessor == am.(string) {
				logger.Debug(
					"vault_totp.getAuthMnt: Found enforcement for auth at mount '%s' ('%s'):\n%s",
					args.Mounts.Auth, authMntAccessor, spew.Sdump(am),
				)
				if isl, ok = enforcement.Data["mfa_method_ids"].([]interface{}); ok {
					for _, mMI := range isl {
						enforcedMfa[mMI.(string)] = struct{}{}
					}
				}
			}
		}
	}
	/*
		if totpList, err = client.Identity.MfaListTotpMethods(ctx); err != nil {
			logger.Err("vault_totp.getAuthMnt: Received error getting list of enabled TOTP methods: %v", err)
			return
		}
		for _, t := range totpList.Data.Keys {
	*/
	for t := range enforcedMfa {
		if totpResp, err = client.Identity.MfaReadTotpMethod(ctx, t); err != nil {
			logger.Err("vault_totp.getAuthMnt: Received error reading TOTP method ID '%s': %v", t, err)
			return
		}
		totpMethod = totpResp.Data
		totpMethodId = t
		logger.Debug(
			"vault_totp.getAuthMnt: Found TOTP MFA method ID '%s' for '%s':\n%s",
			t, authMntAccessor, spew.Sdump(totpMethod),
		)
		break
	}

	logger.Debug("vault_totp.getAuthMnt: Set auth ID to '%s' for auth mount '%s'.", authMntId, args.Mounts.Auth)

	return
}

func getEnt() (err error) {

	var resp *vault.Response[map[string]interface{}]
	var req schema.EntityLookUpRequest

	logger.Debug("vault_totp.getEnt: Fetching/finding entity.")

	if err = getAuthMnt(); err != nil {
		logger.Err("vault_totp.getEnt: Received error fetching auth mount accessor: %v", err)
		return
	}

	for k, v := range args.Gen.EntityLookup {
		switch k {
		case "name":
			req.Name = v
		case "id":
			req.Id = v
		case "alias":
			req.AliasName = v
			req.AliasMountAccessor = authMntAccessor
		case "alias_id":
			req.AliasId = v
		}
	}

	if resp, err = client.Identity.EntityLookUp(ctx, req); err != nil {
		logger.Err("vault_totp.getEnt: Received error looking up entity with %#v: %v", req, err)
		return
	}

	entity = resp.Data

	logger.Debug("vault_totp.getEnt: Found entity:\n%s", spew.Sdump(entity))

	return
}

func getTok(tok *string) (err error) {

	var b []byte
	var b2 []byte
	var state *term.State
	var fd int = int(os.Stdin.Fd())

	if tok != nil && *tok != "" {
		apiTok = *tok
		return
	}

	if state, err = term.GetState(fd); err != nil {
		logger.Debug("vault_totp.getTok: Received error getting current state: %v", err)
		return
	}
	defer func() {
		var tErr error
		if state != nil {
			if tErr = term.Restore(fd, state); tErr != nil {
				logger.Err("vault_totp.getTok: Received error restoring terminal state: %v", tErr)
			}
		}
		fmt.Println()
	}()

	fmt.Print("Please enter your Vault token (will not echo back): ")
	if b, err = term.ReadPassword(fd); err != nil {
		logger.Err("vault_totp.getTok: Received error reading Vault token (initial): %v", err)
		return
	}
	fmt.Print("\nConfirm: ")
	if b2, err = term.ReadPassword(fd); err != nil {
		logger.Err("vault_totp.getTok: Received error reading Vault token (confirm): %v", err)
		return
	}
	fmt.Println()
	if !bytes.Equal(b, b2) {
		err = fmt.Errorf(
			"original token (CRC32 %x, length %d) "+
				"does not match confirmation (CRC32 %x, length %d)",
			crc32.ChecksumIEEE(b), len(string(b)),
			crc32.ChecksumIEEE(b2), len(string(b2)),
		)
		logger.Err("vault_totp.getTok: Error when confirming token: %v")
		return
	}

	apiTok = string(b)

	return
}

func printEntity() (err error) {

	var b []byte
	var rendered bool
	var sb strings.Builder
	var seen map[string]struct{} = make(map[string]struct{})

	sb.WriteString("----\n")
	for _, i := range args.Gen.LookupFmt {
		if _, rendered = seen[i]; rendered {
			continue
		}
		switch i {
		case "dump":
			spew.Fdump(os.Stdout, entity)
		case "json":
			if b, err = json.MarshalIndent(entity, "", "  "); err != nil {
				logger.Err("vault_totp.printEntity: Received error rendering to JSON: %v", err)
				return
			}
			sb.Write(b)
		default:
			continue
		}
		sb.WriteString("\n----\n")
		seen[i] = struct{}{}
	}

	fmt.Print(sb.String())

	return
}

func renderVPath() (mnt, vpath string, err error) {

	var sl []string
	var buf bytes.Buffer
	var tpl *template.Template = template.New("").Option("missingkey=error")

	logger.Debug("vault_totp.renderVPath: Rendering path for entity TOTP storage.")

	if args.Mounts.TotpPath == "" {
		logger.Warning("vault_totp.renderVPath: TOTP persistent Vault storage disabled.")
		return
	}

	if _, err = tpl.Parse(args.Mounts.TotpPath); err != nil {
		logger.Err("vault_totp.renderVPath: Error parsing template: %v", err)
		return
	}

	if err = tpl.Execute(&buf, entity); err != nil {
		logger.Err("vault_totp.renderVPath: Error executing template: %v", err)
		return
	}

	sl = strings.SplitN(buf.String(), ":", 2)

	switch len(sl) {
	case 0:
		err = fmt.Errorf("no mount and/or path; empty TOTP pathspec")
		return
	case 1:
		err = fmt.Errorf("only mount or path specified; invalid TOTP pathspec")
		return
	}
	if sl[0] == "" {
		err = fmt.Errorf("no mount specified; invalid TOTP pathspec")
		return
	}
	if sl[1] == "" {
		err = fmt.Errorf("no path specified; invalid TOTP pathspec")
		return
	}

	mnt = sl[0]
	vpath = sl[1]

	logger.Debug("vault_totp.renderVPath: Rendered vpath for TOTP: '%s:%s'", mnt, vpath)

	return
}

func writeTotp() (err error) {

	var b []byte
	var fnm string

	logger.Debug("vault_totp.writeTotp: Writing TOTP to file/Vault.")

	if totpSeed == nil {
		logger.Debug("vault_totp.writeTotp: No TOTP seed/secret; NO-OP.")
		return
	}

	if !args.Gen.NoQr {
		if err = paths.RealPath(&args.Gen.QrDir); err != nil {
			logger.Err("vault_totp.writeTotp: Received error canonizing QR directory '%s': %v", args.Gen.QrDir, err)
			return
		}
		if err = os.MkdirAll(args.Gen.QrDir, 0o0700); err != nil {
			logger.Err("vault_totp.writeTotp: Received error creating QR directory '%s': %v", args.Gen.QrDir, err)
			return
		}
		fnm = filepath.Join(args.Gen.QrDir, fmt.Sprintf("totp_%s.png", entity["id"].(string)))
		// wildly misnamed field
		if b, err = base64.StdEncoding.DecodeString(totpSeed["barcode"].(string)); err != nil {
			logger.Err("vault_totp.writeTotp: Received error decoding QR code base64: %v", err)
			return
		}
		if err = os.WriteFile(fnm, b, 0o0600); err != nil {
			logger.Err("vault_totp.writeTotp: Received error writing QR code to disk: %v", err)
			return
		}
	}

	if strings.TrimSpace(args.Mounts.TotpPath) != "" {
		if err = writeTotpVault(); err != nil {
			logger.Err("vault_totp.writeTotp: Received error writing TOTP seed/secret to Vault: %v", err)
			return
		}
	}

	logger.Debug("vault_totp.writeTotp: Wrote TOTP to persistent storage (if specified/configured).")

	return
}

func writeTotpVault() (err error) {

	// var b []byte
	var mnt string
	var vpath string
	var baseUri *url.URL
	var uiItemUri *url.URL
	var apiItemUri *url.URL
	var presp *vAPI.KVSecret
	var dat map[string]interface{}
	var datUri map[string]interface{}
	var rresp *vault.Response[schema.KvV2ReadResponse]
	var wresp *vault.Response[schema.KvV2WriteResponse]
	var respErr *vault.ResponseError = new(vault.ResponseError)
	var params []vault.RequestOption = make([]vault.RequestOption, 0)

	logger.Debug("vault_totp.writeTotpVault: Rendering path for entity TOTP storage.")

	if mnt, vpath, err = renderVPath(); err != nil {
		logger.Err("vault_totp.writeTotpVault: Received error rendering Vault path: %v", err)
		return
	}

	if baseUri, err = url.Parse(api.Address()); err != nil {
		logger.Err("vault_totp.writeTotpVault: Received error parsing Vault non-specialized base URI: %v", err)
		return
	}
	uiItemUri = baseUri.JoinPath("ui", "vault", "secrets", mnt, "kv", url.PathEscape(vpath))
	apiItemUri = baseUri.JoinPath("v1", mnt, "data", url.PathEscape(vpath))

	dat = map[string]interface{}{
		"uri_str": totpUrl.String(),
		"qr_b64":  totpSeed["barcode"].(string),
	}
	datUri = make(map[string]interface{})
	for k, v := range totpUrl.Query() {
		datUri[k] = v[0] // I don't think the otp:// URLs have multi-val query params...
	}
	dat["uri"] = datUri

	params = append(
		params,
		vault.WithMountPath(mnt),
	)

	if rresp, err = client.Secrets.KvV2Read(ctx, vpath, params...); err != nil {
		if errors.As(err, &respErr) {
			if respErr.StatusCode != 404 {
				spew.Dump(err)
				logger.Err("vault_totp.writeTotpVault: Received error reading Vault secret at mount '%s' path '%s': %v", mnt, vpath, err)
				return
			} else {
				err = nil
			}
		} else {
			return
		}
	}

	if rresp != nil && rresp.Data.Data != nil {
		/*
			// Yeah this is dumb. This is what happens when you try to do too much magic and try to "protect" developers from themselves.
			// The proper solution is we need smarter developers, not more Fisher-Price Babby's First Code nonsense.
			if b, err = json.Marshal(dat); err != nil {
				logger.Err("vault_totp.writeTotpVault: Received error with merge encoding of existing data: %v", err)
				return
			}
			dat = rresp.Data.Data
			if err = json.Unmarshal(b, &dat); err != nil {
				logger.Err("vault_totp.writeTotpVault: Received error while unmarshaling merge data: %v", err)
				return
			}
		*/
		// We'll do a patch instead (https://developer.hashicorp.com/vault/api-docs/secret/kv/kv-v2#patch-secret).
		// Fuck it.
		if presp, err = api.KVv2(mnt).Patch(ctx, vpath, dat); err != nil {
			logger.Err("vault_totp.writeTotpVault: Received error patching/updating secret '%s' on mount '%s': %v", vpath, mnt, err)
			return
		}
		logger.Debug("vault_totp.writeTotpVault: Updated secret '%s:%s':\n%s", mnt, vpath, spew.Sdump(presp))
	} else {
		if wresp, err = client.Secrets.KvV2Write(
			ctx,
			vpath,
			schema.KvV2WriteRequest{
				Data: dat,
			},
			params...,
		); err != nil {
			logger.Err("vault_totp.writeTotpVault: Received error creating new secret '%s' on mount '%s': %v", vpath, mnt, err)
			return
		}
		logger.Debug("vault_totp.writeTotpVault: Wrote secret '%s:%s':\n%s", mnt, vpath, spew.Sdump(wresp))
	}

	if !args.Gen.Silent {
		fmt.Printf(
			"TOTP stored at %s:%s\n\n"+
				"Web UI URL:\t%s\n"+
				"API URL:\t%s\n",
			mnt, vpath,
			uiItemUri.String(),
			apiItemUri.String(),
		)
	}

	logger.Debug("vault_totp.writeTotpVault: Wrote TOTP seed/secret to Vault.")

	return
}