/*
 * Copyright (c) 2021 Arnaud Ysmal.  All Rights Reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS
 * OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 * SUCH DAMAGE.
 */

package main

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/hmac"
	"crypto/md5"
	"crypto/sha1"
	"encoding/base32"
	"encoding/base64"
	"encoding/binary"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"math"
	"os"
	"os/user"
	"strings"
	"time"

	"database/sql"

	_ "github.com/mattn/go-sqlite3"

	"gopkg.in/ini.v1"
)

type entry struct {
	Encrypted bool   `json:"encrypted"`
	Hash      string `json:"hash"`
	Index     int    `json:"index"`
	Type      string `json:"type"`
	Issuer    string `json:"issuer"`
	Secret    string `json:"secret"`
	Enc       string `json:"enc"`
	Period    int    `json:"period"`
	DigitsNo  int    `json:"digits"`
}

func deriveKeyFromPassword(p string, s []byte) ([]byte, []byte) {
	var m []byte
	prev := []byte{}

	for len(m) < 48 {
		a := make([]byte, len(prev)+len(p)+len(s))
		copy(a, prev)
		copy(a[len(prev):], p)
		copy(a[len(prev)+len(p):], s)

		nprev := md5.Sum(a)
		prev = nprev[:]
		m = append(m, prev...)
	}
	return m[:32], m[32:48]
}

func decryptAESCBC(p string, enc string) ([]byte, error) {
	var tmp []byte
	var c cipher.Block
	var m cipher.BlockMode
	var err error

	if tmp, err = base64.StdEncoding.DecodeString(enc); err != nil {
		return []byte{}, err
	}

	if !bytes.Equal(tmp[:8], []byte{0x53, 0x61, 0x6c, 0x74, 0x65, 0x64, 0x5f, 0x5f}) {
		return []byte{}, errors.New("magic not found")
	}

	key, iv := deriveKeyFromPassword(p, tmp[8:16])

	tmp = tmp[16:]

	if c, err = aes.NewCipher(key); err != nil {
		return []byte{}, err
	}

	m = cipher.NewCBCDecrypter(c, iv)
	m.CryptBlocks(tmp, tmp)

	lastval := int(tmp[len(tmp)-1])
	if lastval <= 16 {
		tmp = tmp[:len(tmp)-lastval]
	}
	return tmp, nil
}

func getFirefoxSyncv2Path() (string, error) {
	var u *user.User
	var err error
	var f *ini.File
	var p string

	if u, err = user.Current(); err != nil {
		return "", err
	}
	if f, err = ini.Load(u.HomeDir + "/.mozilla/firefox/profiles.ini"); err != nil {
		return "", err
	}
	p = f.Section("Profile0").Key("Path").String()
	return u.HomeDir + "/.mozilla/firefox/" + p + "/storage-sync-v2.sqlite", nil
}

func getEntriesFromFirefoxAuthenticator(fname string, password string) ([]*entry, error) {
	var data string
	var db *sql.DB
	var err error
	var entries map[string]*entry
	var pwddb string
	var s []byte
	var i int

	if fname == "" {
		fname, err = getFirefoxSyncv2Path()
		if err != nil {
			return nil, err
		}
	}
	if db, err = sql.Open("sqlite3", fname); err != nil {
		return nil, err
	}
	err = db.QueryRow("select data from storage_sync_data where ext_id=\"authenticator@mymindstorm\";").Scan(&data)
	db.Close()
	if err != nil {
		return nil, err
	}

	if data == "" {
		return nil, errors.New("no data found in database")
	}
	if err = json.Unmarshal([]byte(data), &entries); err != nil {
		return nil, err
	}

	if pwd, ok := entries["key"]; ok {
		if password != "" {
			keyPwd, err := decryptAESCBC(password, pwd.Enc)
			if err != nil {
				return nil, err
			}
			pwddb = fmt.Sprintf("%x", keyPwd)
		}
		delete(entries, "key")
	}

	rv := make([]*entry, len(entries))
	for _, v := range entries {
		if v.Encrypted && pwddb != "" {
			if s, err = decryptAESCBC(pwddb, v.Secret); err == nil {
				v.Secret = string(s)
				v.Encrypted = false
			}
		}
		rv[i] = v
		i++
	}
	return rv, err
}

func (e *entry) ComputeOTP() (string, error) {
	var secretBytes []byte
	var err error
	var secret string
	var period int
	var digits int
	const (
		DefaultPeriod   = 30
		DefaultDigitsNo = 6
	)

	secret = strings.ToUpper(strings.TrimSpace(e.Secret))
	if n := len(secret) % 8; n != 0 {
		secret = secret + strings.Repeat("=", 8-n)
	}

	if secretBytes, err = base32.StdEncoding.DecodeString(secret); err != nil {
		return "", errors.New("decoding of secret as base32 failed")
	}

	period = e.Period
	if period == 0 {
		period = DefaultPeriod
	}
	digits = e.DigitsNo
	if digits == 0 {
		digits = DefaultDigitsNo
	}

	mac := hmac.New(sha1.New, secretBytes)

	binary.Write(mac, binary.BigEndian,
		uint64(math.Floor(float64(time.Now().Unix())/float64(period))))
	sum := mac.Sum(nil)

	offset := sum[len(sum)-1] & 0xf
	value := binary.BigEndian.Uint32(sum[offset:offset+4]) & 0x7fffffff
	return fmt.Sprintf("%0*d", digits, uint64(value)%uint64(math.Pow10(digits))), nil
}

func getEntriesFromConfigFile() (entries []*entry, err error) {
	var u *user.User
	var data []byte
	var fname string

	if u, err = user.Current(); err != nil {
		return
	}
	fname = u.HomeDir + "/.config/totp/config.json"
	if err = os.Chmod(fname, 0600); err != nil {
		return
	}
	if data, err = os.ReadFile(fname); err != nil {
		return
	}
	if err = json.Unmarshal([]byte(data), &entries); err != nil {
		return
	}
	return
}

func main() {
	var sentry string
	var dbf string
	var passwd string
	var err error
	var entries []*entry
	var firefox bool

	flag.StringVar(&dbf, "d", "", "Database file")
	flag.StringVar(&sentry, "e", "", "Select entry")
	flag.BoolVar(&firefox, "f", false, "Get tokens from Firefox Authenticator Extension module")
	flag.StringVar(&passwd, "p", "", "Database password")
	flag.Parse()

	if firefox || dbf != "" {
		if entries, err = getEntriesFromFirefoxAuthenticator(dbf, passwd); err != nil {
			fmt.Printf("%s\n", err)
			return
		}
	} else {
		if entries, err = getEntriesFromConfigFile(); err != nil {
			fmt.Printf("%s\n", err)
			return
		}
	}

	for _, v := range entries {
		if v.Encrypted {
			continue
		}

		if sentry == "" {
			p, err := v.ComputeOTP()
			if err != nil {
				fmt.Printf("%s: %s\n", v.Issuer, err)
			} else {
				fmt.Printf("%s: %s\n", v.Issuer, p)
			}
		} else if v.Issuer == sentry {
			p, err := v.ComputeOTP()
			if err != nil {
				fmt.Printf("%s\n", err)
			} else {
				fmt.Printf("%s\n", p)
			}
		}
	}
}