vault_totp/cmd/add/funcs.go

package main

import (
	`bufio`
	`bytes`
	`context`
	`encoding/base64`
	`fmt`
	`image`
	`image/color`
	_ `image/gif`
	_ `image/jpeg`
	`image/png`
	_ `image/png`
	`io`
	`os`
	`strconv`
	`strings`
	`sync`

	`github.com/creachadair/otp/otpauth`
	`github.com/hashicorp/vault-client-go`
	`github.com/hashicorp/vault-client-go/schema`
	`github.com/makiuchi-d/gozxing`
	`github.com/makiuchi-d/gozxing/qrcode`
	`github.com/pquerna/otp`
	`r00t2.io/gosecret`
	`r00t2.io/sysutils/paths`
	`r00t2.io/vault_totp/common`
	`r00t2.io/vault_totp/internal`
)

// convertToGray converts image.Image `img` to an 8-bit grayscale image for QR codes.
func convertToGray(img image.Image) (g *image.Gray) {

	var x int
	var y int
	var rgba color.Color
	var bounds image.Rectangle

	bounds = img.Bounds()

	g = image.NewGray(bounds)

	for x = 0; x < bounds.Max.X; x++ {
		for y = 0; y < bounds.Max.Y; y++ {
			rgba = img.At(x, y)
			// g.Set(x, y, rgba)
			g.Set(x, y, color.GrayModel.Convert(rgba))
		}
	}

	return
}

/*
convertToMono converts image.Gray `img` (see convertToGray) to a monochrome (2-bit) grayscale image for QR codes.

If `invert` is true, the bits will be flipped.
(This works best if the original image is sent through convertToMono with invert == false *first*).
*/
func convertToMono(img image.Image, invert bool) (bw *image.Paletted) {

	var x int
	var y int
	var g color.Gray
	var palette color.Palette
	var bounds image.Rectangle
	var hiBit color.Color = color.White
	var loBit color.Color = color.Black

	if invert {
		hiBit = color.Black
		loBit = color.White
	}
	palette = color.Palette{
		loBit,
		hiBit,
	}

	bounds = img.Bounds()
	// bw = image.NewGray(bounds)
	bw = image.NewPaletted(bounds, palette)

	for x = 0; x < bounds.Max.X; x++ {
		for y = 0; y < bounds.Max.Y; y++ {
			g = img.At(x, y).(color.Gray)
			// between 0-255 (uint8)
			if g.Y > 128 {
				bw.Set(x, y, hiBit)
			} else {
				bw.Set(x, y, loBit)
			}
		}
	}

	return
}

// decQR reduces some repetition in parseQrFile. err == nil indicates successfull decode.
func decQR(img image.Image) (res *gozxing.Result, err error) {

	var lum gozxing.LuminanceSource
	var bmp *gozxing.BinaryBitmap
	var binarizer gozxing.Binarizer
	var dec gozxing.Reader = qrcode.NewQRCodeReader()
	var hints map[gozxing.DecodeHintType]interface{} = map[gozxing.DecodeHintType]interface{}{
		gozxing.DecodeHintType_TRY_HARDER: true,
	}

	lum = gozxing.NewLuminanceSourceFromImage(img)
	binarizer = gozxing.NewGlobalHistgramBinarizer(lum)
	if bmp, err = gozxing.NewBinaryBitmap(binarizer); err != nil {
		return
	}
	if res, err = dec.Decode(bmp, hints); err != nil {
		return
	}

	return
}

// normalizeName returns a normalized Vault-safe version of the name.
func normalizeName(issuer, account string) (s string) {

	var i string = issuer
	var a string = account

	common.VaultEscape(&i)
	common.VaultEscape(&a)

	s = fmt.Sprintf("%s%s%s", i, vaultSep, a)

	return
}

// parseExplicit parses an explicit OTP specification.
func parseExplicit(e ExplicitOtp) {

	var ok bool
	var p parsedUrl = parsedUrl{
		u:          nil,
		source:     fmt.Sprintf("%#v", e),
		sourceType: "explicit",
	}

	defer wg.Done()

	p.vaultName = normalizeName(e.Issuer, e.Account)
	p.u = &otpauth.URL{
		Type:      strings.ToLower(e.Type),
		Issuer:    e.Issuer,
		Account:   e.Account,
		RawSecret: e.Secret,
		Algorithm: e.Algorithm,
		Digits:    e.Digits,
		Period:    int(e.Period.Seconds()),
		Counter:   0,
	}

	if _, ok = existingOtp[p.vaultName]; ok {
		logger.Debug("parseExplicit: Found existing otp entry for '%s'; skipping.", p.vaultName)
		return
	}
	// Simple as.
	urlChan <- p

	return
}

