Scrypt

import hashlib, base64

#Scrypt(password, salt, itteration_count, block_size, parallel, output_size)
password = b'p@$Sw0rD~7'
salt = b'aa1f2d3f4d23ac44e9c5a6c3d8f9ee8c'
itteration_count = 2048
block_size = 8
parallel = 1
output_size = 32

key = hashlib.scrypt(password, salt=salt, n=itteration_count, r=block_size, p=parallel, dklen=output_size)
print("Derived key:", key.hex())
#Derived key: e813a6f6ccc4e9110193bf9efb7c0a489d76655f9e36629dccbeaf2a73bc0c6f

print("{}${}${}${}${}".format(itteration_count, block_size, parallel, base64.b64encode(salt).decode('utf-8'), base64.b64encode(key).decode('utf-8')))

from cryptopals_lib import *  
from pbkdf2 import pbkdf2
from chacha import salsa_key_schedule
import os, hashlib



class Scrypt():
	"""docstring for Scrypt"""
	def __init__(self, itterations=16384, memory_cost=8, parallel_cost=1, keylength=64):
		#Check that memory factor * parallel factor is not greater than 2^30
		if memory_cost * parallel_cost > 2 ** 30:
			raise Exception("Too much memory in use.")
		self.memory_cost =  memory_cost
		self.parallel_cost =  parallel_cost

		#Check if Itterations is not a power of 2 greater than 2^0
		if itterations < 2 or (itterations & (itterations -1)):
			raise Exception("itterations not a power of 2.")
		self.itterations = itterations

		self.keylength = keylength

		#Initalize other temp buffers
		self.memory_buffer = [ 0 ] * (self.memory_cost << 6)
		self.itter_buffer  = [ 0 ] * ((self.memory_cost * self.itterations) << 5)


	def _smix(self, buffer, round_num):
		index_from = (round_num * self.memory_cost) << 5
		block_size = (self.memory_cost<< 5)

		#Populate the Memory buffer from the input buffer
		self.memory_buffer[:block_size] = buffer[index_from: index_from + block_size]
		#print(self.memory_buffer)

		#Mix the memory buffer and update the itteration buffer
		for i in range(self.itterations):
			index_to = i * block_size
			self.itter_buffer[index_to:index_to + block_size] = self.memory_buffer[:block_size]

			#Do a Block mix
			self.memory_buffer = self._block_mix(self.memory_buffer)

		#
		for i in range(self.itterations):
			start_index = self.memory_buffer[(2 * self.memory_cost -1) << 4] & (self.itterations -1)

			for j in range(block_size):
				self.memory_buffer[j] ^= self.itter_buffer[(start_index * block_size) + j]

			#Do a Block mix
			self.memory_buffer = self._block_mix(self.memory_buffer)

		#print(self.memory_buffer)

		buffer[index_from:index_from + block_size] = self.memory_buffer[:block_size]

		return buffer

	def _block_mix(self, buffer):
		start_index = (2 * self.memory_cost - 1) << 4

		temp = buffer[start_index:start_index+16]

		for i in range(2*self.memory_cost):

			for j in range(16):
				temp[j] ^= buffer[(i <<4) + j]

			#Salsa round
			temp = salsa_key_schedule(temp, rounds=8)

			#Choose to replace in buffer
			index_to = (self.memory_cost << 5) + (i << 4)
			buffer[index_to:index_to+16] = temp[:16]


		#Copy Blocks arround
		for i in range(self.memory_cost): 
			index_from = (self.memory_cost + i) << 5
			index_to = i << 4
			buffer[index_to:index_to + 16] = buffer[index_from:index_from + 16]


		for i in range(self.memory_cost):
			index_from = ((self.memory_cost + i) << 5) + 16
			index_to = (self.memory_cost + i) << 4
			buffer[index_to:index_to + 16] = buffer[index_from:index_from + 16]

		return buffer

	def hash(self, password, salt=None):
		self.password = password

		if salt == None:
			self.salt = os.urandom(16)
		else:
			self.salt = salt

		#???
		buffer = pbkdf2(self.password, self.salt, itterations=1, keylength=((self.parallel_cost*self.memory_cost)<< 7), hashobj=hashlib.sha256)
		int_buffer = bytes_to_intarray(buffer, 4,  byte_order="little")

		#smix rounds
		for round_idx in range(self.parallel_cost):
			int_buffer = self._smix(int_buffer, round_num=round_idx)

		buffer = intarray_to_bytes(int_buffer, 4, byte_order="little")

		return pbkdf2(self.password, buffer, itterations=1, keylength=self.keylength, hashobj=hashlib.sha256)

if __name__ == '__main__':
	test = Scrypt(itterations=1024, memory_cost=1, parallel_cost=1, keylength=64)
	out = test.hash(password = b"correct horse battery staple", salt = b"seasalt")
	print(out.hex())
	#8dc98cddcf52dd725d52b913f7bf8386fa44e1406795aa661487f434007dff1680be6baddd724659316f7ff4663174a7a4ead1c95d5175cf284ac9ae8703e1fb
	
	test = Scrypt(itterations=1024, memory_cost=2, parallel_cost=1, keylength=64)
	out = test.hash(password = b"correct horse battery staple", salt = b"seasalt")
	print(out.hex())
	#77053f0f354002c0f2a240ce9c7b17625fb1440f87a714451217e901f7d03d748411b0bc8e4c150a573f40b98dfa816cf12bb6b01a7567970f1448d7d2a1367a

	test = Scrypt(itterations=1024, memory_cost=1, parallel_cost=2, keylength=64)
	out = test.hash(password = b"correct horse battery staple", salt = b"seasalt")
	print(out.hex())
	#cc93bb017c38aaf54901146bdc5d21c21be2314ea63ec0a4466ea44af50c8a5c87b5cd567b8205f69f601fc8ed66e4108c5b8f12474e06520de57b8fdcc484bc

	test = Scrypt(itterations=16384, memory_cost=8, parallel_cost=2, keylength=64)
	out = test.hash(password = b"correct horse battery staple", salt = b"seasalt")
	print(out.hex())
	#e3e97ec22c635ca626a6e977ae90c69845ee4c716b57e9c00757e508822fedd83d1d0539d2de1c241b830d4ce59d0bcba72d482217f193af07a125eb1c67455f