#include <stdio.h> 
#include <stdlib.h> 
#include <string.h>
#include <assert.h>
#include "mbedtls/aes.h"

const uint8_t encrypted_buffer[64] = {
	0x16, 0xfe, 0x96, 0xe1, 0xb9, 0x54, 0x64, 0xe4, 0x16, 0x64, 0x24, 0xa7, 0x41, 0xf0, 0xd9, 0xa6,
	0x39, 0x76, 0xcf, 0xc3, 0xb4, 0xc9, 0x40, 0x72, 0x97, 0xfe, 0x97, 0x7c, 0xb6, 0x99, 0x72, 0x2b, 
	0x5e, 0xd5, 0x73, 0xbf, 0xbb, 0xa6, 0xd2, 0xd0, 0x84, 0x67, 0xa8, 0x35, 0xf8, 0x60, 0x9e, 0x02, 
	0xce, 0x9b, 0x72, 0x0d, 0x00, 0xf3, 0x50, 0xe5, 0xdc, 0x99, 0x25, 0x2d, 0x72, 0x33, 0x37, 0x0e, 
};

const uint8_t decrypted_buffer[64] = {
	0x54, 0x65, 0x78, 0x74, 0x5f, 0x36, 0x37, 0x38, 0x39, 0x30, 0x5f, 0x54, 0x65, 0x78, 0x74, 0x5f,
	0x36, 0x37, 0x38, 0x39, 0x30, 0x5f, 0x54, 0x65, 0x78, 0x74, 0x5f, 0x36, 0x37, 0x38, 0x39, 0x30,
	0x5f, 0x54, 0x65, 0x78, 0x74, 0x5f, 0x36, 0x37, 0x38, 0x39, 0x30, 0x5f, 0x54, 0x65, 0x78, 0x74,
	0x5f, 0x36, 0x37, 0x38, 0x39, 0x30, 0x5f, 0x54, 0x65, 0x78, 0x74, 0x5f, 0x36, 0x37, 0x38, 0x39,
};

const uint8_t tweak_key[32] = {
	0x9c, 0x8e, 0x2b, 0xa7, 0x51, 0x37, 0x01, 0x76, 0x8c, 0x3c, 0x22, 0xdb, 0x4b, 0x0f, 0xe0, 0x56,
	0xda, 0xf8, 0xab, 0xfc, 0x9a, 0xcf, 0xf4, 0x3a, 0x6f, 0x93, 0x1d, 0x83, 0x73, 0xac, 0x1f, 0x93,
};

const uint8_t crypt_key[32] = {
	0xff, 0x69, 0x89, 0xe4, 0xfa, 0xfd, 0x2d, 0x8c, 0xb3, 0x13, 0xaa, 0x57, 0x43, 0xf0, 0x17, 0x3f,
	0x9b, 0x8a, 0x2c, 0x85, 0x87, 0x87, 0x28, 0x32, 0x13, 0x46, 0x82, 0xe0, 0xe0, 0x53, 0x8d, 0xbd,
};

int decrpyt_bytes(mbedtls_aes_xts_context *ctx, uint8_t *input, uint8_t *output, size_t len, uint32_t offset)
{
	uint8_t iv[16];
	memset(iv, 0, sizeof(iv));
    memcpy(iv, &offset, sizeof(offset));

	return mbedtls_aes_crypt_xts(ctx, MBEDTLS_AES_DECRYPT, len, iv, input, output);
}

int main()
{
	uint8_t xts_aes_key[64] = { 0 };
	memcpy(xts_aes_key, crypt_key, 32);
	memcpy(xts_aes_key + 32, tweak_key, 32);

	uint8_t input_buffer[64] = { 0 };
	uint8_t output_buffer[64] = { 0 };

	memcpy(input_buffer, encrypted_buffer, 64);
	memset(output_buffer, 0, 64);

	mbedtls_aes_xts_context xts_aes_ctx;

	assert(mbedtls_aes_xts_setkey_dec(&xts_aes_ctx, xts_aes_key, 64 * 8) == 0);

	// decrypting 64 bytes at once
	uint32_t offset = 288;
	assert(decrpyt_bytes(&xts_aes_ctx, input_buffer, output_buffer, 64, offset) == 0);

	if (memcmp(output_buffer, decrypted_buffer, 64) != 0)
		printf("Decrypting all 64 bytes at once fail\n");    // FAIL
	else
		printf("Decrypting all 64 bytes at once pass\n");

	memset(output_buffer, 0, 64);

	// decrypting the first 32 bytes at once
	offset = 288;
	assert(decrpyt_bytes(&xts_aes_ctx, input_buffer, output_buffer, 32, offset) == 0);

	// next 32 bytes
	offset = 288 + 32;
	assert(decrpyt_bytes(&xts_aes_ctx, input_buffer + 32, output_buffer + 32, 32, offset) == 0);

	if (memcmp(output_buffer, decrypted_buffer, 64) != 0)
		printf("Decrypting 32 bytes at once fail\n");
	else
		printf("Decrypting 32 bytes at once pass\n");    // PASS

	return 0;
}