use std::{
collections::HashMap,
env,
fs::File,
io::{self, BufRead, BufReader, Error, ErrorKind, Read},
path::Path,
str::FromStr,
};
use crate::{
MergeOption,
paths::permissions::{self, Permissions},
};
pub const TRUE_AS_STR: &str = "true";
pub const FALSE_AS_STR: &str = "false";
#[derive(Default, Debug, Clone)]
pub struct Args {
cmds: Option<Vec<String>>,
args: HashMap<String, Vec<String>>,
sub_args: Option<Vec<String>>,
use_stdin: bool,
}
impl AsRef<Args> for Args {
fn as_ref(&self) -> &Self {
self
}
}
impl Args {
pub fn cmd(&self) -> Option<&str> {
match self.cmds {
Some(ref cmds) => cmds.first().map(|s| s.as_str()),
None => None,
}
}
pub fn cmds(&self) -> Option<Vec<&str>> {
self.cmds.as_ref().map(|v| v.iter().map(|s| s.as_str()).collect())
}
pub fn args(&self) -> &HashMap<String, Vec<String>> {
&self.args
}
pub fn sub_args(&self) -> Option<Vec<&str>> {
self.sub_args.as_ref().map(|v| v.iter().map(|s| s.as_str()).collect())
}
pub fn is_empty(&self) -> bool {
self.cmds.is_none() && self.args.is_empty() && self.sub_args.is_none()
}
pub fn has_args(&self) -> bool {
self.args.is_empty() == false || self.sub_args.is_some()
}
pub fn use_stdin(&self) -> bool {
self.use_stdin
}
pub fn sub_cmd(self) -> (Option<String>, Self) {
let mut cmds = self.cmds;
(
match cmds.as_mut() {
Some(cmds) => match cmds.is_empty() {
true => None,
false => Some(cmds.remove(0)),
},
None => None,
},
Self {
cmds: match cmds {
Some(cmds) => match cmds.is_empty() {
true => None,
false => Some(cmds),
},
None => None,
},
..self
}
)
}
pub fn get<T, S>(&self, keys: &[S]) -> io::Result<Option<T>> where T: FromStr, S: AsRef<str> {
let mut result = None;
for key in keys {
let key = key.as_ref();
if let Some(values) = self.args.get(key) {
match result.is_none() {
true => match values.len() {
0 => return Err(Error::new(ErrorKind::InvalidData, __!("Internal parser error: no values found"))),
1 => result = Some(
T::from_str(&values[0])
.map_err(|_| Error::new(ErrorKind::Other, format!("Failed parsing value {:?} of {:?}", &values[0], &key)))?
),
_ => return Err(Error::new(ErrorKind::InvalidData, format!("Expected 1 value, got: {:?}", &values))),
},
false => return Err(Error::new(ErrorKind::InvalidData, format!("Duplicate value for {:?}", key))),
};
}
}
Ok(result)
}
pub fn take<T, S>(&mut self, keys: &[S]) -> io::Result<Option<T>> where T: FromStr, S: AsRef<str> {
let result = self.get(keys);
if result.as_ref().map(|r| r.is_some()).unwrap_or(false) {
for key in keys {
self.args.remove(key.as_ref());
}
}
result
}
pub fn get_vec<T, S>(&self, keys: &[S]) -> io::Result<Option<Vec<T>>> where T: FromStr, S: AsRef<str> {
let mut result: Option<Vec<_>> = None;
for key in keys {
let key = key.as_ref();
if let Some(values) = self.args.get(key) {
for v in values.iter() {
let v = T::from_str(v).map_err(|_| Error::new(ErrorKind::Other, format!("Failed parsing value {:?} of {:?}", v, key)))?;
match result {
Some(ref mut result) => result.push(v),
None => result = Some(vec![v]),
};
}
}
}
Ok(result)
}
pub fn take_vec<T, S>(&mut self, keys: &[S]) -> io::Result<Option<Vec<T>>> where T: FromStr, S: AsRef<str> {
let result = self.get_vec(keys);
if result.as_ref().map(|r| r.is_some()).unwrap_or(false) {
for key in keys {
self.args.remove(key.as_ref());
}
}
result
}
pub fn take_cmds(&mut self) -> Option<Vec<String>> {
self.cmds.take()
}
pub fn take_sub_args(&mut self) -> Option<Vec<String>> {
self.sub_args.take()
}
pub fn merge_args(&mut self, other: &mut Self, merge_option: MergeOption) -> usize {
let mut count = 0;
let mut not_merged: Option<HashMap<_, _>> = None;
for (key, other_value) in other.args.drain() {
match self.args.get_mut(&key) {
Some(value) => match merge_option {
MergeOption::TakeAll => {
*value = other_value;
count += 1;
},
MergeOption::IgnoreExisting => match not_merged.as_mut() {
Some(not_merged) => drop(not_merged.insert(key, other_value)),
None => not_merged = Some(vec![(key, other_value)].into_iter().collect()),
},
},
None => {
self.args.insert(key, other_value);
count += 1;
},
};
}
if let Some(not_merged) = not_merged {
other.args = not_merged;
}
count
}
}
#[derive(Debug, Eq, PartialEq)]
enum ArgKind {
Command,
ShortOption,
ShortOptionWithValue,
SomeBoolShortOptions,
LongOption,
LongOptionWithValue,
SubArgsSeparator,
StdinIndicator,
}
pub fn parse_strings<S, I>(args: I) -> io::Result<Args> where S: AsRef<str>, I: Iterator<Item=S> {
let mut result = Args {
cmds: None,
args: HashMap::new(),
sub_args: None,
use_stdin: false,
};
fn add_option(arg: &str, value: String, mut options: Option<&mut Vec<String>>, allow_duplicates: bool) -> io::Result<Option<Vec<String>>> {
if let Some(options) = options.as_mut() {
if allow_duplicates || options.contains(&value) == false {
options.push(value);
return Ok(None);
} else {
return Err(Error::new(ErrorKind::InvalidData, format!("Duplicate value for {:?}", &arg)));
}
}
Ok(Some(vec![value]))
}
let mut args = args.peekable();
while let Some(arg) = args.next() {
let arg = arg.as_ref();
match ArgKind::parse(arg)? {
ArgKind::Command => if let Some(cmds) = add_option(arg, arg.to_string(), result.cmds.as_mut(), true)? {
result.cmds = Some(cmds);
},
ArgKind::SomeBoolShortOptions => {
let parts = arg.splitn(2, '=').collect::<Vec<&str>>();
let value = match parts.len() {
1 => TRUE_AS_STR.to_string(),
_ => match parts[1].trim().to_lowercase().as_str() {
self::TRUE_AS_STR => self::TRUE_AS_STR.to_string(),
self::FALSE_AS_STR => self::FALSE_AS_STR.to_string(),
_ => return Err(Error::new(ErrorKind::InvalidData, format!("Invalid some boolean short options: {:?}", &arg))),
},
};
for chr in parts[0].chars() {
let arg = format!("-{}", &chr);
if let Some(options) = add_option(&arg, value.clone(), result.args.get_mut(&arg), false)? {
result.args.insert(arg, options);
}
}
},
ArgKind::ShortOption | ArgKind::LongOption => {
let value = match args.peek().map(|s| s.as_ref().starts_with('-')) {
Some(true) | None => TRUE_AS_STR.to_string(),
Some(false) => args.next().map(|s| s.as_ref().to_string()).ok_or_else(||
Error::new(ErrorKind::InvalidData, format!("Missing value for {:?}", &arg))
)?,
};
if let Some(options) = add_option(arg, value, result.args.get_mut(arg), false)? {
result.args.insert(arg.to_string(), options);
}
},
ArgKind::ShortOptionWithValue | ArgKind::LongOptionWithValue => {
let parts = arg.splitn(2, '=').collect::<Vec<&str>>();
match parts.len() {
2 => {
let (arg, value) = (parts[0].to_string(), parts[1].to_string());
if let Some(options) = add_option(&arg, value, result.args.get_mut(&arg), false)? {
result.args.insert(arg, options);
}
},
_ => return Err(Error::new(ErrorKind::InvalidData, format!("Invalid option: {:?}", &arg))),
};
},
ArgKind::SubArgsSeparator => {
let sub_args = args.map(|some| some.as_ref().to_string()).collect::<Vec<String>>();
if sub_args.is_empty() == false {
result.sub_args = Some(sub_args);
}
break;
},
ArgKind::StdinIndicator => {
match result.use_stdin {
true => return Err(Error::new(ErrorKind::InvalidData, format!("Duplicate {:?}", &arg))),
false => result.use_stdin = true,
};
},
};
}
Ok(result)
}
pub fn parse() -> io::Result<Args> {
parse_strings(env::args().skip(1))
}
pub fn parse_file<P>(file: Option<P>, permissions: Option<Permissions>) -> io::Result<Args> where P: AsRef<Path> {
let current_exe = env::current_exe()?;
let default_file = match current_exe.parent() {
Some(dir) => {
let default_file = dir.join(crate::DIA_ARGS_FILE_NAME);
match default_file.exists() {
true => default_file.canonicalize()?,
false => default_file,
}
},
None => return Err(Error::new(ErrorKind::NotFound, "Could not find parent directory of the program")),
};
let file = match file.map(|f| f.as_ref().to_path_buf()) {
Some(file) => match file.canonicalize() {
Ok(file) => file,
Err(err) => return Err(Error::new(err.kind(), format!("Failed accessing {:?}: {:?}", file, err))),
},
None => default_file.clone(),
};
if file == default_file && default_file.exists() == false {
return Ok(Args::default());
}
if file.exists() == false {
return Err(Error::new(ErrorKind::NotFound, format!("File not found: {:?}", file)));
}
if file.is_file() == false {
return Err(Error::new(ErrorKind::InvalidInput, format!("Not a file: {:?}", file)));
}
if file.metadata()?.len() > crate::MAX_DIA_ARGS_FILE_SIZE {
return Err(Error::new(ErrorKind::InvalidInput, format!("File too large: {:?} (max allowed: {})", file, crate::MAX_DIA_ARGS_FILE_SIZE)));
}
match permissions {
Some(permissions) => permissions::matches(&file, permissions)?,
None => permissions::ensure_second_is_safe(¤t_exe, &file)?,
};
let mut args: Vec<String> = vec![];
for line in read_file_to_string(file, Some(crate::MAX_DIA_ARGS_FILE_SIZE))?.lines() {
let line = line.trim();
if line.is_empty() || line.starts_with('#') {
continue;
}
if line.starts_with('-') == false || false == ['=', ' ', '\t'].iter().any(|separator| match line.find(*separator) {
Some(idx) => {
args.push(format!("{}={}", line[..idx].trim(), line[idx + 1..].trim()));
true
},
None => false,
}) {
args.push(line.to_owned());
}
}
parse_strings(args.into_iter())
}
fn read_file_to_string<P>(file: P, limit: Option<u64>) -> io::Result<String> where P: AsRef<Path> {
const BUF_SIZE: usize = 8 * 1024;
let file = file.as_ref();
let limit = file.metadata()?.len().min(limit.unwrap_or(u64::max_value()));
let mut reader = BufReader::new(File::open(file)?).take(limit);
let mut buf = [0; BUF_SIZE];
let mut data = Vec::with_capacity(limit as usize);
loop {
match reader.read(&mut buf)? {
0 => return String::from_utf8(data).map_err(|e| Error::new(ErrorKind::InvalidData, e)),
read => data.extend(&buf[..read]),
};
}
}
impl ArgKind {
fn parse_long_options<S>(s: S) -> io::Result<ArgKind> where S: AsRef<str> {
let s = s.as_ref();
for (index, chr) in s.chars().skip(2).enumerate() {
match chr {
'=' => match index {
0 => return Err(Error::new(ErrorKind::InvalidData, format!("Invalid long option: {:?}", &s))),
_ => return Ok(ArgKind::LongOptionWithValue),
},
'a'...'z' | 'A'...'Z' | '0'...'9' | '-' => continue,
_ => return Err(Error::new(ErrorKind::InvalidData, format!("Invalid long option: {:?}", &s))),
};
}
Ok(ArgKind::LongOption)
}
fn parse_short_options<S>(s: S) -> io::Result<ArgKind> where S: AsRef<str> {
let s = s.as_ref();
let mut must_be_some_bool_options = false;
let mut chars = s.char_indices().skip(1);
while let Some((index, chr)) = chars.next() {
match chr {
'=' => match index {
0 => return Err(Error::new(ErrorKind::Other, __!("Invalid index: {} -> {}", &index, &chr))),
1 => return Err(Error::new(ErrorKind::InvalidData, format!("Invalid short option: {:?}", &s))),
2 => return Ok(ArgKind::ShortOptionWithValue),
_ => match chars.map(|(_, chr)| chr).collect::<String>().trim().to_lowercase().as_str() {
self::TRUE_AS_STR | self::FALSE_AS_STR => return Ok(ArgKind::SomeBoolShortOptions),
_ => return Err(Error::new(ErrorKind::InvalidData, format!("Invalid some boolean short options: {:?}", &s))),
},
},
'a'...'z' | 'A'...'Z' | '0'...'9' => if index > 1 {
must_be_some_bool_options = true;
},
_ => return Err(Error::new(ErrorKind::InvalidData, format!("Invalid short option: {:?}", &s))),
};
}
match must_be_some_bool_options {
true => Ok(ArgKind::SomeBoolShortOptions),
false => Ok(ArgKind::ShortOption),
}
}
fn parse<S>(s: S) -> io::Result<ArgKind> where S: AsRef<str> {
let s = s.as_ref();
if s.starts_with("--") {
return match s.len() {
2 => Ok(ArgKind::SubArgsSeparator),
_ => Self::parse_long_options(s),
};
}
if s.starts_with('-') {
return match s.len() {
1 => Ok(ArgKind::StdinIndicator),
_ => Self::parse_short_options(s),
};
}
Ok(ArgKind::Command)
}
}
pub fn read_line<T>() -> io::Result<T> where T: FromStr {
let stdin = io::stdin();
let mut stdin = stdin.lock();
let mut buf = String::with_capacity(64);
stdin.read_line(&mut buf)?;
let buf = buf.trim();
T::from_str(&buf).map_err(|_| Error::new(ErrorKind::InvalidData, format!("Invalid {:?}", buf)))
}