text: #include <stdlib.h>
  #include <stdio.h>
  #include <string.h>
  #include "aes.h"
  #include "types.h"
  #include "utils.h"
  
  /* Allocate a new context. */
  aes_ctx_t *new_aes_ctx(const void *key, unsigned int key_size, aes_mode_t mode) {
      aes_ctx_t *ctx;
      
      if ((ctx = malloc(sizeof(*ctx))) == NULL) {
          FATAL_ERROR("Failed to allocate aes_ctx_t!");
      }
  
      mbedtls_cipher_init(&ctx->cipher_dec);
      mbedtls_cipher_init(&ctx->cipher_enc);
      
      if (mbedtls_cipher_setup(&ctx->cipher_dec, mbedtls_cipher_info_from_type(mode))
          || mbedtls_cipher_setup(&ctx->cipher_enc, mbedtls_cipher_info_from_type(mode))) {
          FATAL_ERROR("Failed to set up AES context!");
      }
          
      if (mbedtls_cipher_setkey(&ctx->cipher_dec, key, key_size * 8, AES_DECRYPT)
          || mbedtls_cipher_setkey(&ctx->cipher_enc, key, key_size * 8, AES_ENCRYPT)) {
          FATAL_ERROR("Failed to set key for AES context!");
      }
      
      return ctx;
  }
  
  /* Free an allocated context. */
  void free_aes_ctx(aes_ctx_t *ctx) {
      /* Explicitly allow NULL. */
      if (ctx == NULL) {
          return;
      }
      
      mbedtls_cipher_free(&ctx->cipher_dec);
      mbedtls_cipher_free(&ctx->cipher_enc);
      free(ctx);
  }
  
  /* Set AES CTR or IV for a context. */
  void aes_setiv(aes_ctx_t *ctx, const void *iv, size_t l) {
      if (mbedtls_cipher_set_iv(&ctx->cipher_dec, iv, l)
          || mbedtls_cipher_set_iv(&ctx->cipher_enc, iv, l)) {
          FATAL_ERROR("Failed to set IV for AES context!");
      }
  }
  
  /* Calculate CMAC. */
  void aes_calculate_cmac(void *dst, void *src, size_t size, const void *key) {
      mbedtls_cipher_context_t m_ctx;
      mbedtls_cipher_init(&m_ctx);
      if (mbedtls_cipher_setup(&m_ctx, mbedtls_cipher_info_from_type(MBEDTLS_CIPHER_AES_128_ECB))
          || mbedtls_cipher_cmac_starts(&m_ctx, key, 0x80)
          || mbedtls_cipher_cmac_update(&m_ctx, src, size)
          || mbedtls_cipher_cmac_finish(&m_ctx, dst)) {
              FATAL_ERROR("Failed to calculate CMAC!");
      }
  }
  
  
  /* Encrypt with context. */
  void aes_encrypt(aes_ctx_t *ctx, void *dst, const void *src, size_t l) {
      size_t out_len = 0;
      
      /* Prepare context */
      mbedtls_cipher_reset(&ctx->cipher_enc);
      
      /* XTS doesn't need per-block updating */
      if (mbedtls_cipher_get_cipher_mode(&ctx->cipher_enc) == MBEDTLS_MODE_XTS)
          mbedtls_cipher_update(&ctx->cipher_enc, (const unsigned char * )src, l, (unsigned char *)dst, &out_len);
      else
      {
          unsigned int blk_size = mbedtls_cipher_get_block_size(&ctx->cipher_enc);
          
          /* Do per-block updating */
          for (int offset = 0; (unsigned int)offset < l; offset += blk_size)
          {
              int len = ((unsigned int)(l - offset) > blk_size) ? blk_size : (unsigned int) (l - offset);
              mbedtls_cipher_update(&ctx->cipher_enc, (const unsigned char * )src + offset, len, (unsigned char *)dst + offset, &out_len);
          }
      }
      
      /* Flush all data */
      mbedtls_cipher_finish(&ctx->cipher_enc, NULL, NULL);
  }
  
  /* Decrypt with context. */
  void aes_decrypt(aes_ctx_t *ctx, void *dst, const void *src, size_t l) 
  {
      bool src_equals_dst = false; 
  
      if (src == dst)
      {
          src_equals_dst = true;
  
          dst = malloc(l);
          if (dst == NULL) {
              fprintf(stderr, "Error: AES buffer allocation failure!\n");
              exit(EXIT_FAILURE);
          }
      }
  
      size_t out_len = 0;
      
      /* Prepare context */
      mbedtls_cipher_reset(&ctx->cipher_dec);
      
      /* XTS doesn't need per-block updating */
      if (mbedtls_cipher_get_cipher_mode(&ctx->cipher_dec) == MBEDTLS_MODE_XTS)
          mbedtls_cipher_update(&ctx->cipher_dec, (const unsigned char * )src, l, (unsigned char *)dst, &out_len);
      else
      {
          unsigned int blk_size = mbedtls_cipher_get_block_size(&ctx->cipher_dec);
          
          /* Do per-block updating */
          for (int offset = 0; (unsigned int)offset < l; offset += blk_size)
          {
              int len = ((unsigned int)(l - offset) > blk_size) ? blk_size : (unsigned int) (l - offset);
              mbedtls_cipher_update(&ctx->cipher_dec, (const unsigned char * )src + offset, len, (unsigned char *)dst + offset, &out_len);
          }
      }
      
      /* Flush all data */
      mbedtls_cipher_finish(&ctx->cipher_dec, NULL, NULL);
  
      if (src_equals_dst)
      {
          memcpy((void*)src, dst, l);
          free(dst);
      }
  }
  
  
  static void get_tweak(unsigned char *tweak, size_t sector) {
      for (int i = 0xF; i >= 0; i--) { /* Nintendo LE custom tweak... */
          tweak[i] = (unsigned char)(sector & 0xFF);
          sector >>= 8;
      }
  }
  
  /* Encrypt with context for XTS. */
  void aes_xts_encrypt(aes_ctx_t *ctx, void *dst, const void *src, size_t l, size_t sector, size_t sector_size) {
      unsigned char tweak[0x10];
  
      if (l % sector_size != 0) {
          FATAL_ERROR("Length must be multiple of sectors!");
      }
  
      for (size_t i = 0; i < l; i += sector_size) {
          /* Workaround for Nintendo's custom sector...manually generate the tweak. */
          get_tweak(tweak, sector++);
          aes_setiv(ctx, tweak, 16);
          aes_encrypt(ctx, (char *)dst + i, (const char *)src + i, sector_size);
      }
  }
  
  /* Decrypt with context for XTS. */
  void aes_xts_decrypt(aes_ctx_t *ctx, void *dst, const void *src, size_t l, size_t sector, size_t sector_size) {
      unsigned char tweak[0x10];
  
      if (l % sector_size != 0) {
          FATAL_ERROR("Length must be multiple of sectors!");
      }
  
      for (size_t i = 0; i < l; i += sector_size) {
          /* Workaround for Nintendo's custom sector...manually generate the tweak. */
          get_tweak(tweak, sector++);
          aes_setiv(ctx, tweak, 16);
          aes_decrypt(ctx, (char *)dst + i, (const char *)src + i, sector_size);
      }
  }