1
0
Fork 0
mirror of https://github.com/DarkMatterCore/nxdumptool.git synced 2024-11-29 21:52:22 +00:00
nxdumptool/source/aes.c
2019-05-01 16:24:13 -04:00

325 lines
10 KiB
C

#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include "aes.h"
#include "ui.h"
#include "util.h"
/* Extern variables */
extern int breaks;
extern int font_height;
extern char strbuf[NAME_BUF_LEN * 4];
/* Allocate a new context. */
aes_ctx_t *new_aes_ctx(const void *key, unsigned int key_size, aes_mode_t mode)
{
int ret;
aes_ctx_t *ctx;
if ((ctx = malloc(sizeof(*ctx))) == NULL)
{
uiDrawString("Error: failed to allocate aes_ctx_t!", 0, breaks * font_height, 255, 0, 0);
return NULL;
}
mbedtls_cipher_init(&ctx->cipher_dec);
mbedtls_cipher_init(&ctx->cipher_enc);
ret = mbedtls_cipher_setup(&ctx->cipher_dec, mbedtls_cipher_info_from_type(mode));
if (ret)
{
free_aes_ctx(ctx);
snprintf(strbuf, sizeof(strbuf) / sizeof(strbuf[0]), "Error: failed to set up AES decryption context! (%d)", ret);
uiDrawString(strbuf, 0, breaks * font_height, 255, 0, 0);
return NULL;
}
ret = mbedtls_cipher_setup(&ctx->cipher_enc, mbedtls_cipher_info_from_type(mode));
if (ret)
{
free_aes_ctx(ctx);
snprintf(strbuf, sizeof(strbuf) / sizeof(strbuf[0]), "Error: failed to set up AES encryption context! (%d)", ret);
uiDrawString(strbuf, 0, breaks * font_height, 255, 0, 0);
return NULL;
}
ret = mbedtls_cipher_setkey(&ctx->cipher_dec, key, key_size * 8, AES_DECRYPT);
if (ret)
{
free_aes_ctx(ctx);
snprintf(strbuf, sizeof(strbuf) / sizeof(strbuf[0]), "Error: failed to set key for AES decryption context! (%d)", ret);
uiDrawString(strbuf, 0, breaks * font_height, 255, 0, 0);
return NULL;
}
ret = mbedtls_cipher_setkey(&ctx->cipher_enc, key, key_size * 8, AES_ENCRYPT);
if (ret)
{
free_aes_ctx(ctx);
snprintf(strbuf, sizeof(strbuf) / sizeof(strbuf[0]), "Error: failed to set key for AES encryption context! (%d)", ret);
uiDrawString(strbuf, 0, breaks * font_height, 255, 0, 0);
return NULL;
}
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. */
int aes_setiv(aes_ctx_t *ctx, const void *iv, size_t l)
{
int ret;
ret = mbedtls_cipher_set_iv(&ctx->cipher_dec, iv, l);
if (ret)
{
snprintf(strbuf, sizeof(strbuf) / sizeof(strbuf[0]), "Error: failed to set IV for AES decryption context! (%d)", ret);
uiDrawString(strbuf, 0, breaks * font_height, 255, 0, 0);
return 0;
}
ret = mbedtls_cipher_set_iv(&ctx->cipher_enc, iv, l);
if (ret)
{
snprintf(strbuf, sizeof(strbuf) / sizeof(strbuf[0]), "Error: failed to set IV for AES encryption context! (%d)", ret);
uiDrawString(strbuf, 0, breaks * font_height, 255, 0, 0);
return 0;
}
return 1;
}
/* Calculate CMAC. */
int aes_calculate_cmac(void *dst, void *src, size_t size, const void *key)
{
int ret;
mbedtls_cipher_context_t m_ctx;
mbedtls_cipher_init(&m_ctx);
ret = mbedtls_cipher_setup(&m_ctx, mbedtls_cipher_info_from_type(MBEDTLS_CIPHER_AES_128_ECB));
if (ret)
{
mbedtls_cipher_free(&m_ctx);
snprintf(strbuf, sizeof(strbuf) / sizeof(strbuf[0]), "Error: failed to set up AES context for CMAC calculation! (%d)", ret);
uiDrawString(strbuf, 0, breaks * font_height, 255, 0, 0);
return 0;
}
ret = mbedtls_cipher_cmac_starts(&m_ctx, key, 0x80);
if (ret)
{
mbedtls_cipher_free(&m_ctx);
snprintf(strbuf, sizeof(strbuf) / sizeof(strbuf[0]), "Error: failed to start CMAC calculation! (%d)", ret);
uiDrawString(strbuf, 0, breaks * font_height, 255, 0, 0);
return 0;
}
ret = mbedtls_cipher_cmac_update(&m_ctx, src, size);
if (ret)
{
mbedtls_cipher_free(&m_ctx);
snprintf(strbuf, sizeof(strbuf) / sizeof(strbuf[0]), "Error: failed to update CMAC calculation! (%d)", ret);
uiDrawString(strbuf, 0, breaks * font_height, 255, 0, 0);
return 0;
}
ret = mbedtls_cipher_cmac_finish(&m_ctx, dst);
if (ret)
{
mbedtls_cipher_free(&m_ctx);
snprintf(strbuf, sizeof(strbuf) / sizeof(strbuf[0]), "Error: failed to finish CMAC calculation! (%d)", ret);
uiDrawString(strbuf, 0, breaks * font_height, 255, 0, 0);
return 0;
}
mbedtls_cipher_free(&m_ctx);
return 1;
}
/* Encrypt with context. */
int aes_encrypt(aes_ctx_t *ctx, void *dst, const void *src, size_t l)
{
int ret;
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)
{
ret = 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));
ret = mbedtls_cipher_update(&ctx->cipher_enc, (const unsigned char *)src + offset, len, (unsigned char *)dst + offset, &out_len);
if (ret) break;
}
}
if (ret)
{
snprintf(strbuf, sizeof(strbuf) / sizeof(strbuf[0]), "Error: AES encryption failed! (%d)", ret);
uiDrawString(strbuf, 0, breaks * font_height, 255, 0, 0);
return 0;
}
/* Flush all data */
size_t strbuf_size = sizeof(strbuf);
ret = mbedtls_cipher_finish(&ctx->cipher_enc, (unsigned char *)strbuf, &strbuf_size); // Looks ugly, but using NULL,NULL with mbedtls on Switch is a no-no
if (ret)
{
snprintf(strbuf, sizeof(strbuf) / sizeof(strbuf[0]), "Error: failed to finalize cipher for AES encryption! (%d)", ret);
uiDrawString(strbuf, 0, breaks * font_height, 255, 0, 0);
return 0;
}
return 1;
}
/* Decrypt with context. */
int aes_decrypt(aes_ctx_t *ctx, void *dst, const void *src, size_t l)
{
int ret;
bool src_equals_dst = false;
if (src == dst)
{
src_equals_dst = true;
dst = malloc(l);
if (dst == NULL)
{
uiDrawString("Error: failed to allocate buffer for AES decryption!", 0, breaks * font_height, 255, 0, 0);
return 0;
}
}
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)
{
ret = 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));
ret = mbedtls_cipher_update(&ctx->cipher_dec, (const unsigned char *)src + offset, len, (unsigned char *)dst + offset, &out_len);
if (ret) break;
}
}
if (ret)
{
snprintf(strbuf, sizeof(strbuf) / sizeof(strbuf[0]), "Error: AES decryption failed! (%d)", ret);
uiDrawString(strbuf, 0, breaks * font_height, 255, 0, 0);
return 0;
}
/* Flush all data */
size_t strbuf_size = sizeof(strbuf);
ret = mbedtls_cipher_finish(&ctx->cipher_dec, (unsigned char *)strbuf, &strbuf_size); // Looks ugly, but using NULL,NULL with mbedtls on Switch is a no-no
if (ret)
{
snprintf(strbuf, sizeof(strbuf) / sizeof(strbuf[0]), "Error: failed to finalize cipher for AES decryption! (%d)", ret);
uiDrawString(strbuf, 0, breaks * font_height, 255, 0, 0);
return 0;
}
if (src_equals_dst)
{
memcpy((void*)src, dst, l);
free(dst);
}
return 1;
}
static void get_tweak(unsigned char *tweak, size_t sector)
{
/* Nintendo LE custom tweak... */
for (int i = 0xF; i >= 0; i--)
{
tweak[i] = (unsigned char)(sector & 0xFF);
sector >>= 8;
}
}
/* Encrypt with context for XTS. */
int aes_xts_encrypt(aes_ctx_t *ctx, void *dst, const void *src, size_t l, size_t sector, size_t sector_size)
{
int ret = 0;
unsigned char tweak[0x10];
if ((l % sector_size) != 0)
{
uiDrawString("Error: length must be a multiple of sector size in AES-XTS.", 0, breaks * font_height, 255, 0, 0);
return 0;
}
for (size_t i = 0; i < l; i += sector_size)
{
/* Workaround for Nintendo's custom sector...manually generate the tweak. */
get_tweak(tweak, sector++);
ret = aes_setiv(ctx, tweak, 16);
if (!ret) break;
ret = aes_encrypt(ctx, (char *)dst + i, (const char *)src + i, sector_size);
if (!ret) break;
}
return ret;
}
/* Decrypt with context for XTS. */
int aes_xts_decrypt(aes_ctx_t *ctx, void *dst, const void *src, size_t l, size_t sector, size_t sector_size)
{
int ret = 0;
unsigned char tweak[0x10];
if ((l % sector_size) != 0)
{
uiDrawString("Error: length must be a multiple of sector size in AES-XTS.", 0, breaks * font_height, 255, 0, 0);
return 0;
}
for (size_t i = 0; i < l; i += sector_size)
{
/* Workaround for Nintendo's custom sector...manually generate the tweak. */
get_tweak(tweak, sector++);
ret = aes_setiv(ctx, tweak, 16);
if (!ret) break;
ret = aes_decrypt(ctx, (char *)dst + i, (const char *)src + i, sector_size);
if (!ret) break;
}
return ret;
}