bitbottle 0.10.0

use clap::{App, ArgGroup};
use std::fs::File;
use std::io::{BufRead, BufReader, Read, Write};
use std::process::exit;
use bitbottle::*;
use bitbottle::cli::{pad_truncate, ProgressLine, to_binary_si};

const ABOUT: &str =
"Run a bitbottle encryption test (rabbet): Encrypt or decrypt arbitrary data
to/from an encrypted bottle.";

const ARGS: &str = "
-e --encrypt                'encrypt data into a bitbottle'
-d --decrypt                'decrypt an encrypted bitbottle'
-i --info                   'display info about an encrypted bitbottle'

-r --pub [FILE]...          '(encrypt) add ssh-ed25519 public key file to the list of possible recipients'
-s --secret [FILE]          '(decrypt) use an ssh-ed2519 private key file'
-p --password [PASSWORD]    'use a password to encrypt/decrypt the bitbottle'
--xchacha                   'use xchacha20-poly1305 to encrypt, instead of aes-128-gcm'

-z --snappy                 'compress with snappy (fast) before encrypting'
-y --lzma2                  'compress with lzma2 (intense) before encrypting'

-v --verbose                'describe what's happening as the bitbottle is read/written'
-o [FILE]                   'write output to file instead of stdout'
[FILE]                      'read input from file instead of stdin'
";

static mut VERBOSE: bool = false;

macro_rules! verbose {
    ($($arg:tt)*) => ({
        if unsafe { VERBOSE } {
            eprintln!($($arg)*);
        }
    })
}


fn copy_to<W: Write, R: Read>(mut writer: W, mut reader: R) -> std::io::Result<()> {
    let mut buffer = vec![0u8; 0x1_0000];
    loop {
        let n = reader.read(&mut buffer)?;
        if n == 0 { break; }
        writer.write_all(&buffer[..n])?;
    }
    Ok(())
}

fn drain<R: Read>(mut reader: R) -> std::io::Result<usize> {
    let mut byte_count = 0;
    let mut buffer = vec![0u8; 0x1_0000];
    loop {
        let n = reader.read(&mut buffer)?;
        if n == 0 { return Ok(byte_count); }
        byte_count += n;
    }
}

fn encrypt(
    mut reader: Box<dyn Read>,
    writer: Box<dyn Write>,
    password: Option<&str>,
    pk_filenames: Option<Vec<String>>,
    use_xchacha: bool,
    compression: Option<CompressionAlgorithm>,
) -> BottleResult<()> {
    // load any public keys
    let pks: BottleResult<Vec<Ed25519PublicKey>> = pk_filenames.iter().flatten().map(|filename| {
        Ed25519PublicKey::from_ssh_file(filename).map(|pk| {
            verbose!("Encrypting for {} ({})", pad_truncate(pk.name(), 16), hex::encode(pk.as_bytes()));
            pk
        })
    }).collect();
    let pks = pks.unwrap_or_else(|e| {
        eprintln!("ERROR: Unable to read public key file: {}", e);
        exit(1);
    });

    let mut options = EncryptedBottleWriterOptions::new();
    if let Some(password) = password {
        options.key = EncryptionKey::Password(password.to_string());
    }
    if use_xchacha {
        options.algorithm = EncryptionAlgorithm::XCHACHA20_POLY1305;
    }

    let progress = ProgressLine::new().bobble();
    progress.borrow_mut().show_bar = false;
    let writer = CountingWriter::new(writer, |byte_count, done| {
        let mut progress = progress.borrow_mut();
        progress.update(0f64, format!("Encrypting: {:>5}", to_binary_si(byte_count as f64)));
        if done { progress.force_update(); }
        progress.display();
    });

    let mut bottle_writer = if pks.is_empty() {
        EncryptedBottleWriter::new(writer, options)?
    } else {
        EncryptedBottleWriter::for_recipients(writer, options, pks.as_slice())?
    };

    let bottle_writer = if let Some(compression) = compression {
        let mut snappy_writer = CompressedBottleWriter::new(bottle_writer, compression)?;
        copy_to(&mut snappy_writer, &mut reader)?;
        snappy_writer.close()?
    } else {
        copy_to(&mut bottle_writer, &mut reader)?;
        bottle_writer
    };

    let writer = bottle_writer.close()?;
    let mut progress = progress.borrow_mut();
    progress.update(0f64, format!("Encrypting: {:>5} -- done.", to_binary_si(writer.count as f64)));
    progress.force_update().display();

    eprintln!();
    Ok(())
}

