metalssh 0.0.1

use criterion::BatchSize;
use criterion::Criterion;
use criterion::Throughput;
use criterion::black_box;
use criterion::criterion_group;
use criterion::criterion_main;
use metalssh::crypto::cipher::Cipher;
use metalssh::crypto::cipher::aes128gcm::Aes128Gcm;
use metalssh::crypto::cipher::aes256gcm::Aes256Gcm;
use metalssh::crypto::cipher::chacha20poly1305::ChaCha20Poly1305;
use metalssh::crypto::cipher::none::None;
use metalssh::wire::Packet;

// Static key material for ChaCha20-Poly1305 (64 bytes)
const CHAPOLY_KEY: [u8; 64] = [
    0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
    0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f,
    0x20, 0x21, 0x22, 0x23, 0x24, 0x25, 0x26, 0x27, 0x28, 0x29, 0x2a, 0x2b, 0x2c, 0x2d, 0x2e, 0x2f,
    0x30, 0x31, 0x32, 0x33, 0x34, 0x35, 0x36, 0x37, 0x38, 0x39, 0x3a, 0x3b, 0x3c, 0x3d, 0x3e, 0x3f,
];

// Static key material for AES-128-GCM (16 bytes key + 4 bytes IV)
const AES128_KEY: [u8; 16] = [
    0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
];

const AES128_IV: [u8; 4] = [0x10, 0x11, 0x12, 0x13];

// Static key material for AES-256-GCM (32 bytes key + 4 bytes IV)
const AES256_KEY: [u8; 32] = [
    0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
    0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f,
];

const AES256_IV: [u8; 4] = [0x20, 0x21, 0x22, 0x23];

const TAG_LEN: u8 = 16; // AEAD tag length (same for ChaCha20-Poly1305 and AES-GCM)
const PACKET_SIZE: usize = 207; // Average packet size from none-exec dataset

/// Create a test packet with random-ish data
fn create_test_packet(mac_length: u8) -> Vec<u8> {
    // SSH packet structure: packet_length(4) + padding_length(1) + payload +
    // padding + mac For simplicity, we'll use a fixed structure that's valid
    let payload_size = PACKET_SIZE - 5 - 8 - mac_length as usize; // leave room for header + padding + mac
    let padding_size = 8u8;
    let packet_length = 1 + payload_size as u32 + padding_size as u32;

    let mut packet = Vec::with_capacity(PACKET_SIZE + mac_length as usize);

    // Packet length (4 bytes, big endian)
    packet.extend_from_slice(&packet_length.to_be_bytes());

    // Padding length (1 byte)
    packet.push(padding_size);

    // Payload (pseudo-random pattern)
    for i in 0..payload_size {
        packet.push((i as u8).wrapping_mul(123).wrapping_add(45));
    }

    // Padding
    for i in 0..padding_size {
        packet.push(i.wrapping_mul(67).wrapping_add(89));
    }

    // MAC space (if needed)
    packet.resize(packet.len() + mac_length as usize, 0);

    packet
}

fn bench_none_encrypt(c: &mut Criterion) {
    let cipher = None {};
    let mut group = c.benchmark_group("none");
    group.throughput(Throughput::Bytes(PACKET_SIZE as u64));

    group.bench_function("encrypt", |b| {
        b.iter_batched(
            || create_test_packet(0),
            |mut data| {
                let seq = black_box(0u32);
                let mut packet = Packet::new(black_box(&mut data), 0);
                cipher
                    .encrypt_packet(&mut packet, seq)
                    .expect("encryption should succeed");
                black_box(data)
            },
            BatchSize::SmallInput,
        );
    });

    group.finish();
}

fn bench_none_decrypt(c: &mut Criterion) {
    let cipher = None {};
    let mut group = c.benchmark_group("none");
    group.throughput(Throughput::Bytes(PACKET_SIZE as u64));

    group.bench_function("decrypt", |b| {
        b.iter_batched(
            || create_test_packet(0),
            |mut data| {
                let seq = black_box(0u32);
                let mut packet = Packet::new(black_box(&mut data), 0);
                cipher
                    .decrypt_packet(&mut packet, seq)
                    .expect("decryption should succeed");
                black_box(data)
            },
            BatchSize::SmallInput,
        );
    });

    group.finish();
}

fn bench_chapoly_encrypt(c: &mut Criterion) {
    let cipher = ChaCha20Poly1305::new(CHAPOLY_KEY);
    let mut group = c.benchmark_group("chacha20poly1305");
    group.throughput(Throughput::Bytes(PACKET_SIZE as u64));

    group.bench_function("encrypt", |b| {
        b.iter_batched(
            || create_test_packet(TAG_LEN),
            |mut data| {
                let seq = black_box(0u32);
                let mut packet = Packet::new(black_box(&mut data), TAG_LEN);
                cipher
                    .encrypt_packet(&mut packet, seq)
                    .expect("encryption should succeed");
                black_box(data)
            },
            BatchSize::SmallInput,
        );
    });

    group.finish();
}

