// Package totp implements time-based one-time passwords (RFC 6238) for
// two-factor login, with the settings every authenticator app supports:
// HMAC-SHA1, 6 digits, a 30-second step. It also makes the one-time
// recovery codes and the QR code that sets up an app.
package totp

import (
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha1"
	"crypto/subtle"
	"encoding/base32"
	"encoding/base64"
	"encoding/binary"
	"errors"
	"fmt"
	"math/big"
	"net/url"
	"strings"
	"time"

	"rsc.io/qr"
)

const (
	// Digits is the length of a code.
	Digits = 6
	// Period is the length of a time step.
	Period = 30 * time.Second
	// Skew is how many steps before and after the current one are
	// accepted, for clocks that are a little off.
	Skew = 1
	// SecretBytes is the length of a secret (160 bits, as RFC 4226
	// recommends).
	SecretBytes = 20
	// RecoveryCodes is how many recovery codes a user gets.
	RecoveryCodes = 10
)

var b32 = base32.StdEncoding.WithPadding(base32.NoPadding)

// NewSecret returns a new random secret, base32 encoded (the form
// authenticator apps take).
func NewSecret() string {
	b := make([]byte, SecretBytes)
	if _, err := rand.Read(b); err != nil {
		panic(err)
	}
	return b32.EncodeToString(b)
}

func decode(secret string) ([]byte, error) {
	s := strings.ToUpper(strings.ReplaceAll(strings.TrimSpace(secret), " ", ""))
	k, err := b32.DecodeString(strings.TrimRight(s, "="))
	if err != nil || len(k) == 0 {
		return nil, errors.New("invalid TOTP secret")
	}
	return k, nil
}

// Step returns the time step t falls in.
func Step(t time.Time) int64 { return t.Unix() / int64(Period/time.Second) }

// Code returns the code for a step.
func Code(secret string, step int64) (string, error) {
	k, err := decode(secret)
	if err != nil {
		return "", err
	}
	var msg [8]byte
	binary.BigEndian.PutUint64(msg[:], uint64(step))
	m := hmac.New(sha1.New, k)
	m.Write(msg[:])
	sum := m.Sum(nil)
	off := sum[len(sum)-1] & 0x0f
	v := binary.BigEndian.Uint32(sum[off:off+4]) & 0x7fffffff
	return fmt.Sprintf("%0*d", Digits, v%1_000_000), nil
}

// Normalize removes the spaces and dashes people type in codes.
func Normalize(code string) string {
	return strings.Map(func(r rune) rune {
		if r == ' ' || r == '-' || r == '\t' {
			return -1
		}
		return r
	}, strings.TrimSpace(code))
}

// IsCode reports whether s looks like a TOTP code (as opposed to a
// recovery code).
func IsCode(s string) bool {
	s = Normalize(s)
	if len(s) != Digits {
		return false
	}
	for _, c := range s {
		if c < '0' || c > '9' {
			return false
		}
	}
	return true
}

// Verify checks a code against the steps around now and returns the step
// it matched. Steps up to and including last (the newest step already
// used) are refused, so a code works only once.
func Verify(secret, code string, now time.Time, last int64) (int64, bool) {
	code = Normalize(code)
	if !IsCode(code) {
		return 0, false
	}
	cur := Step(now)
	for d := -Skew; d <= Skew; d++ {
		step := cur + int64(d)
		if step <= last {
			continue
		}
		want, err := Code(secret, step)
		if err != nil {
			return 0, false
		}
		if subtle.ConstantTimeCompare([]byte(want), []byte(code)) == 1 {
			return step, true
		}
	}
	return 0, false
}

// URI is the otpauth:// URI that authenticator apps read from the QR code.
func URI(issuer, account, secret string) string {
	label := url.PathEscape(issuer) + ":" + url.PathEscape(account)
	q := url.Values{}
	q.Set("secret", secret)
	q.Set("issuer", issuer)
	q.Set("algorithm", "SHA1")
	q.Set("digits", fmt.Sprint(Digits))
	q.Set("period", fmt.Sprint(int(Period/time.Second)))
	return "otpauth://totp/" + label + "?" + q.Encode()
}

// QRDataURI renders text as a QR code: an SVG image as a data: URI, ready
// for an <img> tag.
func QRDataURI(text string) (string, error) {
	c, err := qr.Encode(text, qr.M)
	if err != nil {
		return "", err
	}
	const quiet = 4 // the white border scanners need, in modules
	n := c.Size + 2*quiet
	var b strings.Builder
	fmt.Fprintf(&b, `<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 %d %d" shape-rendering="crispEdges">`, n, n)
	fmt.Fprintf(&b, `<rect width="%d" height="%d" fill="#fff"/><path fill="#000" d="`, n, n)
	for y := 0; y < c.Size; y++ {
		for x := 0; x < c.Size; x++ {
			if c.Black(x, y) {
				fmt.Fprintf(&b, "M%d %dh1v1h-1z", x+quiet, y+quiet)
			}
		}
	}
	b.WriteString(`"/></svg>`)
	return "data:image/svg+xml;base64," + base64.StdEncoding.EncodeToString([]byte(b.String())), nil
}

// recoveryAlphabet has no 0/o, 1/l/i, so codes are easy to read and type.
const recoveryAlphabet = "abcdefghjkmnpqrstuvwxyz23456789"

// NewRecoveryCodes returns n random recovery codes like "k7dm2-xq9pw"
// (about 49 bits each).
func NewRecoveryCodes(n int) []string {
	out := make([]string, n)
	for i := range out {
		var s strings.Builder
		for j := range 10 {
			if j == 5 {
				s.WriteByte('-')
			}
			v, err := rand.Int(rand.Reader, big.NewInt(int64(len(recoveryAlphabet))))
			if err != nil {
				panic(err)
			}
			s.WriteByte(recoveryAlphabet[v.Int64()])
		}
		out[i] = s.String()
	}
	return out
}

// NormalizeRecovery puts a recovery code as typed into the form it is
// stored in (lower case, no spaces or dashes).
func NormalizeRecovery(code string) string {
	return strings.ToLower(Normalize(code))
}