// parseAsync is a helper for other parsers.
func parseAsync(uriStr, src, srcType string, pwg *sync.WaitGroup) {

	var ok bool
	var sL string
	var err error
	var u *otpauth.URL
	var multi []*otpauth.URL
	var p parsedUrl = parsedUrl{
		vaultName:  "",
		u:          nil,
		source:     src,
		sourceType: srcType,
	}

	defer pwg.Done()

	sL = strings.ToLower(uriStr)
	if strings.HasPrefix(sL, stdScheme) {
		if p.u, err = otpauth.ParseURL(uriStr); err != nil {
			logger.Err("parseAsync: Failed to parse standard URL '%s' for %s '%s': %v", uriStr, srcType, src, err)
			return
		}
		p.vaultName = normalizeName(p.u.Issuer, p.u.Account)
		if _, ok = existingOtp[p.vaultName]; ok {
			logger.Debug("parseAsync: Found existing otp entry for '%s'; skipping.", p.vaultName)
			return
		}
		urlChan <- p
	} else if strings.HasPrefix(sL, googleScheme) {
		if multi, err = otpauth.ParseMigrationURL(uriStr); err != nil {
			logger.Err("parseAsync: Failed to parse Google URLs for %s '%s': %v", srcType, src, err)
			return
		}
		for _, u = range multi {
			p = parsedUrl{
				vaultName:  normalizeName(u.Issuer, u.Account),
				u:          new(otpauth.URL),
				source:     src,
				sourceType: srcType,
			}
			if _, ok = existingOtp[p.vaultName]; ok {
				logger.Debug("parseAsync: Found existing otp entry for '%s'; skipping.", p.vaultName)
				continue
			}
			*p.u = *u
			urlChan <- p
		}
	} else {
		logger.Err("parseAsync: invalid uriStr '%s' for %s '%s'", uriStr, srcType, src)
		return
	}

	return
}

// parseOtpFile reads in file at filepath `fpath`, reading a TOTP URI per-line.
func parseOtpFile(fpath string) {

	var err error
	var f *os.File
	var line string
	var buf *bufio.Reader
	var pwg sync.WaitGroup
	var fp string = fpath

	defer wg.Done()

	if err = paths.RealPath(&fp); err != nil {
		logger.Err("parseOtpFile: Failed to canonize filepath '%s': %v", fp, err)
		return
	}

	if f, err = os.Open(fp); err != nil {
		logger.Err("parseOtpFile: Failed to open file '%s': %v", fp, err)
		return
	}
	defer func() {
		var fErr error
		if fErr = f.Close(); fErr != nil {
			logger.Err("parseOtpFile: Failed to close file '%s': %v", fp, fErr)
		}
	}()
	buf = bufio.NewReader(f)

	for {
		line, err = buf.ReadString('\n')
		if err != nil && err != io.EOF {
			logger.Err("parseOtpFile: Failed to read from file '%s': %v", fp, err)
			return
		}
		line = strings.TrimSpace(line)
		if line != "" {
			pwg.Add(1)
			go parseAsync(line, fp, "file", &pwg)
		}
		if err == io.EOF {
			break
		}
	}

	pwg.Wait()

	return
}

// parseOtpUri parses in an already-formatted TOTP URI.
func parseOtpUri(uriStr string) {

	var pwg sync.WaitGroup

	defer wg.Done()

	pwg.Add(1)
	go parseAsync(uriStr, "URI", "<CLI>", &pwg)
	pwg.Wait()

	// parseAsync handles the channel send etc.

	return
}

// parseQrFile reads in QR code image at filepath `fpath`.
func parseQrFile(fpath string) {

	var err error
	var f *os.File
	var s string
	var bw image.Image
	var img image.Image
	var pwg sync.WaitGroup
	var res *gozxing.Result
	var fp string = fpath

	defer wg.Done()

	if err = paths.RealPath(&fp); err != nil {
		logger.Err("parseQrFile: Failed to canonize filepath '%s': %v", fp, err)
		return
	}

	if f, err = os.Open(fp); err != nil {
		logger.Err("parseQrFile: Failed to open file '%s': %v", fpath, err)
		return
	}

	if img, s, err = image.Decode(f); err != nil {
		logger.Err("parseQrFile: Failed to decode image '%s': %v", fp, err)
		return
	}
	logger.Debug("parseQrFile: Decoded image '%s', decoded as '%s'", fp, s)

	/*
		PHASE 1: LUMINANCE
		Luminance is performed at every phase, but this is the only modification
		this phase.
		This is done to reduce noise in QR images and does a much
		better job with "vanity" QR, where they abuse EC to put a logo in
		the center or summat. Note that it still doesn't work for QR codes
		surrounded by a black boundary and uses black encoding.
		https://github.com/makiuchi-d/gozxing/issues/76
		We have a couple things to try for that though if this fails.
	*/
	if res, err = decQR(img); err != nil {
		/*
			PHASE 2: CONVERT TO GRAYSCALE
			If the above wasn't successful, try with a grayscale image.
		*/
		bw = convertToGray(img)
		if res, err = decQR(bw); err != nil {
			/*
				PHASE 3: CONVERT TO MONOCHROME
				If the above wasn't successful, try with a monochrome image.
			*/
			bw = convertToMono(bw, false)
			if res, err = decQR(bw); err != nil {
				/*
					PHASE 4
					INVERTED MONOCHROME
					If the above wasn't successful, try with the monochrome inverted.
				*/
				bw = convertToMono(bw, true)
				if res, err = decQR(bw); err != nil {
					// I'm all out of tricks.
					logger.Err("parseQrFile: Failed to decode image '%s' as original, grayscale, monochrome, or inverted monochrome: %v", fp, err)
					return
				}
			}
		}
	}

	pwg.Add(1)
	go parseAsync(res.GetText(), fp, "QR code file", &pwg)
	pwg.Wait()

	// parseAsync handles the channel send etc.

	return
}

