#   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 Crypto.Cipher import AES
from cryptography.hazmat.primitives.ciphers.aead import AESSIV

def aes_siv_encrypt(primary_master_key, hmac_master_key, plaintext, associated_data):
	"""
	Encrypt the given bytes using AES-SIV
	"""
	
	if len(plaintext) == 0:
		# Must use PyCryptodome
		# https://github.com/pyca/cryptography/issues/10958 - cryptography AESSIV does not accept empty plaintext (e.g. root directory has empty directory ID)
		if associated_data:
			if len(associated_data) > 1:
				# Incompatible with PyCryptodome
				raise ValueError('Cannot encrypt zero-length plaintext with AES-SIV with >1 associated data')
			
			if not associated_data[0]:
				# Incompatible with PyCryptodome
				raise ValueError('Cannot encrypt zero-length plaintext with AES-SIV with zero-length associated data')
			
			# If there is only one associated data, this is equivalent to the nonce, so we can use PyCryptodome
			ciphertext, tag = AES.new(hmac_master_key + primary_master_key, AES.MODE_SIV, nonce=associated_data[0]).encrypt_and_digest(plaintext)
			return tag + ciphertext
		
		# Zero-length plaintext with no AAD - encrypt with PyCryptodome
		ciphertext, tag = AES.new(hmac_master_key + primary_master_key, AES.MODE_SIV).encrypt_and_digest(plaintext)
		return tag + ciphertext
	
	# In all other cases, use cryptography AESSIV
	tag_and_ciphertext = AESSIV(hmac_master_key + primary_master_key).encrypt(plaintext, associated_data)
	return tag_and_ciphertext

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