aes.c

#include <stdlib.h>
#include <stdio.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!");
    }
}

/* 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) {
    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);
}

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, (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, (char *)src + i, sector_size);
    }
}