// parseSs reads in TOTP information from the TOTP GNOME extension storage in SecretService.
func parseSs() {

	var err error
	var pwg sync.WaitGroup
	var item *gosecret.Item
	var svc *gosecret.Service
	var items []*gosecret.Item
	var coll *gosecret.Collection

	defer wg.Done()

	if svc, err = gosecret.NewService(); err != nil {
		logger.Err("parseSs: Failed to attach to gosecret service: %v", err)
		return
	}
	defer func() {
		var sErr error
		if sErr = svc.Close(); sErr != nil {
			logger.Warning("parseSs: Received error closing gosecret service: %v", sErr)
		}
	}()

	if coll, err = svc.GetCollection(ssCollectionNm); err != nil {
		logger.Err("parseSs: Failed to fetch gosecret collection '%s': %v", ssCollectionNm, err)
		return
	}

	if items, err = coll.Items(); err != nil {
		logger.Err("parseSs: Failed to fetch gosecret items from '%s': %v", ssCollectionNm, err)
		return
	}

	for _, item = range items {
		pwg.Add(1)
		go parseSsItemAsync(*item, &pwg)
	}

	pwg.Wait()

	return
}

// parseSsItemAsync is used by parseSs.
func parseSsItemAsync(item gosecret.Item, pwg *sync.WaitGroup) {

	var err error
	var i int
	var ok bool
	var s string
	var iss string
	var acct string
	var p parsedUrl
	var spl []string

	defer pwg.Done()

	// Fallback values from the name, though these *should* be attributes in the Item.
	// OTP plugin uses secret names like:
	//   <id>:<issuer>:<account>
	// where <id> is a zero-indexed unsigned integer that determines a secret's position.
	// No, I don't know why he didn't just set the ordering/weight as an item attribute.
	spl = strings.SplitN(item.LabelName, ":", 3)

	if iss, ok = item.Attrs["issuer"]; !ok {
		switch len(spl) {
		case 1:
			iss = spl[0]
		default:
			iss = spl[1]
		}
	}
	if acct, ok = item.Attrs["name"]; !ok {
		switch len(spl) {
		case 1:
			acct = spl[0]
		case 2:
			acct = spl[1]
		default:
			acct = spl[2]
		}
	}

	p = parsedUrl{
		vaultName:  normalizeName(iss, acct),
		u:          nil,
		source:     fmt.Sprintf("%s/%s", ssCollectionNm, item.LabelName),
		sourceType: "SecretService",
	}
	if _, ok = existingOtp[p.vaultName]; ok {
		logger.Debug("parseSsItemAsync: Found existing otp entry for '%s'; skipping.", p.vaultName)
		return
	}

	p.u = &otpauth.URL{
		Type:      "",
		Issuer:    iss,
		Account:   acct,
		RawSecret: string(item.Secret.Value),
		Algorithm: "",
		Digits:    0,
		Period:    0,
		Counter:   0,
	}

	if s, ok = item.Attrs["type"]; ok {
		p.u.Type = strings.ToLower(s)
	} else {
		p.u.Type = "totp"
	}

	if s, ok = item.Attrs["algorithm"]; ok {
		p.u.Algorithm = strings.ToUpper(
			strings.ReplaceAll(s, "-", ""),
		)
	} else {
		p.u.Algorithm = "SHA256" // Probably the most common.
	}

	if s, ok = item.Attrs["digits"]; ok {
		if i, err = strconv.Atoi(s); err != nil {
			logger.Err("parseSsItemAsync: Failed to parse digits '%s' for %s '%s': %v", s, p.sourceType, p.source, err)
			return
		}
		p.u.Digits = i
	} else {
		p.u.Digits = 6
	}

	if s, ok = item.Attrs["period"]; ok {
		if i, err = strconv.Atoi(s); err != nil {
			logger.Err("parseSsItemAsync: Failed to parse perid '%s' for %s '%s': %v", s, p.sourceType, p.source, err)
			return
		}
	} else {
		p.u.Period = 30
	}

	urlChan <- p

	return
}

