pixelqr/main.go

package main

import (
	"bytes"
	"encoding/base32"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net/url"
	"os"
	"time"

	"gopkg.in/alecthomas/kingpin.v2"

	"github.com/qpliu/qrencode-go/qrencode"

	"git.icedream.tech/icedream/pixelqr/internal"
)

const (
	Algorithm_SHA1    = "sha1"
	Algorithm_SHA256  = "sha256"
	Algorithm_SHA512  = "sha512"
	Algorithm_Default = Algorithm_SHA1

	ECLevel_L = "l"
	ECLevel_M = "m"
	ECLevel_Q = "q"
	ECLevel_H = "h"

	Type_TOTP = "totp"
	Type_HOTP = "hotp"

	Digits_Default = 6

	Period_Default = 30 * time.Second
)

var (
	ecLevelMapping = map[string]qrencode.ECLevel{
		ECLevel_L: qrencode.ECLevelL,
		ECLevel_M: qrencode.ECLevelM,
		ECLevel_Q: qrencode.ECLevelQ,
		ECLevel_H: qrencode.ECLevelH,
	}
)

var (
	app = kingpin.New("pixelqr", "Generates QR code from any input and displays it in your ANSI true-color compatible terminal.")

	flagANSImage = app.Flag("draw", "Use ANSI drawing code using ANSImage.").
			Default("false").
			Bool()
	flagECLevel = app.Flag("ec", "Sets the error correction level for the generated code.").
			Default(ECLevel_L).
			Enum(ECLevel_L, ECLevel_M, ECLevel_Q, ECLevel_H)

	cmdText         = app.Command("text", "Generate QR code from plain text file or stdin.")
	cmdTextArgInput = cmdText.Arg("input", "The input file to encode.").
			ExistingFile()

	cmdOtp         = app.Command("otp", "Generate QR code for ")
	cmdOtpFlagType = cmdOtp.Flag("type", "The type of the key.").
			Default(Type_TOTP).
			Enum(Type_TOTP, Type_HOTP)
	cmdOtpFlagLabel = cmdOtp.Flag("label", "The name for the OTP key.").
			String()
	cmdOtpFlagIssuer = cmdOtp.Flag("issuer", "Indicates the provider or service this account is associated with.").
				Default("pixelqr").String()
	cmdOtpFlagAlgorithm = cmdOtp.Flag("algorithm", "Determines the key generation algorithm. Not supported by all authenticator apps.").
				Default(fmt.Sprintf("%s", Algorithm_Default)).
				Enum(Algorithm_SHA1, Algorithm_SHA256, Algorithm_SHA512)
	cmdOtpFlagDigits = cmdOtp.Flag("digits", "Determines how long of a one-time passcode to display to the user.").
				Default(fmt.Sprintf("%d", Digits_Default)).
				Uint8()
	cmdOtpFlagCounter = cmdOtp.Flag("counter", "Required if type is hotp, sets the initial counter value for code generation.").
				Int()
	cmdOtpFlagPeriod = cmdOtp.Flag("period", "Only if type is totp, defines a period that a TOTP code will be valid for. Not supported by all authenticator apps.").
				Default(Period_Default.String()).
				Duration()
	cmdOtpFlagBase32Encoded = cmdOtp.Flag("base32encoded", "Set this to true if the secret master key you pass in has already been base32-encoded.").
				Bool()
	cmdOtpArgSecret = cmdOtp.Arg("secret", "The secret master key.").
			Required().String()
)

var (
	ErrCommandNotFound = errors.New("command not found")
)

func main() {
	var err error
	defer func() {
		if err == nil {
			return
		}
		log.Fatal(err)
	}()

	commandName, err := (app.Parse(os.Args[1:]))
	if err != nil {
		return
	}

	switch commandName {
	case cmdText.FullCommand(): // plain text -> qr
		var inputFile *os.File
		if cmdTextArgInput == nil || *cmdTextArgInput == "-" || *cmdTextArgInput == "" {
			inputFile = os.Stdin
		} else {
			inputFile, err = os.Open(*cmdTextArgInput)
		}
		if err != nil {
			return
		}
		defer inputFile.Close()
		input, err := ioutil.ReadAll(inputFile)
		if err != nil {
			return
		}
		err = generateCliQR(string(input))

	case cmdOtp.FullCommand(): // otp -> qr
		// https://github.com/google/google-authenticator/wiki/Key-Uri-Format
		// otpauth://totp/Example:alice@google.com?secret=JBSWY3DPEHPK3PXP&issuer=Example
		query := url.Values{}
		query.Set("issuer", *cmdOtpFlagIssuer)
		if !*cmdOtpFlagBase32Encoded {
			// base32-encode secret before generating url
			query.Set("secret", base32.StdEncoding.EncodeToString([]byte(*cmdOtpArgSecret)))
		} else {
			// secret already base32-encoded
			query.Set("secret", *cmdOtpArgSecret)
		}
		if cmdOtpFlagAlgorithm != nil && *cmdOtpFlagAlgorithm != Algorithm_Default {
			query.Set("algorithm", *cmdOtpFlagAlgorithm)
		}
		if cmdOtpFlagDigits != nil && *cmdOtpFlagDigits != Digits_Default {
			query.Set("digits", fmt.Sprintf("%d", *cmdOtpFlagDigits))
		}
		/*if cmdOtpFlagLabel != nil && len(*cmdOtpFlagLabel) > 0 {
			query.Set("label"] = *cmdOtpFlagLabel
		}*/
		if cmdOtpFlagPeriod != nil && *cmdOtpFlagPeriod != Period_Default {
			periodSeconds := int(*cmdOtpFlagPeriod / time.Second)
			query.Set("period", fmt.Sprintf("%d", periodSeconds))
		}
		url := &url.URL{
			Scheme:   "otpauth",
			Host:     *cmdOtpFlagType,
			Path:     *cmdOtpFlagLabel,
			RawQuery: query.Encode(),
		}
		log.Println("Generated URL:", url.String())
		err = generateCliQR(url.String())

	default: // other
		err = ErrCommandNotFound
	}
}

func generateCliQR(input string) (err error) {
	grid, err := qrencode.Encode(input, ecLevelMapping[*flagECLevel])
	if err != nil {
		return
	}

	var s string

	if *flagANSImage {

		s, err = internal.ConvertGridToANSImage(grid)
		if err != nil {
			return
		}
	} else {

		w := new(bytes.Buffer)
		err = internal.ConvertGridToUnicode(w, grid)
		if err != nil {
			return
		}
		s = string(w.Bytes())
	}

	fmt.Println(s)

	return
}