*/
static void aes_setup(AESContext * ctx, int blocklen,
	       unsigned char *key, int keylen)
{
    int i, j, Nk, rconst;

    assert(blocklen == 16 || blocklen == 24 || blocklen == 32);
    assert(keylen == 16 || keylen == 24 || keylen == 32);

    /*
     * Basic parameters. Words per block, words in key, rounds.
     */
    Nk = keylen / 4;
    ctx->Nb = blocklen / 4;
    ctx->Nr = 6 + (ctx->Nb > Nk ? ctx->Nb : Nk);

    /*
     * Assign core-function pointers.
     */
    if (ctx->Nb == 8)
	ctx->encrypt = aes_encrypt_nb_8, ctx->decrypt = aes_decrypt_nb_8;
    else if (ctx->Nb == 6)
	ctx->encrypt = aes_encrypt_nb_6, ctx->decrypt = aes_decrypt_nb_6;
    else if (ctx->Nb == 4)
	ctx->encrypt = aes_encrypt_nb_4, ctx->decrypt = aes_decrypt_nb_4;

    /*
     * Now do the key setup itself.
     */
    rconst = 1;
    for (i = 0; i < (ctx->Nr + 1) * ctx->Nb; i++) {
	if (i < Nk)
	    ctx->keysched[i] = GET_32BIT_MSB_FIRST(key + 4 * i);
	else {
	    word32 temp = ctx->keysched[i - 1];
	    if (i % Nk == 0) {
		int a, b, c, d;
		a = (temp >> 16) & 0xFF;
		b = (temp >> 8) & 0xFF;
		c = (temp >> 0) & 0xFF;
		d = (temp >> 24) & 0xFF;
		temp = Sbox[a] ^ rconst;
		temp = (temp << 8) | Sbox[b];
		temp = (temp << 8) | Sbox[c];
		temp = (temp << 8) | Sbox[d];
		rconst = mulby2(rconst);
	    } else if (i % Nk == 4 && Nk > 6) {
		int a, b, c, d;
		a = (temp >> 24) & 0xFF;
		b = (temp >> 16) & 0xFF;
		c = (temp >> 8) & 0xFF;
		d = (temp >> 0) & 0xFF;
		temp = Sbox[a];
		temp = (temp << 8) | Sbox[b];
		temp = (temp << 8) | Sbox[c];
		temp = (temp << 8) | Sbox[d];
	    }
	    ctx->keysched[i] = ctx->keysched[i - Nk] ^ temp;
	}
    }

    /*
     * Now prepare the modified keys for the inverse cipher.
     */
    for (i = 0; i <= ctx->Nr; i++) {
	for (j = 0; j < ctx->Nb; j++) {
	    word32 temp;
	    temp = ctx->keysched[(ctx->Nr - i) * ctx->Nb + j];
	    if (i != 0 && i != ctx->Nr) {
		/*
		 * Perform the InvMixColumn operation on i. The D
		 * tables give the result of InvMixColumn applied
		 * to Sboxinv on individual bytes, so we should
		 * compose Sbox with the D tables for this.
		 */
		int a, b, c, d;
		a = (temp >> 24) & 0xFF;
		b = (temp >> 16) & 0xFF;
		c = (temp >> 8) & 0xFF;
		d = (temp >> 0) & 0xFF;
		temp = D0[Sbox[a]];
		temp ^= D1[Sbox[b]];
		temp ^= D2[Sbox[c]];
		temp ^= D3[Sbox[d]];
	    }
	    ctx->invkeysched[i * ctx->Nb + j] = temp;
	}
    }
}

static void aes_encrypt(AESContext * ctx, word32 * block)
{
    ctx->encrypt(ctx, block);
}

static void aes_decrypt(AESContext * ctx, word32 * block)
{
    ctx->decrypt(ctx, block);
}

static void aes_encrypt_cbc(unsigned char *blk, int len, AESContext * ctx)
{
    word32 iv[4];
    int i;

    assert((len & 15) == 0);

    memcpy(iv, ctx->iv, sizeof(iv));

    while (len > 0) {
	for (i = 0; i < 4; i++)
	    iv[i] ^= GET_32BIT_MSB_FIRST(blk + 4 * i);
	aes_encrypt(ctx, iv);
	for (i = 0; i < 4; i++)
	    PUT_32BIT_MSB_FIRST(blk + 4 * i, iv[i]);
	blk += 16;
	len -= 16;
    }

    memcpy(ctx->iv, iv, sizeof(iv));
}

static void aes_decrypt_cbc(unsigned char *blk, int len, AESContext * ctx)
{
    word32 iv[4], x[4], ct[4];
    int i;

    assert((len & 15) == 0);

    memcpy(iv, ctx->iv, sizeof(iv));

    while (len > 0) {
	for (i = 0; i < 4; i++)
	    x[i] = ct[i] = GET_32BIT_MSB_FIRST(blk + 4 * i);
	aes_decrypt(ctx, x);
	for (i = 0; i < 4; i++) {
	    PUT_32BIT_MSB_FIRST(blk + 4 * i, iv[i] ^ x[i]);
	    iv[i] = ct[i];
	}
	blk += 16;
	len -= 16;
    }

    memcpy(ctx->iv, iv, sizeof(iv));
}

