//
// Syd: rock-solid application kernel
// src/syd-aes.rs: AES-GCM Encryption and Decryption Utility
//
// Copyright (c) 2024, 2025 Ali Polatel <alip@chesswob.org>
//
// SPDX-License-Identifier: GPL-3.0

use std::{
    os::{fd::AsRawFd, unix::ffi::OsStrExt},
    process::ExitCode,
    time::{Duration, Instant},
};

use hex::FromHex;
use nix::{
    fcntl::{splice, OFlag, SpliceFFlags},
    sys::socket::{sendmsg, AlgAddr, ControlMessage, MsgFlags},
    unistd::{isatty, pipe2},
};
use secure_string::SecureBytes;
use syd::{
    config::PIPE_BUF,
    err::SydResult,
    hash::{
        aes_ctr_dec, aes_ctr_enc, aes_ctr_flush, aes_ctr_init, aes_ctr_setup, Key, BLOCK_SIZE, IV,
        IV_SIZE, KEY_SIZE,
    },
};

fn process_data(encrypting: bool, key: Key, iv: IV, verbose: bool) -> SydResult<()> {
    let fd = aes_ctr_setup(&key)?;
    let fd = aes_ctr_init(&fd, false)?;

    if encrypting {
        aes_ctr_enc(&fd, &[], Some(&iv), true)?;
    } else {
        aes_ctr_dec(&fd, &[], Some(&iv), true)?;
    }

    let (pipe_rd, pipe_wr) = pipe2(OFlag::O_CLOEXEC)?;

    let mut nread = 0;
    let stime = Instant::now();
    let mut ltime = stime;
    let mut nbytes = 0;
    let mut nwrite = 0;
    let verbose = if verbose {
        isatty(nix::libc::STDERR_FILENO).unwrap_or(false)
    } else {
        false
    };

    // SAFETY: This buffer holds plaintext,
    // we zero it on free and ensure it never swaps out.
    let (mut bufz, mut bufu) = if !encrypting {
        (Some(SecureBytes::from(vec![0u8; PIPE_BUF])), None)
    } else {
        (None, Some(vec![0u8; PIPE_BUF]))
    };
    let buf = if let Some(ref mut bufz) = bufz {
        bufz.unsecure_mut()
    } else if let Some(ref mut bufu) = bufu {
        bufu
    } else {
        unreachable!()
    };

    loop {
        // Use splice to move data from standard input to pipe.
        match splice(
            std::io::stdin(),
            None,
            &pipe_wr,
            None,
            PIPE_BUF,
            SpliceFFlags::empty(),
        )? {
            0 => break, // EOF
            n => {
                match splice(&pipe_rd, None, &fd, None, n, SpliceFFlags::SPLICE_F_MORE)? {
                    0 => break, // EOF
                    n => nread += n,
                };

                while nread >= BLOCK_SIZE {
                    let nblock = (nread / BLOCK_SIZE) * BLOCK_SIZE;
                    let n = aes_ctr_flush(&fd, &std::io::stdout(), buf, nblock)?;
                    nread -= n;
                    nbytes += n;
                    nwrite += 1;
                }

                if verbose {
                    let now = Instant::now();
                    if now.duration_since(ltime) >= Duration::from_millis(500) {
                        let elapsed = stime.elapsed();
                        let speed = nbytes as f64 / elapsed.as_secs_f64();
                        let output = format!(
                            "{} bytes ({:.2} GB, {:.2} GiB) processed, {:.2?} s, {:.2} MB/s",
                            nbytes,
                            nbytes as f64 / 1_000_000_000.0,
                            nbytes as f64 / (1 << 30) as f64,
                            elapsed,
                            speed / (1 << 20) as f64
                        );
                        eprint!("\r\x1B[K{output}");
                        ltime = now;
                    }
                }
            }
        }
    }

    // Signal the AF_ALG socket the end of the process (No MSG_MORE in flags).
    let cmsgs = [ControlMessage::AlgSetOp(if encrypting {
        &nix::libc::ALG_OP_ENCRYPT
    } else {
        &nix::libc::ALG_OP_DECRYPT
    })];
    sendmsg::<AlgAddr>(fd.as_raw_fd(), &[], &cmsgs, MsgFlags::empty(), None)?;

    if nread > 0 {
        // {En,De}crypt the final batch.
        // SAFETY: Zero-out memory if decrypting.
        aes_ctr_flush(&fd, &std::io::stdout(), buf, nread)?;
        if verbose {
            nbytes += nread;
            nwrite += 1;
        }
    }

    if verbose {
        let elapsed = stime.elapsed();
        eprintln!(
            "\n{} records of each {} bytes processed.\n{} bytes ({:.2} GB, {:.2} GiB) processed, {:.5?} s, {:.2} MB/s",
            nwrite,
            PIPE_BUF,
            nbytes,
            nbytes as f64 / 1_000_000_000.0,
            nbytes as f64 / (1 << 30) as f64,
            elapsed,
            nbytes as f64 / elapsed.as_secs_f64() / (1 << 20) as f64
        );
    }

    Ok(())
}

