package romulus_go

// Converted to go with C2GO, tweaks by 51m0n - 2022.

/*
 * Date: 05 May 2021
 * Contact: Romulus Team (Mustafa Khairallah - mustafa.khairallah@ntu.edu.sg)
 * Romulus-M as compliant with the Romulus v1.3 specifications.
 * This file icludes the functions of Romulus-N
 * It superseeds earlier versions developed by Mustafa Khairallah and maintained
 * by Mustafa Khairallah, Thomas Peyrin and Kazuhiko Minematsu
 */

// Padding function: pads the byte length of the message mod 16 to the last incomplete block.
// For complete blocks it returns the same block.
func pad(m []byte, mp []byte, l int, len8 int) {
	var i int

	for i = 0; i < l; i++ {
		if i < len8 {
			mp[i] = m[i]
		} else if i == l-1 {
			mp[i] = byte(len8 & 0x0f)
		} else {
			mp[i] = 0x00
		}
	}
}

// G(S): generates the key stream from the internal state by multiplying the state S by the constant matrix G
func g8A(s []byte, c []byte) {
	var i int

	for i = 0; i < 16; i++ {
		c[i] = (s[i] >> 1) ^ (s[i] & 0x80) ^ ((s[i] & 0x01) << 7)
	}
}

// Rho(S,A) pads an A block and XORs it to the internal state.
func rho_ad(m []byte, s []byte, len8 int, ver int) {
	var i int
	var mp [16]byte

	pad(m, mp[:], ver, len8)
	for i = 0; i < ver; i++ {
		s[i] = s[i] ^ mp[i]
	}
}

// Rho(S,M): pads an M block and outputs S'= M xor S and C = M xor G(S)
func rho(m []byte, c []byte, s []byte, len8 int, ver int) {
	var i int
	var mp [16]byte

	pad(m, mp[:], ver, len8)

	g8A(s, c)
	for i = 0; i < ver; i++ {
		s[i] = s[i] ^ mp[i]
		if i < len8 {
			c[i] = c[i] ^ mp[i]
		} else {
			c[i] = 0
		}
	}
}

// Inverse-Rho(S,M): pads a C block and outputs S'= C xor G(S) xor S and M = C xor G(S)
func irho(m []byte, c []byte, s []byte, len8 int, ver int) {
	var i int
	var cp [16]byte

	pad(c, cp[:], ver, len8)

	g8A(s, m)
	for i = 0; i < ver; i++ {
		if i < len8 {
			s[i] = s[i] ^ cp[i] ^ m[i]
		} else {
			s[i] = s[i] ^ cp[i]
		}

		if i < len8 {
			m[i] = m[i] ^ cp[i]
		} else {
			m[i] = 0
		}
	}
}

// Resets the value of the counter.
func reset_lfsr_gf56(CNT []byte) {
	CNT[0] = 0x01
	CNT[1] = 0x00
	CNT[2] = 0x00
	CNT[3] = 0x00
	CNT[4] = 0x00
	CNT[5] = 0x00
	CNT[6] = 0x00
}

// Applies CNT'=2 * CNT (mod GF(2^56)), where GF(2^56) is defined using the irreducible polynomial
// x^56 + x^7 + x^4 + x^2 + 1
func lfsr_gf56(CNT []byte) {
	var fb0 byte

	fb0 = CNT[6] >> 7

	CNT[6] = CNT[6]<<1 | CNT[5]>>7
	CNT[5] = CNT[5]<<1 | CNT[4]>>7
	CNT[4] = CNT[4]<<1 | CNT[3]>>7
	CNT[3] = CNT[3]<<1 | CNT[2]>>7
	CNT[2] = CNT[2]<<1 | CNT[1]>>7
	CNT[1] = CNT[1]<<1 | CNT[0]>>7
	if fb0 == 1 {
		CNT[0] = (CNT[0] << 1) ^ 0x95
	} else {
		CNT[0] = CNT[0] << 1
	}
}

// Combines the secret key, nonce (or A block), counter and domain bits to form the full 384-bit tweakey
func compose_tweakey(KT []byte, K []byte, T []byte, CNT []byte, D byte, t int) {
	var i int

	for i = 0; i < 7; i++ {
		KT[i] = CNT[i]
	}

	KT[i] = D
	for i = 8; i < 16; i++ {
		KT[i] = 0x00
	}

	for i = 0; i < t; i++ {
		KT[i+16] = T[i]
	}

	for i = 0; i < 16; i++ {
		KT[i+16+t] = K[i]
	}
}

