mnemonic_key.py

   1 #!/usr/bin/env python3
   2 import hashlib
   3 import math
   4 import sys
   5 import argparse
   6
   7
   8 BITS_PER_WORD = 11
   9
  10
  11 def parse_args():
  12     parser = argparse.ArgumentParser("Encode and decode files as a mnemonic")
  13     parser.add_argument("wordlist", type=str, help="The wordlist to use")
  14     parser.add_argument("--decode", action="store_true")
  15     parser.add_argument(
  16         "--input", type=str, help="The input file when encoding"
  17     )
  18     parser.add_argument(
  19         "--output", type=str, help="The file to write to when decoding"
  20     )
  21     parser.add_argument(
  22         "--length", type=int, help="Length in bytes of the decoded output"
  23     )
  24     return parser.parse_args()
  25
  26
  27 def bits_to_int(bits):
  28     """Convert passed bits into int
  29
  30     Least significant bit first
  31     """
  32     b = 0
  33     for i, bit in enumerate(bits):
  34         b += bit << i
  35
  36     return b
  37
  38
  39 def byte_to_bits(byte):
  40     """Convert byte into bit array
  41
  42     Least significant bit first
  43     """
  44     return [byte >> i & 1 for i in range(8)]
  45
  46
  47 def create_mnemonic(bites, words, bits_per_word=BITS_PER_WORD):
  48     """Create mnemonic from bytes
  49
  50     Create mnemonic phrase from an input byte array.  Each byte
  51     is convert into a bit array (least significant bit first) and
  52     all such bit arrays are concatenated in the order of the input
  53     bytes.  BITS_PER_WORD many bits are consumed from the beginning
  54     of the array and converted into an integer (least significant
  55     bit first) which is used as an index to look up a word in the
  56     given wordlist.  A list of so looked up words is returned.
  57
  58     If necessary, the concatenated bit array is padded with the
  59     beginning bits of the sha256 hash of the input byte array
  60     to get to the next multiple of the word size.
  61
  62     :param bites: The bytes to convert.
  63     :param words: The word list to use, must have 2**n many words
  64     :param bits_per_word: The number of bits to consume per word.  The
  65         word list should be 2**bits_per_word long
  66     :retrun: Mnemonic phrase
  67     """
  68     digest = hashlib.sha256(bites).digest()
  69
  70     bits = []
  71     for b in bites:
  72         bits += byte_to_bits(b)
  73
  74     checksum_bits = []
  75     for b in digest:
  76         checksum_bits += byte_to_bits(b)
  77
  78     n_bits = len(bits)
  79     smallest_n_bits = math.floor((n_bits/BITS_PER_WORD) + 1) * BITS_PER_WORD
  80     bits_missing = smallest_n_bits - n_bits
  81     bits += checksum_bits[0:bits_missing]
  82
  83     mnemonic = []
  84     for i in range(0, len(bits), 11):
  85         word_int = bits_to_int(bits[i:i+11])
  86         mnemonic.append(words[word_int])
  87
  88     return mnemonic
  89
  90
  91 def parse_mnemonic(mnemonic, words):
  92     """Parse mnemonic into bytearray using wordlist
  93
  94     For each word in the mnemonic, find it's 0 indexed position
  95     in the wordlist, convert the the position into a bit array
  96     (lest significant bit first) and concatenate all such bit
  97     arrays. Pad it with 7 * [0] to ensure the last bits fit into the
  98     last byte.  Convert the bit array into a byte array (least
  99     significant bit first)
 100
 101     :param mnemonic: A list of words from the mnemonic
 102     :param words: The (ordered) word list
 103     :return: Decoded bytes
 104     """
 105     bits = []
 106     for word in mnemonic:
 107         i = words.index(word)
 108         bits += [i >> j & 1 for j in range(11)]
 109
 110     n_bits = len(bits)
 111     # Add padding bits to ensure the last chunck has 8 bits
 112     bits += [0] * 7
 113     bites = []
 114     for i in range(0, n_bits, 8):
 115         bites.append(bits_to_int(bits[i:i+8]))
 116
 117     return bytearray(bites)
 118
 119
 120 def run(word_file, encode, in_file, out_file, length):
 121     with open(word_file, "r") as wordlist:
 122         words = [word.strip() for word in wordlist.readlines() if word.strip()]
 123
 124     if encode:
 125         with open(in_file, "rb") as in_file:
 126             bites = bytearray(in_file.read())
 127         mnemonic = create_mnemonic(bites, words)
 128         print("\n".join(mnemonic))
 129
 130     else:
 131         mnemonic = []
 132         for line in sys.stdin.readlines():
 133             mnemonic += line.split()
 134
 135         bites = parse_mnemonic(mnemonic, words)
 136         with open(out_file, "wb") as out_file:
 137             out_file.write(bites[:length])
 138
 139
 140 if __name__ == "__main__":
 141     args = parse_args()
 142     run(
 143         args.wordlist,
 144         not args.decode,
 145         args.input,
 146         args.output,
 147         args.length if args.decode else None,
 148     )