package main

// Self-contained scrypt (RFC 7914) with Node's crypto.scryptSync defaults:
// N=16384, r=8, p=1, keyLen=64. The stdlib has no scrypt, so PBKDF2-HMAC-SHA256,
// the Salsa20/8 core, BlockMix and ROMix are implemented here from the spec.
//
// Important Node quirk: crypto.scryptSync(secret, salt, 64) takes the salt as
// a string and uses its raw UTF-8 bytes — the hex-looking salt produced by
// HashSecret is NOT hex-decoded before use.

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/binary"
	"math/bits"
)

// ScryptHash derives 64 bytes from secret and salt with N=16384, r=8, p=1.
func ScryptHash(secret, salt string) []byte {
	const (
		n      = 16384
		r      = 8
		p      = 1
		keyLen = 64
	)
	password := []byte(secret)
	saltBytes := []byte(salt)
	b := pbkdf2SHA256(password, saltBytes, 1, p*128*r)
	// Work in uint32 words for ROMix/BlockMix.
	words := make([]uint32, len(b)/4)
	for i := range words {
		words[i] = binary.LittleEndian.Uint32(b[i*4:])
	}
	blockWords := 32 * r
	for i := 0; i < p; i++ {
		romix(words[i*blockWords:(i+1)*blockWords], n, r)
	}
	for i := range words {
		binary.LittleEndian.PutUint32(b[i*4:], words[i])
	}
	return pbkdf2SHA256(password, b, 1, keyLen)
}

func pbkdf2SHA256(password, salt []byte, iter, keyLen int) []byte {
	hashLen := sha256.Size
	numBlocks := (keyLen + hashLen - 1) / hashLen
	dk := make([]byte, 0, numBlocks*hashLen)
	var counter [4]byte
	for block := 1; block <= numBlocks; block++ {
		mac := hmac.New(sha256.New, password)
		mac.Write(salt)
		binary.BigEndian.PutUint32(counter[:], uint32(block))
		mac.Write(counter[:])
		u := mac.Sum(nil)
		t := make([]byte, len(u))
		copy(t, u)
		for i := 1; i < iter; i++ {
			mac = hmac.New(sha256.New, password)
			mac.Write(u)
			u = mac.Sum(nil)
			for j := range t {
				t[j] ^= u[j]
			}
		}
		dk = append(dk, t...)
	}
	return dk[:keyLen]
}

// salsa208Core applies the Salsa20/8 core function in place to 16 words.
func salsa208Core(x *[16]uint32) {
	orig := *x
	for i := 0; i < 8; i += 2 {
		// Column round.
		x[4] ^= bits.RotateLeft32(x[0]+x[12], 7)
		x[8] ^= bits.RotateLeft32(x[4]+x[0], 9)
		x[12] ^= bits.RotateLeft32(x[8]+x[4], 13)
		x[0] ^= bits.RotateLeft32(x[12]+x[8], 18)

		x[9] ^= bits.RotateLeft32(x[5]+x[1], 7)
		x[13] ^= bits.RotateLeft32(x[9]+x[5], 9)
		x[1] ^= bits.RotateLeft32(x[13]+x[9], 13)
		x[5] ^= bits.RotateLeft32(x[1]+x[13], 18)

		x[14] ^= bits.RotateLeft32(x[10]+x[6], 7)
		x[2] ^= bits.RotateLeft32(x[14]+x[10], 9)
		x[6] ^= bits.RotateLeft32(x[2]+x[14], 13)
		x[10] ^= bits.RotateLeft32(x[6]+x[2], 18)

		x[3] ^= bits.RotateLeft32(x[15]+x[11], 7)
		x[7] ^= bits.RotateLeft32(x[3]+x[15], 9)
		x[11] ^= bits.RotateLeft32(x[7]+x[3], 13)
		x[15] ^= bits.RotateLeft32(x[11]+x[7], 18)

		// Row round.
		x[1] ^= bits.RotateLeft32(x[0]+x[3], 7)
		x[2] ^= bits.RotateLeft32(x[1]+x[0], 9)
		x[3] ^= bits.RotateLeft32(x[2]+x[1], 13)
		x[0] ^= bits.RotateLeft32(x[3]+x[2], 18)

		x[6] ^= bits.RotateLeft32(x[5]+x[4], 7)
		x[7] ^= bits.RotateLeft32(x[6]+x[5], 9)
		x[4] ^= bits.RotateLeft32(x[7]+x[6], 13)
		x[5] ^= bits.RotateLeft32(x[4]+x[7], 18)

		x[11] ^= bits.RotateLeft32(x[10]+x[9], 7)
		x[8] ^= bits.RotateLeft32(x[11]+x[10], 9)
		x[9] ^= bits.RotateLeft32(x[8]+x[11], 13)
		x[10] ^= bits.RotateLeft32(x[9]+x[8], 18)

		x[12] ^= bits.RotateLeft32(x[15]+x[14], 7)
		x[13] ^= bits.RotateLeft32(x[12]+x[15], 9)
		x[14] ^= bits.RotateLeft32(x[13]+x[12], 13)
		x[15] ^= bits.RotateLeft32(x[14]+x[13], 18)
	}
	for i := range x {
		x[i] += orig[i]
	}
}

// blockMix performs scryptBlockMix over 2r 64-byte blocks held as uint32
// words, writing the result into out (which must not alias in). tmp is the
// running 64-byte block.
func blockMix(in, out []uint32, r int, tmp *[16]uint32) {
	copy(tmp[:], in[(2*r-1)*16:(2*r)*16])
	for i := 0; i < 2*r; i++ {
		block := in[i*16 : (i+1)*16]
		for j := 0; j < 16; j++ {
			tmp[j] ^= block[j]
		}
		salsa208Core(tmp)
		// Y_i goes to the even blocks first, then the odd blocks.
		dest := i / 2
		if i%2 == 1 {
			dest = dest + r
		}
		copy(out[dest*16:(dest+1)*16], tmp[:])
	}
}

// romix is scryptROMix over one 128*r-byte chunk held as uint32 words.
func romix(b []uint32, n, r int) {
	blockWords := 32 * r
	v := make([]uint32, n*blockWords)
	x := make([]uint32, blockWords)
	y := make([]uint32, blockWords)
	copy(x, b)
	for i := 0; i < n; i++ {
		copy(v[i*blockWords:], x)
		blockMix(x, y, r, &[16]uint32{})
		copy(x, y)
	}
	for i := 0; i < n; i++ {
		// Integerify: first 64 bits of the last 64-byte block, little-endian.
		pos := (2*r - 1) * 16
		j := int((uint64(x[pos]) | uint64(x[pos+1])<<32) % uint64(n))
		vj := v[j*blockWords:]
		for k := range x {
			x[k] ^= vj[k]
		}
		blockMix(x, y, r, &[16]uint32{})
		copy(x, y)
	}
	copy(b, x)
}