// An interface between Romulus and the underlying TBC
func block_cipher(s []byte, k []byte, T []byte, CNT []byte, D byte, t int) {
	var KT [48]byte

	compose_tweakey(KT[:], k, T, CNT, D, t)
	skinny_128_384_plus_enc(s, KT[:])
}

// Calls the TBC using the nonce as part of the tweakey
func nonce_encryption(N []byte, CNT []byte, s []byte, k []byte, t int, D byte) {
	var T [16]byte
	var i int
	for i = 0; i < t; i++ {
		T[i] = N[i]
	}

	block_cipher(s, k, T[:], CNT, D, t)
}

// Generates the tag T from the final state S by applying T=G(S).
func generate_tag(c *[]byte, s []byte, n int, clen *uint64) {
	g8A(s, *c)
	// For some reason im going to assert that n-clen is positive
	*c = (*c)[(uint64)(n)-*clen:]
	//*c = *c - *clen
}

// Absorbs and encrypts the message blocks.
func msg_encryption(M *[]byte, c *[]byte, N []byte, CNT []byte, s []byte, k []byte, n uint, t uint, D byte, mlen uint64) uint64 {
	var len8 int

	if mlen >= uint64(n) {
		len8 = int(n)
		mlen = mlen - uint64(n)
	} else {
		len8 = int(mlen)
		mlen = 0
	}

	rho(*M, *c, s, len8, int(n))
	*c = (*c)[len8:]
	*M = (*M)[len8:]
	lfsr_gf56(CNT)
	nonce_encryption(N, CNT, s, k, int(t), D)
	return mlen
}

// Absorbs and decrypts the ciphertext blocks.
func msg_decryption(M *[]byte, c *[]byte, N []byte, CNT []byte, s []byte, k []byte, n uint, t uint, D byte, clen uint64) uint64 {
	var len8 int

	if clen >= uint64(n) {
		len8 = int(n)
		clen = clen - uint64(n)
	} else {
		len8 = int(clen)
		clen = 0
	}

	irho(*M, *c, s, len8, int(n))
	*c = (*c)[len8:]
	*M = (*M)[len8:]
	lfsr_gf56(CNT)
	nonce_encryption(N, CNT, s, k, int(t), D)
	return clen
}

// Handles the special case when the number of blocks of A is odd
func ad2msg_encryption(M *[]byte, CNT []byte, s []byte, k []byte, t uint, D byte, mlen uint64) uint64 {
	var T [16]byte
	var len8 int

	if mlen <= uint64(t) {
		len8 = int(mlen)
		mlen = 0
	} else {
		len8 = int(t)
		mlen = mlen - uint64(t)
	}

	pad(*M, T[:], int(t), len8)

	block_cipher(s, k, T[:], CNT, D, int(t))
	lfsr_gf56(CNT)
	*M = (*M)[len8:]

	return mlen
}

// Absorbs the AD blocks.
func ad_encryption(A *[]byte, s []byte, k []byte, adlen uint64, CNT []byte, D byte, n uint, t uint) uint64 {
	var T [16]byte
	var len8 int

	if adlen >= uint64(n) {
		len8 = int(n)
		adlen = adlen - uint64(n)
	} else {
		len8 = int(adlen)
		adlen = 0
	}

	rho_ad(*A, s, len8, int(n))
	*A = (*A)[len8:]
	lfsr_gf56(CNT)

	if adlen != 0 {
		if adlen >= uint64(t) {
			len8 = int(t)
			adlen = adlen - uint64(t)
		} else {
			len8 = int(adlen)
			adlen = 0
		}

		pad(*A, T[:], int(t), len8)
		*A = (*A)[len8:]
		block_cipher(s, k, T[:], CNT, D, int(t))
		lfsr_gf56(CNT)
	}

	return adlen
}