// readUrls reads in parsed OTPAuth URLs from various sources and writes them to Vault.
func readUrls() {

	var err error
	var u parsedUrl
	var vc *vault.Client
	var pwg sync.WaitGroup
	var ctx context.Context = context.Background()

	if vc, err = internal.GetVaultClient(&args.VaultArgs); err != nil {
		logger.Err("readUrls: Failed to get Vault client: %v", err)
		vaultReady <- false
		return
	}

	if existingOtp, err = internal.ListTotpKeys(ctx, vc, args.AddArgs.VaultTotpMnt); err != nil {
		logger.Err("readUrls: Failed to list totp keys: %v", err)
		vaultReady <- false
		return
	}
	vaultReady <- true

	for u = range urlChan {
		if u.u != nil {
			logger.Debug("readUrls: Read in '%s' from %s '%s'", u.u.String(), u.sourceType, u.source)
			pwg.Add(1)
			go storeUrl(u, ctx, vc, &pwg)
			if _, err = vc.Secrets.TotpCreateKey(
				ctx,
				u.vaultName,
				schema.TotpCreateKeyRequest{
					AccountName: "",
					Algorithm:   "",
					Digits:      0,
					Exported:    false,
					Generate:    false,
					Issuer:      "",
					Key:         "",
					KeySize:     0,
					Period:      "",
					QrSize:      0,
					Skew:        0,
					Url:         u.u.String(),
				},
				vault.WithMountPath(args.AddArgs.VaultTotpMnt),
			); err != nil {
				logger.Err("readUrls: Failed to create key '%s' on '%s': %v", u.vaultName, args.AddArgs.VaultTotpMnt, err)
				err = nil
				continue
			}
		}
	}

	pwg.Wait()

	doneChan <- true
}

func storeUrl(p parsedUrl, vctx context.Context, vc *vault.Client, pwg *sync.WaitGroup) {

	var err error
	var mnt string
	var qrB64 string
	var secPath string
	var img image.Image
	var otpk *otp.Key
	var dat map[string]interface{}
	var buf *bytes.Buffer = new(bytes.Buffer)

	defer pwg.Done()

	if args.AddArgs.VaultKV2MntPath == nil {
		return
	}

	// QR CODE. Any failures here should be non-fatal; capturing the secret is more important.
	if otpk, err = otp.NewKeyFromURL(p.u.String()); err != nil {
		logger.Err("storeUrl: Failed to create OTP key for '%s': %v", p.u.String(), err)
	} else {
		if img, err = otpk.Image(512, 512); err != nil {
			logger.Err("storeUrl: Failed to create QR code for '%s': %v", p.u.String(), err)
		} else {
			// WHY DOES IT CREATE A 16-BIT IMAGE.
			// IT'S MONOCHROME; IT USES A WHOLE-ASS 1 BIT OF FIDELITY.
			// We can save a whole-ass ~17% reduction by just converting to 8-bit,
			// but we save ~32% reduction by going to 1-bit (true monochrome).
			img = convertToMono(convertToGray(img), false)
			if err = png.Encode(buf, img); err != nil {
				logger.Err("storeUrl: Failed to encode QR code for '%s': %v", p.u.String(), err)
			} else {
				qrB64 = base64.StdEncoding.EncodeToString(buf.Bytes())
			}
		}
	}

	mnt, secPath = internal.SplitVaultPathspec2(*args.AddArgs.VaultKV2MntPath)
	if mnt == "" {
		mnt = defStoreMnt
	}
	secPath = path.Join(secPath, p.u.Issuer, p.u.Account)

	dat = map[string]interface{}{
		// PNG image data, 512 x 512, 16-bit grayscale, non-interlaced
		"qr_b64":  qrB64,
		"uri_str": p.u.String(),
		"generator": fmt.Sprintf(
			"[%s]%s",
			args.AddArgs.VaultTotpMnt,
		),
		"uri": map[string]any{
			"account":   p.u.Account,
			"algorithm": p.u.Algorithm,
			"digits":    p.u.Digits,
			"issuer":    p.u.Issuer,
			"period":    p.u.Period,
			"secret":    p.u.RawSecret,
		},
	}

	if _, err = vc.Secrets.KvV2Write(
		vctx,
		secPath,
		schema.KvV2WriteRequest{
			Data: dat,
			// Options: nil,
		},
		vault.WithMountPath(mnt),
	); err != nil {
		logger.Err("storeUrl: Failed to store key '%s' from %s '%s': %v", p.vaultName, p.sourceType, p.source, err)
		return
	}
}