fn main() -> SydResult<ExitCode> {
    use lexopt::prelude::*;

    syd::set_sigpipe_dfl()?;

    // Parse CLI options.
    let mut opt_encrypt = None;
    let mut opt_key_hex = None;
    let mut opt_iv_hex = None;
    let mut opt_tag = None;
    let mut opt_verbose = false;

    let mut parser = lexopt::Parser::from_env();
    while let Some(arg) = parser.next()? {
        match arg {
            Short('h') => {
                help();
                return Ok(ExitCode::SUCCESS);
            }
            Short('v') => opt_verbose = true,
            Short('e') => opt_encrypt = Some(true),
            Short('d') => opt_encrypt = Some(false),
            Short('k') => opt_key_hex = Some(parser.value()?.parse::<String>()?),
            Short('i') => opt_iv_hex = Some(parser.value()?.parse::<String>()?),
            Short('t') => opt_tag = Some(parser.value()?),
            _ => return Err(arg.unexpected().into()),
        }
    }

    if opt_encrypt.is_none() || opt_key_hex.is_none() || opt_iv_hex.is_none() {
        eprintln!("syd-aes: Error: -e or -d and -k, -i options are required.");
        help();
        return Ok(ExitCode::FAILURE);
    }

    #[allow(clippy::disallowed_methods)]
    let mut key = match <[u8; KEY_SIZE]>::from_hex(&opt_key_hex.unwrap()) {
        Ok(key) => Key::new(key),
        Err(err) => {
            eprintln!(
                "syd-aes: Error: Key must be valid hex, and 512 bits (64 bytes) in length: {err}"
            );
            return Ok(ExitCode::FAILURE);
        }
    };

    if let Some(info) = opt_tag {
        key = key.derive(None, info.as_bytes());
    }

    #[allow(clippy::disallowed_methods)]
    let iv = match <[u8; IV_SIZE]>::from_hex(&opt_iv_hex.unwrap()) {
        Ok(iv) => IV::new(iv),
        Err(err) => {
            eprintln!(
                "syd-aes: Error: IV must be valid hex, and 128 bits (16 bytes) in length: {err}"
            );
            return Ok(ExitCode::FAILURE);
        }
    };

    #[allow(clippy::disallowed_methods)]
    process_data(opt_encrypt.unwrap(), key, iv, opt_verbose).map(|_| ExitCode::SUCCESS)
}

fn help() {
    println!("Usage: syd-aes [-h] -e|-d -k <key-hex> -i <iv-hex> -t <tag-str>");
    println!("AES-CTR Encryption and Decryption Utility");
    println!("Reads from standard input and writes to standard output.");
    println!("  -h        Print this help message and exit.");
    println!("  -v        Enable verbose mode.");
    println!("  -e        Encrypt the input data.");
    println!("  -d        Decrypt the input data.");
    println!("  -k <key>  Hex-encoded key (256 bits)");
    println!("  -i <iv>   Hex-encoded IV (128 bits)");
    println!("  -t <tag>  Information tag for HKDF (optional)");
}