func romulus_m_encrypt(c []byte, clen *uint64, m []byte, mlen uint64, ad []byte, adlen uint64, nsec []byte, npub []byte, k []byte) int {
	var s [16]byte
	var CNT [7]byte
	var T [16]byte
	var N []byte
	var n uint
	var t uint
	var i uint
	var w byte
	var xlen uint64

	N = npub
	// sm
	mstart := m[:]
	cstart := c[:]
	n = 16
	t = 16

	xlen = mlen

	for i = 0; i < n; i++ {
		s[i] = 0
	}

	reset_lfsr_gf56(CNT[:])

	// Calculating the domain separation bits for the last block MAC TBC call depending on the length of M and AD
	w = 48

	if adlen == 0 {
		w = w ^ 2
		if xlen == 0 {
			w = w ^ 1
		} else if xlen%uint64(n+t) == 0 {
			w = w ^ 4
		} else if xlen%uint64(n+t) < uint64(t) {
			w = w ^ 1
		} else if xlen%uint64(n+t) == uint64(t) {
			w = w ^ 0
		} else {
			w = w ^ 5
		}
	} else if adlen%uint64(n+t) == 0 {
		w = w ^ 8
		if xlen == 0 {
			w = w ^ 1
		} else if xlen%uint64(n+t) == 0 {
			w = w ^ 4
		} else if xlen%uint64(n+t) < uint64(n) {
			w = w ^ 1
		} else if xlen%uint64(n+t) == uint64(n) {
			w = w ^ 0
		} else {
			w = w ^ 5
		}
	} else if adlen%uint64(n+t) < uint64(n) {
		w = w ^ 2
		if xlen == 0 {
			w = w ^ 1
		} else if xlen%uint64(n+t) == 0 {
			w = w ^ 4
		} else if xlen%uint64(n+t) < uint64(t) {
			w = w ^ 1
		} else if xlen%uint64(n+t) == uint64(t) {
			w = w ^ 0
		} else {
			w = w ^ 5
		}
	} else if adlen%uint64(n+t) == uint64(n) {
		w = w ^ 0
		if xlen == 0 {
			w = w ^ 1
		} else if xlen%uint64(n+t) == 0 {
			w = w ^ 4
		} else if xlen%uint64(n+t) < uint64(t) {
			w = w ^ 1
		} else if xlen%uint64(n+t) == uint64(t) {
			w = w ^ 0
		} else {
			w = w ^ 5
		}
	} else {
		w = w ^ 10
		if xlen == 0 {
			w = w ^ 1
		} else if xlen%uint64(n+t) == 0 {
			w = w ^ 4
		} else if xlen%uint64(n+t) < uint64(n) {
			w = w ^ 1
		} else if xlen%uint64(n+t) == uint64(n) {
			w = w ^ 0
		} else {
			w = w ^ 5
		}
	}

	if adlen == 0 { // AD is an empty string
		lfsr_gf56(CNT[:])
	} else {
		for adlen > 0 {
			adlen = ad_encryption(&ad, s[:], k, adlen, CNT[:], 40, n, t)
		}
	}

	if w&8 == 0 {
		xlen = ad2msg_encryption(&m, CNT[:], s[:], k, t, 44, xlen)
	} else if mlen == 0 {
		lfsr_gf56(CNT[:])
	}

	for xlen > 0 {
		xlen = ad_encryption(&m, s[:], k, xlen, CNT[:], 44, n, t)
	}

	nonce_encryption(N, CNT[:], s[:], k, int(t), w)

	// Tag generation
	g8A(s[:], T[:])

	m = mstart

	reset_lfsr_gf56(CNT[:])

	for i = 0; i < n; i = i + 1 {
		s[i] = T[i]
	}

	n = 16
	*clen = mlen + uint64(n)

	if mlen > 0 {
		nonce_encryption(N, CNT[:], s[:], k, int(t), 36)
		for mlen > uint64(n) {
			mlen = msg_encryption(&m, &c, N, CNT[:], s[:], k, n, t, 36, mlen)
		}

		rho(m, c, s[:], int(mlen), 16)
		c = c[mlen:]
		m = m[mlen:]
	}

	// Tag Concatenation
	for i = 0; i < 16; i = i + 1 {
		(c[i:])[0] = T[i]
	}

	c = cstart

	return 0
}

