Add beginning of wordcodec for bytes
This commit is contained in:
parent
ea4f45828d
commit
2278f566bf
|
@ -0,0 +1,152 @@
|
|||
package keys
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"io/ioutil"
|
||||
"math/big"
|
||||
"os"
|
||||
"strings"
|
||||
|
||||
"github.com/pkg/errors"
|
||||
)
|
||||
|
||||
const BankSize = 2048
|
||||
|
||||
// TODO: add error-checking codecs for invalid phrases
|
||||
|
||||
type Codec interface {
|
||||
BytesToWords([]byte) ([]string, error)
|
||||
WordsToBytes([]string) ([]byte, error)
|
||||
}
|
||||
|
||||
type WordCodec struct {
|
||||
words []string
|
||||
bytes map[string]int
|
||||
}
|
||||
|
||||
var _ Codec = WordCodec{}
|
||||
|
||||
func NewCodec(words []string) (codec WordCodec, err error) {
|
||||
if len(words) != BankSize {
|
||||
return codec, errors.Errorf("Bank must have %d words, found %d", BankSize, len(words))
|
||||
}
|
||||
|
||||
b := map[string]int{}
|
||||
for i, w := range words {
|
||||
if _, ok := b[w]; ok {
|
||||
return codec, errors.Errorf("Duplicate word in list: %s", w)
|
||||
}
|
||||
b[w] = i
|
||||
}
|
||||
|
||||
return WordCodec{words, b}, nil
|
||||
}
|
||||
|
||||
func LoadCodec(bank string) (codec WordCodec, err error) {
|
||||
words, err := loadBank(bank)
|
||||
if err != nil {
|
||||
return codec, err
|
||||
}
|
||||
return NewCodec(words)
|
||||
}
|
||||
|
||||
// loadBank opens a wordlist file and returns all words inside
|
||||
func loadBank(bank string) ([]string, error) {
|
||||
filename := "wordlist/" + bank + ".txt"
|
||||
words, err := getData(filename)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
wordsAll := strings.Split(strings.TrimSpace(words), "\n")
|
||||
return wordsAll, nil
|
||||
}
|
||||
|
||||
// TODO: read from go-bind assets
|
||||
func getData(filename string) (string, error) {
|
||||
f, err := os.Open(filename)
|
||||
if err != nil {
|
||||
return "", errors.WithStack(err)
|
||||
}
|
||||
defer f.Close()
|
||||
|
||||
data, err := ioutil.ReadAll(f)
|
||||
if err != nil {
|
||||
return "", errors.WithStack(err)
|
||||
}
|
||||
|
||||
return string(data), nil
|
||||
}
|
||||
|
||||
// given this many bytes, we will produce this many words
|
||||
func wordlenFromBytes(numBytes int) int {
|
||||
// 8 bits per byte, and we add +10 so it rounds up
|
||||
return (8*numBytes + 10) / 11
|
||||
}
|
||||
|
||||
// given this many words, we will produce this many bytes.
|
||||
// sometimes there are two possibilities.
|
||||
// if maybeShorter is true, then represents len OR len-1 bytes
|
||||
func bytelenFromWords(numWords int) (length int, maybeShorter bool) {
|
||||
// calculate the max number of complete bytes we could store in this word
|
||||
length = 11 * numWords / 8
|
||||
// if one less byte would also generate this length, set maybeShorter
|
||||
if wordlenFromBytes(length-1) == numWords {
|
||||
maybeShorter = true
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// TODO: add checksum
|
||||
func (c WordCodec) BytesToWords(data []byte) (words []string, err error) {
|
||||
// 2048 words per bank, which is 2^11.
|
||||
numWords := wordlenFromBytes(len(data))
|
||||
|
||||
n2048 := big.NewInt(2048)
|
||||
nData := big.NewInt(0).SetBytes(data)
|
||||
nRem := big.NewInt(0)
|
||||
// Alternative, use condition "nData.BitLen() > 0"
|
||||
// to allow for shorter words when data has leading 0's
|
||||
for i := 0; i < numWords; i++ {
|
||||
nData.DivMod(nData, n2048, nRem)
|
||||
rem := nRem.Int64()
|
||||
words = append(words, c.words[rem])
|
||||
}
|
||||
fmt.Println(words)
|
||||
return words, nil
|
||||
}
|
||||
|
||||
func (c WordCodec) WordsToBytes(words []string) ([]byte, error) {
|
||||
// // 2048 words per bank, which is 2^11.
|
||||
// numWords := (8*len(dest) + 10) / 11
|
||||
// if numWords != len(words) {
|
||||
// return errors.New(Fmt("Expected %v words for %v dest bytes", numWords, len(dest)))
|
||||
// }
|
||||
|
||||
l := len(words)
|
||||
n2048 := big.NewInt(2048)
|
||||
nData := big.NewInt(0)
|
||||
// since we output words based on the remainder, the first word has the lowest
|
||||
// value... we must load them in reverse order
|
||||
for i := 1; i <= l; i++ {
|
||||
w := words[l-i]
|
||||
rem, ok := c.bytes[w]
|
||||
if !ok {
|
||||
return nil, errors.Errorf("Unrecognized word: %s", w)
|
||||
}
|
||||
nRem := big.NewInt(int64(rem))
|
||||
nData.Mul(nData, n2048)
|
||||
nData.Add(nData, nRem)
|
||||
fmt.Printf("+%d: %v\n", rem, nData)
|
||||
}
|
||||
|
||||
// we copy into a slice of the expected size, so it is not shorter if there
|
||||
// are lots of leading 0s
|
||||
dataBytes := nData.Bytes()
|
||||
fmt.Printf("%#v\n", dataBytes)
|
||||
|
||||
outLen, _ := bytelenFromWords(len(words))
|
||||
output := make([]byte, outLen)
|
||||
copy(output[outLen-len(dataBytes):], dataBytes)
|
||||
fmt.Printf("%#v\n", output)
|
||||
return output, nil
|
||||
}
|
|
@ -0,0 +1,85 @@
|
|||
package keys
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestLengthCalc(t *testing.T) {
|
||||
assert := assert.New(t)
|
||||
|
||||
cases := []struct {
|
||||
bytes, words int
|
||||
flexible bool
|
||||
}{
|
||||
{1, 1, false},
|
||||
{2, 2, false},
|
||||
// bytes pairs with same word count
|
||||
{3, 3, true},
|
||||
{4, 3, true},
|
||||
{5, 4, false},
|
||||
// bytes pairs with same word count
|
||||
{10, 8, true},
|
||||
{11, 8, true},
|
||||
{12, 9, false},
|
||||
{13, 10, false},
|
||||
{20, 15, false},
|
||||
// bytes pairs with same word count
|
||||
{21, 16, true},
|
||||
{32, 24, true},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
wl := wordlenFromBytes(tc.bytes)
|
||||
assert.Equal(tc.words, wl, "%d", tc.bytes)
|
||||
|
||||
bl, flex := bytelenFromWords(tc.words)
|
||||
assert.Equal(tc.flexible, flex, "%d", tc.words)
|
||||
if !flex {
|
||||
assert.Equal(tc.bytes, bl, "%d", tc.words)
|
||||
} else {
|
||||
// check if it is either tc.bytes or tc.bytes +1
|
||||
choices := []int{tc.bytes, tc.bytes + 1}
|
||||
assert.Contains(choices, bl, "%d", tc.words)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestEncodeDecode(t *testing.T) {
|
||||
assert, require := assert.New(t), require.New(t)
|
||||
|
||||
codec, err := LoadCodec("english")
|
||||
require.Nil(err, "%+v", err)
|
||||
|
||||
cases := [][]byte{
|
||||
// {7, 8, 9}, // TODO: 3 words -> 3 or 4 bytes
|
||||
// {12, 54, 99, 11}, // TODO: 3 words -> 3 or 4 bytes
|
||||
{1, 2, 3, 4, 5}, // normal
|
||||
{0, 0, 0, 0, 5, 22, 123, 55, 22}, // leading 0s (9 chars, clear)
|
||||
{22, 44, 55, 1, 13, 0, 0, 0, 0}, // trailing 0s (9 chars, clear)
|
||||
{0, 5, 253, 2, 0}, // leading and trailing zeros
|
||||
// others?
|
||||
}
|
||||
|
||||
for i, tc := range cases {
|
||||
w, err := codec.BytesToWords(tc)
|
||||
if assert.Nil(err, "%d: %v", i, err) {
|
||||
b, err := codec.WordsToBytes(w)
|
||||
if assert.Nil(err, "%d: %v", i, err) {
|
||||
assert.Equal(len(tc), len(b))
|
||||
assert.Equal(tc, b)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestCheckInvalidLists(t *testing.T) {
|
||||
// assert, require := assert.New(t), require.New(t)
|
||||
|
||||
}
|
||||
|
||||
func TestCheckTypoDetection(t *testing.T) {
|
||||
|
||||
}
|
File diff suppressed because it is too large
Load Diff
Loading…
Reference in New Issue