static void aes_sdctr(unsigned char *blk, int len, AESContext *ctx)
{
    word32 iv[4], b[4], tmp;
    int i;

    assert((len & 15) == 0);

    memcpy(iv, ctx->iv, sizeof(iv));

    while (len > 0) {
	memcpy(b, iv, sizeof(b));
	aes_encrypt(ctx, b);
	for (i = 0; i < 4; i++) {
	    tmp = GET_32BIT_MSB_FIRST(blk + 4 * i);
	    PUT_32BIT_MSB_FIRST(blk + 4 * i, tmp ^ b[i]);
	}
	for (i = 3; i >= 0; i--)
	    if ((iv[i] = (iv[i] + 1) & 0xffffffff) != 0)
		break;
	blk += 16;
	len -= 16;
    }

    memcpy(ctx->iv, iv, sizeof(iv));
}

static void *aes_make_context(void)
{
    return snew(AESContext);
}

static void aes_free_context(void *handle)
{
    sfree(handle);
}

static void aes128_key(void *handle, unsigned char *key)
{
    AESContext *ctx = (AESContext *)handle;
    aes_setup(ctx, 16, key, 16);
}

static void aes192_key(void *handle, unsigned char *key)
{
    AESContext *ctx = (AESContext *)handle;
    aes_setup(ctx, 16, key, 24);
}

static void aes256_key(void *handle, unsigned char *key)
{
    AESContext *ctx = (AESContext *)handle;
    aes_setup(ctx, 16, key, 32);
}

static void aes_iv(void *handle, unsigned char *iv)
{
    AESContext *ctx = (AESContext *)handle;
    int i;
    for (i = 0; i < 4; i++)
	ctx->iv[i] = GET_32BIT_MSB_FIRST(iv + 4 * i);
}

static void aes_ssh2_encrypt_blk(void *handle, unsigned char *blk, int len)
{
    AESContext *ctx = (AESContext *)handle;
    aes_encrypt_cbc(blk, len, ctx);
}

static void aes_ssh2_decrypt_blk(void *handle, unsigned char *blk, int len)
{
    AESContext *ctx = (AESContext *)handle;
    aes_decrypt_cbc(blk, len, ctx);
}

static void aes_ssh2_sdctr(void *handle, unsigned char *blk, int len)
{
    AESContext *ctx = (AESContext *)handle;
    aes_sdctr(blk, len, ctx);
}

void aes256_encrypt_pubkey(unsigned char *key, unsigned char *blk, int len)
{
    AESContext ctx;
    aes_setup(&ctx, 16, key, 32);
    memset(ctx.iv, 0, sizeof(ctx.iv));
    aes_encrypt_cbc(blk, len, &ctx);
    memset(&ctx, 0, sizeof(ctx));
}

void aes256_decrypt_pubkey(unsigned char *key, unsigned char *blk, int len)
{
    AESContext ctx;
    aes_setup(&ctx, 16, key, 32);
    memset(ctx.iv, 0, sizeof(ctx.iv));
    aes_decrypt_cbc(blk, len, &ctx);
    memset(&ctx, 0, sizeof(ctx));
}

static const struct ssh2_cipher ssh_aes128_ctr = {
    aes_make_context, aes_free_context, aes_iv, aes128_key,
    aes_ssh2_sdctr, aes_ssh2_sdctr,
    "aes128-ctr",
    16, 128, 0, "AES-128 SDCTR"
};

static const struct ssh2_cipher ssh_aes192_ctr = {
    aes_make_context, aes_free_context, aes_iv, aes192_key,
    aes_ssh2_sdctr, aes_ssh2_sdctr,
    "aes192-ctr",
    16, 192, 0, "AES-192 SDCTR"
};

static const struct ssh2_cipher ssh_aes256_ctr = {
    aes_make_context, aes_free_context, aes_iv, aes256_key,
    aes_ssh2_sdctr, aes_ssh2_sdctr,
    "aes256-ctr",
    16, 256, 0, "AES-256 SDCTR"
};

static const struct ssh2_cipher ssh_aes128 = {
    aes_make_context, aes_free_context, aes_iv, aes128_key,
    aes_ssh2_encrypt_blk, aes_ssh2_decrypt_blk,
    "aes128-cbc",
    16, 128, SSH_CIPHER_IS_CBC, "AES-128 CBC"
};

static const struct ssh2_cipher ssh_aes192 = {
    aes_make_context, aes_free_context, aes_iv, aes192_key,
    aes_ssh2_encrypt_blk, aes_ssh2_decrypt_blk,
    "aes192-cbc",
    16, 192, SSH_CIPHER_IS_CBC, "AES-192 CBC"
};

static const struct ssh2_cipher ssh_aes256 = {
    aes_make_context, aes_free_context, aes_iv, aes256_key,
    aes_ssh2_encrypt_blk, aes_ssh2_decrypt_blk,
    "aes256-cbc",
    16, 256, SSH_CIPHER_IS_CBC, "AES-256 CBC"
};

static const struct ssh2_cipher ssh_rijndael_lysator = {
    aes_make_context, aes_free_context, aes_iv, aes256_key,
    aes_ssh2_encrypt_blk, aes_ssh2_decrypt_blk,
    "rijndael-cbc@lysator.liu.se",
    16, 256, SSH_CIPHER_IS_CBC, "AES-256 CBC"
};

static const struct ssh2_cipher *const aes_list[] = {
    &ssh_aes256_ctr,
    &ssh_aes256,
    &ssh_rijndael_lysator,
    &ssh_aes192_ctr,
    &ssh_aes192,
    &ssh_aes128_ctr,
    &ssh_aes128,
};

const struct ssh2_ciphers ssh2_aes = {
    sizeof(aes_list) / sizeof(*aes_list),
    aes_list
};