#   cryptomator-utils: Python utilities for inspecting Cryptomator drives
#   Copyright (C) 2024  Lee Yingtong Li (RunasSudo)
#
#   This program is free software: you can redistribute it and/or modify
#   it under the terms of the GNU Affero 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 Affero General Public License for more details.
#
#   You should have received a copy of the GNU Affero General Public License
#   along with this program.  If not, see <https://www.gnu.org/licenses/>.

from cryptography.hazmat.primitives.ciphers.aead import AESSIV
from cryptography.hazmat.primitives.ciphers.algorithms import AES
from cryptography.hazmat.primitives.cmac import CMAC

import struct

def aes_siv_encrypt(primary_master_key, hmac_master_key, plaintext, associated_data):
	"""
	Encrypt the given bytes using AES-SIV
	"""
	
	if len(plaintext) == 0:
		# https://github.com/pyca/cryptography/issues/10958 - cryptography AESSIV does not accept empty plaintext (e.g. root directory has empty directory ID)
		
		# Manually calculate the synthetic IV
		siv = aes_siv_s2v(hmac_master_key, plaintext, associated_data)
		
		# Empty plaintext equals empty ciphertext, so result is just the SIV
		return siv
	
	# In all other cases, use cryptography AESSIV
	siv_and_ciphertext = AESSIV(hmac_master_key + primary_master_key).encrypt(plaintext, associated_data)
	return siv_and_ciphertext

def aes_siv_decrypt(primary_master_key, hmac_master_key, siv_and_ciphertext, associated_data):
	"""
	Decrypt the given AES-SIV ciphertext
	"""
	
	return AESSIV(hmac_master_key + primary_master_key).decrypt(siv_and_ciphertext, associated_data)

def aes_cmac(hmac_master_key, data):
	mac = CMAC(AES(hmac_master_key))
	mac.update(data)
	return mac.finalize()

def aes_siv_dbl(data):
	# Based on miscreant.py by Phil Rogaway, MIT License
	
	overflow = 0
	words = struct.unpack(b'!LLLL', data)
	output_words = []
	
	for word in reversed(words):
		new_word = (word << 1) & 0xFFFFFFFF
		new_word |= overflow
		overflow = int((word & 0x80000000) >= 0x80000000)
		output_words.append(new_word)
	
	result = bytearray(struct.pack(b'!LLLL', *reversed(output_words)))
	if overflow:
		result[-1] ^= 0x87  # Foot-gun! Not constant time
	
	return result

def aes_siv_s2v(hmac_master_key, plaintext, associated_data):
	# Based on miscreant.py by Phil Rogaway, MIT License
	
	# Note: The standalone S2V returns CMAC(1) if the number of passed vectors is zero, however in SIV construction this case is never triggered, since we always pass plaintext as the last vector, so we omit this case.
	
	d = bytes(128//8)  # 128-bit blocks
	d = aes_cmac(hmac_master_key, d)
	
	if associated_data:
		for ad in associated_data:
			d = aes_siv_dbl(d)
			d = bytes(x ^ y for x, y in zip(d, aes_cmac(hmac_master_key, ad)))  # d ^= aes_cmac(hmac_master_key, ad)
	
	if len(plaintext) >= 128//8:  # 128 bits
		mac = CMAC(AES(hmac_master_key))
		difference = len(plaintext) - 128/8
		mac.update(plaintext[:difference])
		d = bytes(x ^ y for x, y in zip(d, plaintext[difference:]))  # d ^= plaintext[difference:]
		mac.update(d)
		return mac.finalize()
	
	d = aes_siv_dbl(d)
	for i in range(len(plaintext)):
		d[i] ^= plaintext[i]
	d[len(plaintext)] ^= 0x80
	return aes_cmac(hmac_master_key, d)