mmgen-wallet/mmgen/sha2.py

#!/usr/bin/env python3
#
# mmgen = Multi-Mode GENerator, command-line Bitcoin cold storage solution
# Copyright (C)2013-2019 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/>.

"""
sha2.py: A non-optimized but very compact implementation of the SHA2 hash
         algorithm for the MMGen suite.  Implements SHA256, SHA512 and
         SHA256Compress (unpadded SHA256, required for Zcash addresses)
"""

from struct import pack,unpack

class Sha2(object):
	'Implementation based on the pseudocode at https://en.wikipedia.org/wiki/SHA-2'
	K = None

	@classmethod
	def initConstants(cls):

		from math import sqrt

		def nextPrime(n=2):
			while True:
				for factor in range(2,int(sqrt(n))+1):
					if n % factor == 0:
						break
				else:
					yield n
				n += 1

		# Find the first nRounds primes
		npgen = nextPrime()
		primes = [next(npgen) for x in range(cls.nRounds)]

		fb_mul = 2 ** cls.wordBits
		def getFractionalBits(n):
			return int((n - int(n)) * fb_mul)

		if cls.use_gmp:
			from gmpy2 import context,set_context,sqrt,cbrt
			set_context(context(precision=75))
		else:
			cbrt = lambda n: pow(n, 1 / 3)

		# First wordBits bits of the fractional parts of the square roots of the first 8 primes
		H = (getFractionalBits(sqrt(n)) for n in primes[:8])

		# First wordBits bits of the fractional parts of the cube roots of the first nRounds primes
		K = (getFractionalBits(cbrt(n)) for n in primes)

		cls.H_init = tuple(H)
		cls.K      = tuple(K)

	def __init__(self,message,preprocess=True):
		'Use preprocess=False for Sha256Compress'
		assert type(message) in (bytes,bytearray,list),'message must be of type bytes, bytearray or list'
		if self.K == None:
			type(self).initConstants()
		self.H = list(self.H_init)
		self.M = message
		self.W = [0] * self.nRounds
		if preprocess:
			self.padMessage()
		self.bytesToWords()
		self.compute()

	def padMessage(self):
		"""
		Pre-processing (padding) (adapted from the standard, translating bits to bytes)
		- Begin with the original message of length MSGLEN bytes
		- Create MLPACK := MSGLEN * 8 encoded as a blkSize-bit big-endian integer
		- Append 0x80 (0b10000000)
		- Append PADLEN null bytes, where PADLEN is the minimum number >= 0 such that
		  MSGLEN + 1 + PADLEN + len(MLPACK) is a multiple of blkSize
		- Append MLPACK
		"""
		mlpack = self.pack_msglen()
		padlen = self.blkSize - (len(self.M) % self.blkSize) - len(mlpack) - 1
		if padlen < 0: padlen += self.blkSize
		self.M = self.M + bytes([0x80] + [0] * padlen) + mlpack

	def bytesToWords(self):
		ws = self.wordSize
		assert len(self.M) % ws == 0
		self.M = tuple([unpack(self.word_fmt,self.M[i*ws:ws+(i*ws)])[0] for i in range(len(self.M) // ws)])

	def digest(self):
		return b''.join((pack(self.word_fmt,w) for w in self.H))

	def hexdigest(self):
		return self.digest().hex()

	def compute(self):
		for i in range(0,len(self.M),16):
			self.processBlock(i)

	def processBlock(self,offset):
		'Process a blkSize-byte chunk of the message'

		def rrotate(a,b): return ((a << self.wordBits-b) & self.wordMask) | (a >> b)
		def addm(a,b): return (a + b) & self.wordMask

		# Copy chunk into first 16 words of message schedule array
		for i in range(16):
			self.W[i] = self.M[offset + i]

		# Extend the first 16 words into the remaining nRounds words of message schedule array
		for i in range(16,self.nRounds):
			g0 = self.W[i-15]
			gamma0 = rrotate(g0,self.g0r1) ^ rrotate(g0,self.g0r2) ^ (g0 >> self.g0r3)
			g1 = self.W[i-2]
			gamma1 = rrotate(g1,self.g1r1) ^ rrotate(g1,self.g1r2) ^ (g1 >> self.g1r3)
			self.W[i] = addm(addm(addm(gamma0,self.W[i-7]),gamma1),self.W[i-16])

		# Initialize working variables from current hash state
		a,b,c,d,e,f,g,h = self.H

		# Compression function main loop
		for i in range(self.nRounds):
			ch = (e & f) ^ (~e & g)
			maj = (a & b) ^ (a & c) ^ (b & c)

			sigma0 = rrotate(a,self.s0r1) ^ rrotate(a,self.s0r2) ^ rrotate(a,self.s0r3)
			sigma1 = rrotate(e,self.s1r1) ^ rrotate(e,self.s1r2) ^ rrotate(e,self.s1r3)

			t1 = addm(addm(addm(addm(h,sigma1),ch),self.K[i]),self.W[i])
			t2 = addm(sigma0,maj)

			h = g
			g = f
			f = e
			e = addm(d,t1)
			d = c
			c = b
			b = a
			a = addm(t1,t2)

		# Save hash state
		for n,v in enumerate((a,b,c,d,e,f,g,h)):
			self.H[n] = addm(self.H[n],v)

class Sha256(Sha2):
	use_gmp = False
	g0r1,g0r2,g0r3 = (7,18,3)
	g1r1,g1r2,g1r3 = (17,19,10)
	s0r1,s0r2,s0r3 = (2,13,22)
	s1r1,s1r2,s1r3 = (6,11,25)
	blkSize = 64
	nRounds = 64
	wordSize = 4
	wordBits = 32
	wordMask = 0xffffffff
	word_fmt = '>I'

	def pack_msglen(self):
		return pack('>Q',len(self.M)*8)

class Sha512(Sha2):
	"""
	SHA-512 is identical in structure to SHA-256, but:
	- the message is broken into 1024-bit chunks,
	- the initial hash values and round constants are extended to 64 bits,
	- there are 80 rounds instead of 64,
	- the message schedule array W has 80 64-bit words instead of 64 32-bit words,
	- to extend the message schedule array W, the loop is from 16 to 79 instead of from 16 to 63,
	- the round constants are based on the first 80 primes 2..409,
	- the word size used for calculations is 64 bits long,
	- the appended length of the message (before pre-processing), in bits, is a 128-bit big-endian integer, and
	- the shift and rotate amounts used are different.
	"""
	use_gmp = True
	g0r1,g0r2,g0r3 = (1,8,7)
	g1r1,g1r2,g1r3 = (19,61,6)
	s0r1,s0r2,s0r3 = (28,34,39)
	s1r1,s1r2,s1r3 = (14,18,41)
	blkSize = 128
	nRounds = 80
	wordSize = 8
	wordBits = 64
	wordMask = 0xffffffffffffffff
	word_fmt = '>Q'

	def pack_msglen(self):
		return  pack('>Q', (len(self.M)*8) // (2**64)) + \
				pack('>Q', (len(self.M)*8) % (2**64))