func romulus_m_decrypt(m []byte, mlen *uint64, nsec []byte, c []byte, clen uint64, ad []byte, adlen uint64, npub []byte, k []byte) int {
	var s [16]byte
	var CNT [7]byte
	var T [16]byte
	var N []byte
	var n uint
	var t uint
	var i uint
	var w byte
	var xlen uint64
	var mauth []byte

	mauth = m

	N = npub

	n = 16
	t = 16

	xlen = clen - 16

	reset_lfsr_gf56(CNT[:])

	copy(T[:], c[len(c)-len(T):])

	for i = 0; i < n; i = i + 1 {
		s[i] = T[i]
	}

	n = 16
	clen = clen - 16
	*mlen = clen

	if clen > 0 {
		nonce_encryption(N, CNT[:], s[:], k, int(t), 36)
		for clen > uint64(n) {
			clen = msg_decryption(&m, &c, N, CNT[:], s[:], k, n, t, 36, clen)
		}

		irho(m, c, s[:], int(clen), 16)
		c = c[clen:]
		m = m[clen:]
	}

	for i = 0; i < n; i++ {
		s[i] = 0
	}

	reset_lfsr_gf56(CNT[:])

	// Calculating the domain separation bits for the last block MAC TBC call depending on the length of M and AD
	w = 48

	if adlen == 0 {
		w = w ^ 2
		if xlen == 0 {
			w = w ^ 1
		} else if xlen%uint64(n+t) == 0 {
			w = w ^ 4
		} else if xlen%uint64(n+t) < uint64(t) {
			w = w ^ 1
		} else if xlen%uint64(n+t) == uint64(t) {
			w = w ^ 0
		} else {
			w = w ^ 5
		}
	} else if adlen%uint64(n+t) == 0 {
		w = w ^ 8
		if xlen == 0 {
			w = w ^ 1
		} else if xlen%uint64(n+t) == 0 {
			w = w ^ 4
		} else if xlen%uint64(n+t) < uint64(n) {
			w = w ^ 1
		} else if xlen%uint64(n+t) == uint64(n) {
			w = w ^ 0
		} else {
			w = w ^ 5
		}
	} else if adlen%uint64(n+t) < uint64(n) {
		w = w ^ 2
		if xlen == 0 {
			w = w ^ 1
		} else if xlen%uint64(n+t) == 0 {
			w = w ^ 4
		} else if xlen%uint64(n+t) < uint64(t) {
			w = w ^ 1
		} else if xlen%uint64(n+t) == uint64(t) {
			w = w ^ 0
		} else {
			w = w ^ 5
		}
	} else if adlen%uint64(n+t) == uint64(n) {
		w = w ^ 0
		if xlen == 0 {
			w = w ^ 1
		} else if xlen%uint64(n+t) == 0 {
			w = w ^ 4
		} else if xlen%uint64(n+t) < uint64(t) {
			w = w ^ 1
		} else if xlen%uint64(n+t) == uint64(t) {
			w = w ^ 0
		} else {
			w = w ^ 5
		}
	} else {
		w = w ^ 10
		if xlen == 0 {
			w = w ^ 1
		} else if xlen%uint64(n+t) == 0 {
			w = w ^ 4
		} else if xlen%uint64(n+t) < uint64(n) {
			w = w ^ 1
		} else if xlen%uint64(n+t) == uint64(n) {
			w = w ^ 0
		} else {
			w = w ^ 5
		}
	}

	if adlen == 0 { // AD is an empty string
		lfsr_gf56(CNT[:])
	} else {
		for adlen > 0 {
			adlen = ad_encryption(&ad, s[:], k, adlen, CNT[:], 40, n, t)
		}
	}

	if w&8 == 0 {
		xlen = ad2msg_encryption(&mauth, CNT[:], s[:], k, t, 44, xlen)
	} else if clen == 0 {
		lfsr_gf56(CNT[:])
	}

	for xlen > 0 {
		xlen = ad_encryption(&mauth, s[:], k, xlen, CNT[:], 44, n, t)
	}

	nonce_encryption(N, CNT[:], s[:], k, int(t), w)

	// Tag generation
	g8A(s[:], T[:])

	// Tag verification
	for i = 0; i < 16; i++ {
		if T[i] != ((c[i:])[0]) {
			return -1
		}
	}

	return 0
}