github.com/GGP1/kure@v0.8.4/commands/2fa/add/add.go (about)

     1  package add
     2  
     3  import (
     4  	"bufio"
     5  	"encoding/base32"
     6  	"fmt"
     7  	"io"
     8  	"net/url"
     9  	"strings"
    10  
    11  	"github.com/GGP1/kure/auth"
    12  	cmdutil "github.com/GGP1/kure/commands"
    13  	"github.com/GGP1/kure/db/totp"
    14  	"github.com/GGP1/kure/pb"
    15  	"github.com/GGP1/kure/terminal"
    16  
    17  	"github.com/pkg/errors"
    18  	"github.com/spf13/cobra"
    19  	bolt "go.etcd.io/bbolt"
    20  )
    21  
    22  const example = `
    23  * Add with setup key
    24  kure 2fa add Sample
    25  
    26  * Add with URL
    27  kure 2fa add -u`
    28  
    29  type addOptions struct {
    30  	digits int32
    31  	url    bool
    32  }
    33  
    34  // NewCmd returns a new command.
    35  func NewCmd(db *bolt.DB, r io.Reader) *cobra.Command {
    36  	opts := addOptions{}
    37  	cmd := &cobra.Command{
    38  		Use:   "add <name>",
    39  		Short: "Add a two-factor authentication code",
    40  		Long: `Add a two-factor authentication code.
    41  
    42  • Using a setup key: services tipically show hyperlinked text like "Enter manually" or "Enter this text code", copy the hexadecimal code given and submit it when requested.
    43  
    44  • Using a URL: extract the URL encoded in the QR code given and submit it when requested. Format: otpauth://totp/{service}:{account}?secret={secret}.`,
    45  		Example: example,
    46  		Args: func(cmd *cobra.Command, args []string) error {
    47  			// When adding with URL the name won't be specified
    48  			if opts.url {
    49  				return nil
    50  			}
    51  
    52  			return cmdutil.MustNotExist(db, cmdutil.TOTP)(cmd, args)
    53  		},
    54  		PreRunE: auth.Login(db),
    55  		RunE:    runAdd(db, r, &opts),
    56  		PostRun: func(cmd *cobra.Command, args []string) {
    57  			opts = addOptions{
    58  				digits: 6,
    59  			}
    60  		},
    61  	}
    62  
    63  	f := cmd.Flags()
    64  	f.Int32VarP(&opts.digits, "digits", "d", 6, "TOTP length {6|7|8}")
    65  	f.BoolVarP(&opts.url, "url", "u", false, "add using a URL")
    66  
    67  	return cmd
    68  }
    69  
    70  func runAdd(db *bolt.DB, r io.Reader, opts *addOptions) cmdutil.RunEFunc {
    71  	return func(cmd *cobra.Command, args []string) error {
    72  		name := strings.Join(args, " ")
    73  		name = cmdutil.NormalizeName(name)
    74  
    75  		if opts.url {
    76  			return addWithURL(db, r)
    77  		}
    78  
    79  		return addWithKey(db, r, name, opts.digits)
    80  	}
    81  }
    82  
    83  func addWithKey(db *bolt.DB, r io.Reader, name string, digits int32) error {
    84  	if digits < 6 || digits > 8 {
    85  		return errors.Errorf("invalid digits number [%d], it must be either 6, 7 or 8", digits)
    86  	}
    87  
    88  	key := terminal.Scanln(bufio.NewReader(r), "Key")
    89  	// Adjust key
    90  	key = strings.ReplaceAll(key, " ", "")
    91  	key += strings.Repeat("=", -len(key)&7)
    92  	key = strings.ToUpper(key)
    93  
    94  	if _, err := base32.StdEncoding.DecodeString(key); err != nil {
    95  		return errors.Wrap(err, "invalid key")
    96  	}
    97  
    98  	return createTOTP(db, name, key, digits)
    99  }
   100  
   101  // addWithURL creates a new TOTP using the values passed in the url.
   102  func addWithURL(db *bolt.DB, r io.Reader) error {
   103  	uri := terminal.Scanln(bufio.NewReader(r), "URL")
   104  	URL, err := url.Parse(uri)
   105  	if err != nil {
   106  		return errors.Wrap(err, "parsing url")
   107  	}
   108  
   109  	query := URL.Query()
   110  	if err := validateURL(URL, query); err != nil {
   111  		return err
   112  	}
   113  
   114  	name := getName(URL.Path)
   115  	if err := cmdutil.Exists(db, name, cmdutil.TOTP); err != nil {
   116  		return err
   117  	}
   118  
   119  	digits := stringDigits(query.Get("digits"))
   120  	secret := query.Get("secret")
   121  	if _, err := base32.StdEncoding.DecodeString(secret); err != nil {
   122  		return errors.Wrap(err, "invalid secret")
   123  	}
   124  
   125  	return createTOTP(db, name, secret, digits)
   126  }
   127  
   128  func createTOTP(db *bolt.DB, name, key string, digits int32) error {
   129  	t := &pb.TOTP{
   130  		Name:   name,
   131  		Raw:    key,
   132  		Digits: digits,
   133  	}
   134  
   135  	if err := totp.Create(db, t); err != nil {
   136  		return err
   137  	}
   138  
   139  	fmt.Printf("\n%q TOTP added\n", name)
   140  	return nil
   141  }
   142  
   143  // getName extracts the service name from the URL.
   144  func getName(path string) string {
   145  	// Given "/Example:account@mail.com", return "Example"
   146  	path = strings.TrimPrefix(path, "/")
   147  	name, _, _ := strings.Cut(path, ":")
   148  	return cmdutil.NormalizeName(name)
   149  }
   150  
   151  // stringDigits returns the digits to use depending on the string passed.
   152  func stringDigits(digits string) int32 {
   153  	switch digits {
   154  	case "8":
   155  		return 8
   156  	case "7":
   157  		return 7
   158  	default:
   159  		return 6
   160  	}
   161  }
   162  
   163  func validateURL(URL *url.URL, query url.Values) error {
   164  	if URL.Scheme != "otpauth" {
   165  		return errors.New("invalid scheme, must be otpauth")
   166  	}
   167  
   168  	if URL.Host != "totp" {
   169  		return errors.New("invalid host, must be totp")
   170  	}
   171  
   172  	algorithm := query.Get("algorithm")
   173  	if algorithm != "" && algorithm != "SHA1" {
   174  		return errors.New("invalid algorithm, must be SHA1")
   175  	}
   176  
   177  	period := query.Get("period")
   178  	if period != "" && period != "30" {
   179  		return errors.New("invalid period, must be 30 seconds")
   180  	}
   181  
   182  	return nil
   183  }