extends RefCounted

const STORAGE_VERSION: int = 2
const BLOCK_SIZE: int = 16
const KEY_CONTEXT: String = "WhaleTown-V2/session-store/v1"

func save_refresh_token(path: String, refresh_token: String) -> bool:
	if refresh_token.is_empty():
		clear(path)
		return true
	var keys := _derive_keys()
	if keys.is_empty():
		return false
	var iv := Crypto.new().generate_random_bytes(BLOCK_SIZE)
	var encryption_key: PackedByteArray = keys.get("encryption", PackedByteArray())
	var authentication_key: PackedByteArray = keys.get("authentication", PackedByteArray())
	var encrypted := _encrypt(refresh_token.to_utf8_buffer(), encryption_key, iv)
	if encrypted.is_empty():
		return false
	var authenticated_data := iv.duplicate()
	authenticated_data.append_array(encrypted)
	var mac := Crypto.new().hmac_digest(HashingContext.HASH_SHA256, authentication_key, authenticated_data)
	var payload := {
		"version": STORAGE_VERSION,
		"iv": Marshalls.raw_to_base64(iv),
		"ciphertext": Marshalls.raw_to_base64(encrypted),
		"mac": Marshalls.raw_to_base64(mac),
		"saved_at": Time.get_unix_time_from_system(),
	}
	var file := FileAccess.open(path, FileAccess.WRITE)
	if file == null:
		return false
	file.store_string(JSON.stringify(payload))
	file.close()
	return true

func load_refresh_token(path: String) -> String:
	if not FileAccess.file_exists(path):
		return ""
	var keys := _derive_keys()
	if keys.is_empty():
		clear(path)
		return ""
	var json := JSON.new()
	if json.parse(FileAccess.get_file_as_string(path)) != OK or not (json.data is Dictionary):
		clear(path)
		return ""
	var payload: Dictionary = json.data
	if int(payload.get("version", 0)) != STORAGE_VERSION:
		clear(path)
		return ""
	var iv := Marshalls.base64_to_raw(str(payload.get("iv", "")))
	var encrypted := Marshalls.base64_to_raw(str(payload.get("ciphertext", "")))
	var expected_mac := Marshalls.base64_to_raw(str(payload.get("mac", "")))
	if iv.size() != BLOCK_SIZE or encrypted.is_empty() or expected_mac.is_empty():
		clear(path)
		return ""
	var authenticated_data := iv.duplicate()
	authenticated_data.append_array(encrypted)
	var encryption_key: PackedByteArray = keys.get("encryption", PackedByteArray())
	var authentication_key: PackedByteArray = keys.get("authentication", PackedByteArray())
	var actual_mac := Crypto.new().hmac_digest(HashingContext.HASH_SHA256, authentication_key, authenticated_data)
	if not _constant_time_equals(actual_mac, expected_mac):
		clear(path)
		return ""
	var decrypted := _decrypt(encrypted, encryption_key, iv)
	return decrypted.get_string_from_utf8().strip_edges()

func clear(path: String) -> void:
	if FileAccess.file_exists(path):
		DirAccess.remove_absolute(ProjectSettings.globalize_path(path))

func _derive_keys() -> Dictionary:
	if OS.get_name() == "Web":
		return {}
	var device_id := OS.get_unique_id().strip_edges()
	if device_id.is_empty():
		return {}
	var root_key := _sha256((KEY_CONTEXT + ":" + device_id).to_utf8_buffer())
	if root_key.is_empty():
		return {}
	return {
		"encryption": _sha256(root_key + ":encryption".to_utf8_buffer()),
		"authentication": _sha256(root_key + ":authentication".to_utf8_buffer()),
	}

func _sha256(bytes: PackedByteArray) -> PackedByteArray:
	var hash_context := HashingContext.new()
	if hash_context.start(HashingContext.HASH_SHA256) != OK:
		return PackedByteArray()
	if hash_context.update(bytes) != OK:
		return PackedByteArray()
	return hash_context.finish()

func _encrypt(plain_text: PackedByteArray, key: PackedByteArray, iv: PackedByteArray) -> PackedByteArray:
	var padded := plain_text.duplicate()
	var padding_size := BLOCK_SIZE - (padded.size() % BLOCK_SIZE)
	for _index in range(padding_size):
		padded.append(padding_size)
	var aes := AESContext.new()
	if aes.start(AESContext.MODE_CBC_ENCRYPT, key, iv) != OK:
		return PackedByteArray()
	var encrypted := aes.update(padded)
	aes.finish()
	return encrypted

func _decrypt(encrypted: PackedByteArray, key: PackedByteArray, iv: PackedByteArray) -> PackedByteArray:
	var aes := AESContext.new()
	if aes.start(AESContext.MODE_CBC_DECRYPT, key, iv) != OK:
		return PackedByteArray()
	var padded := aes.update(encrypted)
	aes.finish()
	if padded.is_empty():
		return PackedByteArray()
	var padding_size := int(padded[padded.size() - 1])
	if padding_size < 1 or padding_size > BLOCK_SIZE or padding_size > padded.size():
		return PackedByteArray()
	for index in range(padded.size() - padding_size, padded.size()):
		if int(padded[index]) != padding_size:
			return PackedByteArray()
	padded.resize(padded.size() - padding_size)
	return padded

func _constant_time_equals(left: PackedByteArray, right: PackedByteArray) -> bool:
	if left.size() != right.size():
		return false
	var difference := 0
	for index in range(left.size()):
		difference |= int(left[index]) ^ int(right[index])
	return difference == 0