fn bench_chapoly_decrypt(c: &mut Criterion) {
    let cipher = ChaCha20Poly1305::new(CHAPOLY_KEY);
    let mut group = c.benchmark_group("chacha20poly1305");
    group.throughput(Throughput::Bytes(PACKET_SIZE as u64));

    group.bench_function("decrypt", |b| {
        b.iter_batched(
            || {
                // Setup: Pre-encrypt the packet so we can decrypt it (not measured)
                let mut data = create_test_packet(TAG_LEN);
                {
                    let mut packet = Packet::new(&mut data, TAG_LEN);
                    cipher
                        .encrypt_packet(&mut packet, 0)
                        .expect("pre-encryption should succeed");
                }
                data
            },
            |mut data| {
                // Benchmark: Only measure the decrypt operation
                let seq = black_box(0u32);
                let mut packet = Packet::new(black_box(&mut data), TAG_LEN);
                cipher
                    .decrypt_packet(&mut packet, seq)
                    .expect("decryption should succeed");
                black_box(data)
            },
            BatchSize::SmallInput,
        );
    });

    group.finish();
}

fn bench_aes128gcm_encrypt(c: &mut Criterion) {
    let cipher = Aes128Gcm::new(AES128_KEY, AES128_IV).expect("cipher creation should succeed");
    let mut group = c.benchmark_group("aes128gcm");
    group.throughput(Throughput::Bytes(PACKET_SIZE as u64));

    group.bench_function("encrypt", |b| {
        b.iter_batched(
            || create_test_packet(TAG_LEN),
            |mut data| {
                let seq = black_box(0u32);
                let mut packet = Packet::new(black_box(&mut data), TAG_LEN);
                cipher
                    .encrypt_packet(&mut packet, seq)
                    .expect("encryption should succeed");
                black_box(data)
            },
            BatchSize::SmallInput,
        );
    });

    group.finish();
}

fn bench_aes128gcm_decrypt(c: &mut Criterion) {
    let cipher = Aes128Gcm::new(AES128_KEY, AES128_IV).expect("cipher creation should succeed");
    let mut group = c.benchmark_group("aes128gcm");
    group.throughput(Throughput::Bytes(PACKET_SIZE as u64));

    group.bench_function("decrypt", |b| {
        b.iter_batched(
            || {
                // Setup: Pre-encrypt the packet so we can decrypt it (not measured)
                let mut data = create_test_packet(TAG_LEN);
                {
                    let mut packet = Packet::new(&mut data, TAG_LEN);
                    cipher
                        .encrypt_packet(&mut packet, 0)
                        .expect("pre-encryption should succeed");
                }
                data
            },
            |mut data| {
                // Benchmark: Only measure the decrypt operation
                let seq = black_box(0u32);
                let mut packet = Packet::new(black_box(&mut data), TAG_LEN);
                cipher
                    .decrypt_packet(&mut packet, seq)
                    .expect("decryption should succeed");
                black_box(data)
            },
            BatchSize::SmallInput,
        );
    });

    group.finish();
}

fn bench_aes256gcm_encrypt(c: &mut Criterion) {
    let cipher = Aes256Gcm::new(AES256_KEY, AES256_IV).expect("cipher creation should succeed");
    let mut group = c.benchmark_group("aes256gcm");
    group.throughput(Throughput::Bytes(PACKET_SIZE as u64));

    group.bench_function("encrypt", |b| {
        b.iter_batched(
            || create_test_packet(TAG_LEN),
            |mut data| {
                let seq = black_box(0u32);
                let mut packet = Packet::new(black_box(&mut data), TAG_LEN);
                cipher
                    .encrypt_packet(&mut packet, seq)
                    .expect("encryption should succeed");
                black_box(data)
            },
            BatchSize::SmallInput,
        );
    });

    group.finish();
}

fn bench_aes256gcm_decrypt(c: &mut Criterion) {
    let cipher = Aes256Gcm::new(AES256_KEY, AES256_IV).expect("cipher creation should succeed");
    let mut group = c.benchmark_group("aes256gcm");
    group.throughput(Throughput::Bytes(PACKET_SIZE as u64));

    group.bench_function("decrypt", |b| {
        b.iter_batched(
            || {
                // Setup: Pre-encrypt the packet so we can decrypt it (not measured)
                let mut data = create_test_packet(TAG_LEN);
                {
                    let mut packet = Packet::new(&mut data, TAG_LEN);
                    cipher
                        .encrypt_packet(&mut packet, 0)
                        .expect("pre-encryption should succeed");
                }
                data
            },
            |mut data| {
                // Benchmark: Only measure the decrypt operation
                let seq = black_box(0u32);
                let mut packet = Packet::new(black_box(&mut data), TAG_LEN);
                cipher
                    .decrypt_packet(&mut packet, seq)
                    .expect("decryption should succeed");
                black_box(data)
            },
            BatchSize::SmallInput,
        );
    });

    group.finish();
}

criterion_group!(
    benches,
    bench_none_encrypt,
    bench_none_decrypt,
    bench_chapoly_encrypt,
    bench_chapoly_decrypt,
    bench_aes128gcm_encrypt,
    bench_aes128gcm_decrypt,
    bench_aes256gcm_encrypt,
    bench_aes256gcm_decrypt
);
criterion_main!(benches);