fn decrypt(
    reader: Box<dyn Read>,
    writer: Box<dyn Write>,
    password: Option<&str>,
    sk_filename: Option<&str>,
) -> BottleResult<()> {
    let bottle_reader = BottleReader::new(reader)?;

    // load the secret key?
    let sk = sk_filename.map(|filename| {
        Ed25519SecretKey::from_ssh_file(filename, None).map(|sk| {
            verbose!("Decrypting with key: {}", sk.name());
            sk
        })
    }).transpose().unwrap_or_else(|e| {
        eprintln!("ERROR: Unable to read secret key file: {}", e);
        exit(1);
    });

    let progress = ProgressLine::new().bobble();
    progress.borrow_mut().show_bar = false;
    let mut writer = CountingWriter::new(writer, |byte_count, done| {
        let mut progress = progress.borrow_mut();
        progress.update(0f64, format!("Decrypting: {:>5}", to_binary_si(byte_count as f64)));
        if done { progress.force_update(); }
        progress.display();
    });

    let encryption_key = password.map(|p| EncryptionKey::Password(p.to_string())).unwrap_or(EncryptionKey::Generate);
    let reader = match sk {
        Some(sk) => EncryptedBottleReader::build(bottle_reader, encryption_key, &[ sk ]),
        None => EncryptedBottleReader::new(bottle_reader, encryption_key),
    }?;
    let info = &reader.info;
    let key_count = info.public_encrypted_keys.len();
    verbose!("Bitbottle encrypted with {:?}{}{}",
        info.algorithm,
        if info.argon.is_some() { ", password" } else { "" },
        if key_count > 0 {
            format!(", {} public key{}", key_count, if key_count > 1 { "s" } else { "" })
        } else {
            String::new()
        },
    );

    // is it a nested bottle? a *compressed* nested bottle? dig in!
    let mut reader = BufReader::new(reader);
    let data = reader.fill_buf()?;
    let reader = if is_bottle(data) && data[5] == BottleType::Compressed as u8 {
        let mut compressed_reader = CompressedBottleReader::new(BottleReader::new(reader)?)?;
        verbose!("Nested bitbottle compressed with {:?}", compressed_reader.algorithm);
        copy_to(&mut writer, &mut compressed_reader)?;
        compressed_reader.close()?
    } else {
        copy_to(&mut writer, &mut reader)?;
        reader
    };
    reader.into_inner().close()?;

    let mut progress = progress.borrow_mut();
    progress.update(0f64, format!("Decrypting: {:>5} -- done.", to_binary_si(writer.count as f64)));
    progress.force_update().display();

    eprintln!();
    Ok(())
}

fn dump_info(info: &EncryptedBottleInfo) {
    eprintln!("Encryption:             {:?}", info.algorithm);
    if let Some(argon) = info.argon {
        eprintln!("Password encryption:    ARGON2ID (time={}, mem={}, par={})",
            argon.time_cost, to_binary_si((1 << argon.memory_cost_bits) as f64), argon.parallelism);
    }
    if !info.public_encrypted_keys.is_empty() {
        eprintln!("Public key algorithm:   {:?}", info.public_key_algorithm);
        for encrypted_key in &info.public_encrypted_keys {
            eprintln!("Encrypted for:          {} ({})",
                pad_truncate(encrypted_key.public_key.name(), 16), hex::encode(encrypted_key.public_key.as_bytes()));
        }
    }
    eprintln!("Block size:             {}", to_binary_si((1 << info.block_size_bits) as f64));
}

