mmgen-wallet/mmgen/baseconv.py

#!/usr/bin/env python3
#
# mmgen = Multi-Mode GENerator, command-line Bitcoin cold storage solution
# Copyright (C)2013-2022 The MMGen Project <mmgen@tuta.io>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

"""
baseconv.py:  base conversion class for the MMGen suite
"""

from hashlib import sha256
from .exception import *
from .util import die

def is_b58_str(s):
	return set(list(s)) <= set(baseconv.digits['b58'])

def is_b32_str(s):
	return set(list(s)) <= set(baseconv.digits['b32'])

class baseconv(object):

	desc = {
		'b58':   ('base58',            'base58-encoded data'),
		'b32':   ('MMGen base32',      'MMGen base32-encoded data created using simple base conversion'),
		'b16':   ('hexadecimal string','base16 (hexadecimal) string data'),
		'b10':   ('base10 string',     'base10 (decimal) string data'),
		'b8':    ('base8 string',      'base8 (octal) string data'),
		'b6d':   ('base6d (die roll)', 'base6 data using the digits from one to six'),
#		'tirosh':('Tirosh mnemonic',   'base1626 mnemonic using truncated Tirosh wordlist'), # not used by wallet
		'mmgen': ('MMGen native mnemonic',
		'MMGen native mnemonic seed phrase created using old Electrum wordlist and simple base conversion'),
	}
	# https://en.wikipedia.org/wiki/Base32#RFC_4648_Base32_alphabet
	# https://tools.ietf.org/html/rfc4648
	digits = {
		'b58': tuple('123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz'),
		'b32': tuple('ABCDEFGHIJKLMNOPQRSTUVWXYZ234567'), # RFC 4648 alphabet
		'b16': tuple('0123456789abcdef'),
		'b10': tuple('0123456789'),
		'b8':  tuple('01234567'),
		'b6d': tuple('123456'),
	}
	mn_base = 1626
	wl_chksums = {
		'mmgen':  '5ca31424',
#		'tirosh': '48f05e1f', # tirosh truncated to mn_base
#		'tirosh1633': '1a5faeff' # tirosh list is 1633 words long!
	}
	seedlen_map = {
		'b58': { 16:22, 24:33, 32:44 },
		'b6d': { 16:50, 24:75, 32:100 },
		'mmgen': { 16:12, 24:18, 32:24 },
	}
	seedlen_map_rev = {
		'b58': { 22:16, 33:24, 44:32 },
		'b6d': { 50:16, 75:24, 100:32 },
		'mmgen': { 12:16, 18:24, 24:32 },
	}

	def __init__(self,wl_id):

		if wl_id == 'mmgen':
			from .mn_electrum import words
			self.digits[wl_id] = words
		elif wl_id not in self.digits:
			raise ValueError(f'{wl_id}: unrecognized mnemonic ID')

		self.wl_id = wl_id

	def get_wordlist(self):
		return self.digits[self.wl_id]

	def get_wordlist_chksum(self):
		return sha256(' '.join(self.digits[self.wl_id]).encode()).hexdigest()[:8]

	def check_wordlist(self):

		wl = self.digits[self.wl_id]
		from .util import qmsg,compare_chksums
		ret = f'Wordlist: {self.wl_id}\nLength: {len(wl)} words'
		new_chksum = self.get_wordlist_chksum()

		compare_chksums(
			new_chksum,
			'generated',
			self.wl_chksums[self.wl_id],
			'saved',
			die_on_fail = True )

		if tuple(sorted(wl)) == wl:
			return ret + '\nList is sorted'
		else:
			die(3,'ERROR: List is not sorted!')

	@staticmethod
	def get_pad(pad,seed_pad_func):
		"""
		'pad' argument to baseconv conversion methods must be either None, 'seed' or an integer.
		If None, output of minimum (but never zero) length will be produced.
		If 'seed', output length will be mapped from input length using data in seedlen_map.
		If an integer, the string, hex string or byte output will be padded to this length.
		"""
		if pad == None:
			return 0
		elif type(pad) == int:
			return pad
		elif pad == 'seed':
			return seed_pad_func()
		else:
			raise BaseConversionPadError(f"{pad!r}: illegal value for 'pad' (must be None,'seed' or int)")

	def tohex(self,words_arg,pad=None):
		"convert string or list data of instance base to hex string"
		return self.tobytes(words_arg,pad//2 if type(pad)==int else pad).hex()

	def tobytes(self,words_arg,pad=None):
		"convert string or list data of instance base to byte string"

		words = words_arg if isinstance(words_arg,(list,tuple)) else tuple(words_arg.strip())
		desc = self.desc[self.wl_id][0]

		if len(words) == 0:
			raise BaseConversionError(f'empty {desc} data')

		def get_seed_pad():
			assert self.wl_id in self.seedlen_map_rev, f'seed padding not supported for base {self.wl_id!r}'
			d = self.seedlen_map_rev[self.wl_id]
			if not len(words) in d:
				raise BaseConversionError(
					f'{len(words)}: invalid length for seed-padded {desc} data in base conversion' )
			return d[len(words)]

		pad_val = max(self.get_pad(pad,get_seed_pad),1)
		wl = self.digits[self.wl_id]
		base = len(wl)

		if not set(words) <= set(wl):
			raise BaseConversionError(
				( 'seed data' if pad == 'seed' else f'{words_arg!r}:' ) +
				f' not in {desc} format' )

		ret = sum([wl.index(words[::-1][i])*(base**i) for i in range(len(words))])
		bl = ret.bit_length()
		return ret.to_bytes(max(pad_val,bl//8+bool(bl%8)),'big')

	def fromhex(self,hexstr,pad=None,tostr=False):
		"convert hex string to list or string data of instance base"

		from .util import is_hex_str
		if not is_hex_str(hexstr):
			raise HexadecimalStringError(
				( 'seed data' if pad == 'seed' else f'{hexstr!r}:' ) +
				' not a hexadecimal string' )

		return self.frombytes( bytes.fromhex(hexstr), pad, tostr )

	def frombytes(self,bytestr,pad=None,tostr=False):
		"convert byte string to list or string data of instance base"

		if not bytestr:
			raise BaseConversionError('empty data not allowed in base conversion')

		def get_seed_pad():
			assert self.wl_id in self.seedlen_map, f'seed padding not supported for base {self.wl_id!r}'
			d = self.seedlen_map[self.wl_id]
			if not len(bytestr) in d:
				raise SeedLengthError(
					f'{len(bytestr)}: invalid byte length for seed data in seed-padded base conversion' )
			return d[len(bytestr)]

		pad = max(self.get_pad(pad,get_seed_pad),1)
		wl = self.digits[self.wl_id]

		def gen():
			num = int.from_bytes(bytestr,'big')
			base = len(wl)
			while num:
				yield num % base
				num //= base

		ret = list(gen())
		o = [wl[n] for n in [0] * (pad-len(ret)) + ret[::-1]]

		return (' ' if self.wl_id == 'mmgen' else '').join(o) if tostr else o