rae/src/args/parse/generate.rs

//! Parser for generate mode
use super::*;
use tokio::sync::mpsc::Receiver;
use std::collections::HashSet;

#[instrument]
pub(super) async fn handle(mut rx: Receiver<Atom>) -> Result<config::Operation, error::ConstructError>
{
    let mut had = smallmap::Map::new();

    #[derive(Default)]
    struct Output
    {
	kind: Option<GenerateKeyKind>,
	input: Option<String>,
	password: Password,
	input_password: Password,
	format: KeyFormat,
	output: Option<String>,
	output_public: Option<String>,
    }

    #[cold]
    #[inline(never)]
    fn _invalidate<D: std::fmt::Display>(msg: D) -> !
    {
	panic!("{}. This should've been handled by `parse()`", msg);
    }
    
    macro_rules! invalidate {
	
	($msg:literal) => {
	    _invalidate($msg);
	};
	($msg:literal $($tt:tt)*) => {
	    _invalidate(format!($msg $($tt)*));
	}
    }
    let mut stage = Output{
	kind: Some(match rx.recv().await.unwrap() {
	    Atom::Generate(GenerateAtom::Type(ty)) => ty,
	    _ => invalidate!("First atom is not type."),
	}),
	..Output::default()
    };

    while let Some(Atom::Generate(atom)) = rx.recv().await
    {
	if GenerateAtom::hash_into(&atom, &mut had) {
	    return Err(error::ConstructError::AlreadyExists(format!("TODO: {:?}", atom)));
	}
	match atom {
	    GenerateAtom::Type(_) => invalidate!("Type cannot be specified twice."),
	    GenerateAtom::Input(input) => stage.input = Some(input),
	    GenerateAtom::Password(passwd) => stage.password = passwd,
	    GenerateAtom::Format(format) => stage.format =format,
	    GenerateAtom::Output(output) => stage.output = Some(output),
	    GenerateAtom::PublicOutput(output) => stage.output_public = Some(output),
	    GenerateAtom::InputPassword(passwd) => stage.input_password = passwd,
	}
    }

    let input_pw = stage.input_password;
    let output = if let Some(GenerateKeyKind::Rsa) = stage.kind {
	//RSA
	config::op::GenerateKey::Rsa(config::op::Rsa {
	    input: stage.input.map(|inp| (inp, input_pw)),
	    output_priv: stage.output.ok_or_else(|| error::ConstructError::ExpectedNotPresent("output"))?,
	    output_pub: stage.output_public,
	    password: stage.password,
	    description: Default::default(), /* TODO */
	})
    } else {
	//AES
	config::op::GenerateKey::Aes(config::op::Aes {
	    input: stage.input.map(|inp| (inp, input_pw)),
	    output: stage.output.ok_or_else(|| error::ConstructError::ExpectedNotPresent("output"))?,
	    password: stage.password,
	    description: Default::default(), /* TODO */
	})
    };
    Ok(config::Operation::GenerateKey(output))
}

fn split<'a>(input: &'a str) -> (&'a str, Option<&'a str>)
{
    if let Some(index) = input.find('=') {
	(&input[..index], Some(&input[(index+1)..]))
    } else {
	(input, None)
    }
}


pub(super) async fn parse<I: IntoIterator<Item=String>>(args: I, atoms: &mut mpsc::Sender<Atom>) -> Result<(), error::Error>
{
    let mut args = args.into_iter();
    let mut reading = true;
    let mut sent_output=0usize;
    
    macro_rules! send {
	($atom:expr) => {
	    atoms.send(Atom::Generate($atom))
	}
    }
    
    let (mode, _) = {
	let arg = args.next();
	match arg.as_ref().map(|x| x.trim().to_uppercase()).as_deref() {
	    Some("RSA") => (GenerateKeyKind::Rsa, send!(GenerateAtom::Type(GenerateKeyKind::Rsa)).await?),
	    Some("AES") => (GenerateKeyKind::Aes, send!(GenerateAtom::Type(GenerateKeyKind::Aes)).await?),
	    _ => return Err(error::Error::ExpectedKind(arg)),
	}
    };
    

    while let Some(arg) = args.next()
    {
	macro_rules! take_one {
	    ($msg:literal $($tt:tt)*) => {
		match args.next() {
		    Some(v) => v,
		    None => return Err((arg, error::ParseErrorKind::ExpectedArg(format!($msg $($tt)*))).into()),
		}
	    }
	}

	let mut had = HashSet::new();
	if reading && arg.starts_with("--") {
	    let (key, value) = split(&arg);
	    if had.contains(&key[..]) {
		return Err((error::ParseErrorKind::NotUniqueLong(Some(key.to_owned())), arg).swap().into());
	    } else {
		had.insert(key.to_owned());
	    }
	    match &key.to_lowercase().trim()[2..] {
		"input" => send!(GenerateAtom::Input(take_one!("--input: Expected filename"))).await?,
		"password" if value.is_some() => {
		    let value = value.unwrap();
		    send!(GenerateAtom::Password(Password::Specific(value.to_owned()))).await?;
		},
		"password" => send!(GenerateAtom::Password(Password::Yes)).await?,
		"format" if mode.is_rsa() => 
		    send!(GenerateAtom::Format(match take_one!("--format: Expected output format").to_uppercase().trim() {
			"PEM" => KeyFormat::PEM,
			"BIN" => KeyFormat::Bin,
			"TEXT" => KeyFormat::Text,
			__ => return Err((arg, error::ParseErrorKind::UnknownKeyFormat("rsa")).into()),
		    })).await?,
		"format" =>
		    send!(GenerateAtom::Format(match take_one!("--format: Expected output format").to_uppercase().trim() {
			"BIN" => KeyFormat::Bin,
			"TEXT" => KeyFormat::Text,
			__ => return Err((arg, error::ParseErrorKind::UnknownKeyFormat("aes")).into()),
		    })).await?,
		_ => return Err((arg, error::ParseErrorKind::UnexpectedArg(None)).into()),
	    }
	} else {
	    
	    if had.contains(&arg[..]) {
		return Err((arg, error::ParseErrorKind::NotUniqueLong(None)).into());
	    } else {
		had.insert(arg.to_owned());
	    }
	    reading=false;
	    match sent_output
	    {
		0 => send!(GenerateAtom::Output(arg)).await?,
		1 if mode.is_rsa() => send!(GenerateAtom::PublicOutput(arg)).await?,
		_ => return Err((arg, error::ParseErrorKind::UnexpectedArg(None)).into()),
	    }
	    sent_output +=1;
	}
	
    }
    
    Ok(())
}