fn info(
    reader: Box<dyn Read>,
    password: Option<&str>,
    sk_filename: Option<&str>,
) -> BottleResult<()> {
    // load the secret key?
    let sk = sk_filename.map(|filename| {
        Ed25519SecretKey::from_ssh_file(filename, None).map(|sk| {
            verbose!("Decrypting with key: {}", sk.name());
            sk
        })
    }).transpose().unwrap_or_else(|e| {
        eprintln!("ERROR: Unable to read secret key file: {}", e);
        exit(1);
    });

    let encryption_key = password.map(|p| EncryptionKey::Password(p.to_string())).unwrap_or(EncryptionKey::Generate);

    let mut bottle_reader = BottleReader::new(reader)?;
    let info = EncryptedBottleReader::unpack_info(&mut bottle_reader)?;
    if bottle_reader.bottle_cap.bottle_type != BottleType::Encrypted {
        eprintln!("Not encrypted: {:?}", bottle_reader.bottle_cap.bottle_type);
        return Ok(());
    }
    dump_info(&info);
    let byte_count: usize;

    // if they gave us a password or secret key, try to decrypt it:
    if password.is_some() || sk.is_some() {
        let reader = match sk {
            Some(sk) => EncryptedBottleReader::build_with_info(bottle_reader, info, encryption_key, &[ sk ]),
            None => EncryptedBottleReader::new_with_info(bottle_reader, info, encryption_key),
        }?;

        // is it a nested bottle? a *compressed* nested bottle? dig in!
        let mut reader = BufReader::new(CountingReader::new(reader, |_, _| ()));
        let data = reader.fill_buf().unwrap_or_else(|_| {
            eprintln!("ERROR: Incorrect password or key.");
            exit(1);
        });

        let reader = if is_bottle(data) && data[5] == BottleType::Compressed as u8 {
            let mut compressed_reader = CompressedBottleReader::new(BottleReader::new(reader)?)?;
            eprintln!("Compressed:             {:?}", compressed_reader.algorithm);
            byte_count = drain(&mut compressed_reader)?;
            compressed_reader.close()?
        } else {
            byte_count = drain(&mut reader)?;
            reader
        };
        let counting_reader = reader.into_inner();
        eprintln!("Compressed size:        {}", to_binary_si(counting_reader.count as f64));
        counting_reader.into_inner().close()?;
    } else {
        // skip the entire encrypted stream -- we can't read it.
        bottle_reader.next_stream()?;
        byte_count = drain(&mut bottle_reader.data_stream()?)?;
        bottle_reader.close_stream()?;
        bottle_reader.next_stream()?;
        bottle_reader.close()?;
    }

    eprintln!("Total size:             {}", to_binary_si(byte_count as f64));

    Ok(())
}


pub fn main() {
    let args = App::new("rabbet")
        .version("1.0")
        .author("Robey Pointer <robey@lag.net>")
        .about(ABOUT)
        .args_from_usage(ARGS)
        .group(
            ArgGroup::with_name("task")
                .args(&[ "encrypt", "decrypt", "info" ])
                .multiple(false)
                .required(true)
        )
        .get_matches();

    if args.is_present("verbose") {
        unsafe { VERBOSE = true; }
    }

    // make sure there's some way to get this key
    let password = args.value_of("password");
    let pk_filenames = args.values_of_lossy("pub");
    let sk_filename = args.value_of("secret");
    if args.is_present("encrypt") && password.is_none() && pk_filenames.is_none() {
        eprintln!("ERROR: If you encrypt with no password and no public keys, nobody will ever be able to decrypt the bitbottle.");
        exit(1);
    }
    if args.is_present("decrypt") && password.is_none() && sk_filename.is_none() {
        eprintln!("ERROR: You need a password or a secret key to decrypt.");
        exit(1);
    }

    let reader: Box<dyn Read> = match args.value_of("FILE") {
        Some(filename) => Box::new(File::open(filename).unwrap_or_else(|e| {
            eprintln!("ERROR: Unable to read file {}: {}", filename, e);
            exit(1);
        })),
        None => Box::new(std::io::stdin()),
    };
    let writer: Box<dyn Write> = match args.value_of("o") {
        Some(filename) => Box::new(File::create(filename).unwrap_or_else(|e| {
            eprintln!("ERROR: Unable to write to file {}: {}", filename, e);
            exit(1);
        })),
        None => Box::new(std::io::stdout()),
    };

    let rv = if args.is_present("encrypt") {
        let compression = if args.is_present("snappy") {
            Some(CompressionAlgorithm::SNAPPY)
        } else if args.is_present("lzma2") {
            Some(CompressionAlgorithm::LZMA2)
        } else {
            None
        };
        encrypt(reader, writer, password, pk_filenames, args.is_present("xchacha"), compression)
    } else if args.is_present("decrypt") {
        decrypt(reader, writer, password, sk_filename)
    } else if args.is_present("info") {
        info(reader, password, sk_filename)
    } else {
        // impossible.
        Err(BottleError::CipherError)
    };
    rv.unwrap_or_else(|e| {
        eprintln!("ERROR: {}", e);
    });
}