marg/
lib.rs

1mod feature;
2pub mod token;
3pub mod key;
4
5use std::collections::HashMap;
6use std::fs::File;
7use std::io::{BufRead, BufReader};
8use std::path::Path;
9use uuid::Uuid;
10use crate::feature::{featured, SupportedDb};
11use crate::key::KeyFile;
12
13const CMD_FILE: &str = "file";
14const CMD_DB: &str = "db";
15const CMD_TBL: &str = "config";
16const CMD_TOKEN: &str = "token";
17const CMD_TTL: &str = "ttl";
18const CMD_KEY: &str = "key";
19const CMD_SECRET: &str = "secret";
20const CMD_PASS: &str = "PASSPHRASE";
21/// self host or node id
22const CMD_UUID: &str = "uuid";
23
24
25/// App startup args:
26/// - db connection url: usually the first arg
27///   - or prefix with: '--db '
28///     (optional, default postgres:// )
29///
30/// - config table name in format: schema.table, usually the second arg.
31///   - or prefix with: '--config '
32///     (optional, default public.{the_appname})
33///
34/// - token (db pwd) script name usually the third arg (required feature 'token')
35///   - or prefix with '--token '
36///
37/// - token live in minutes, usually the forth arg (required feature 'token')
38///   - or prefix with '--ttl '
39///
40/// - UUID this app instance to use as a node id or config recognition. Identified as UUID formatted string.
41///   - or prefix with '--uuid '
42///
43/// - (Private) Key text file name to use with RSA OR AES encryption (required feature 'rsa')
44///   - or prefix with '--key '
45///
46/// - cipher secret for AES taken from env \[SECRET\] OR cml --secret
47///  compile with keep_env_secret feature to not remove from env
48///
49/// Alternative configuration:
50/// - file name, usually the first arg
51///   - or prefix with '--file '
52///
53/// File format:
54///  - db: OR db=
55///  - config: OR config=
56///  - token: OR token=
57///  - uuid: OR uuid=
58///  - ttl: OR ttl=
59///  - pk: OR pk=
60///
61/// params passed in cmd line override params loaded from file & env.
62///
63/// env:
64/// PGPASSWORD, in case of postgres db url, use to connect to the DB
65/// PASSPHRASE, in case of RSA private key required a passphrase
66///
67#[derive(Debug, Clone)]
68pub struct ArgConfig {
69    /// instance ID
70    pub uuid: Uuid,
71    /// indicate the instance uuid was set or autogenerated on (every) start
72    pub uuid_gen: bool,
73    /// database connection string
74    pub db_url: String,
75
76    /// Format: schema.table  
77    pub table: String,
78    /// key=value loaded from file if present
79    pub cfg: HashMap<String, String>,
80    /// token (i.e. db pwd) script name usually the third arg (required feature 'token')
81    pub token: token::Token,
82    /// RSA private key file name to use with RSA OR AES encryption (required feature 'rsa')
83    pub pk: Option<KeyFile>,
84    /// cipher secret for AES taken from env \[SECRET\] OR cml --secret
85    /// use keep_env_secret feature to not remove from env
86    pub secret: Option<String>,
87}
88
89
90impl ArgConfig {
91
92    pub fn from_args() -> Result<ArgConfig, String> {
93        let user = match std::env::var_os("USER") {
94            Some(a) => a.to_str().unwrap_or("postgres").to_string(),
95            _ => "postgres".to_string(),
96        };
97        let f = featured();
98        let pwd = if !f.env_pwd().is_empty() {
99            match std::env::var_os(f.env_pwd().as_str()).map(|v| v) {
100                Some(a) => a.to_str().map(|v| v.to_string()),
101                _ => None,
102            }
103        } else {
104            None
105        };
106        let input: Vec<String> = std::env::args_os().map(|e| e.to_string_lossy().to_string()).collect();
107
108        ArgConfig::new(input, f, user, pwd)
109    }
110
111    // first arg is an app name itself
112    fn new(input: Vec<String>, feature: SupportedDb, user: String, pwd: Option<String>) -> Result<Self, String> {
113        let mut cfg = HashMap::new();
114        let mut db: Option<String> = None;
115        let mut tbl: Option<String>  = None;
116        let mut token: Option<String>  = None;
117
118        let mut ttl: Option<String>  = None;
119        let mut pk: Option<String>  = None;
120        let mut secret: Option<String>  = None;
121        let mut uuid: Option<Uuid> = None;
122        let mut ignore_next = true;
123        for i in 1..input.len() {
124            if ignore_next {continue}
125            if input[i].starts_with("--") {
126                if i < input.len() - 1 {
127                    let v = &input[i].as_str()[2..];
128                    if v == CMD_FILE {
129                        let _ = load(v, &mut cfg)?;
130                        ignore_next = true;
131                    } else if v == CMD_DB {
132                        db = Some(input[i + 1].to_string());
133                        ignore_next = true;
134                    } else if v == CMD_TBL {
135                        tbl = Some(input[i + 1].to_string());
136                        ignore_next = true;
137                    } else if v == CMD_TOKEN {
138                        token = Some(input[i + 1].to_string());
139                        ignore_next = true;
140                    } else if v == CMD_TTL {
141                        ttl = Some(input[i + 1].to_string());
142                        ignore_next = true;
143                    } else if v == CMD_KEY {
144                        let file = input[i + 1].to_string();
145                        if key::is_key_file(&file) {
146                            pk = Some(file);
147                            ignore_next = true;
148                        }
149                    } else if v == CMD_UUID {
150                        uuid = Uuid::parse_str(input[i + 1].as_str()).ok();
151                        ignore_next = true;
152                    } else if v == CMD_SECRET {
153                        secret = Some(input[i + 1].to_string());
154                        ignore_next = true;
155                    }
156                }
157            } else {
158            }
159        }
160        // first was a check by tag names, then try to guess
161        for i in &input {
162            if i.starts_with("--") {
163                continue
164            }
165            if db.is_none() && feature.is_valid_url(i) {
166                db = Some(i.to_string());
167                continue
168            }
169            if tbl.is_none() && is_sound_schema_table(i) {
170                tbl = Some(i.to_string());
171                continue
172            }
173            if token.is_none() && i.len() > 1 {
174                token = Some(i.to_string());
175                continue
176            }
177            if ttl.is_none() && i.parse::<u16>().is_ok() {
178                ttl = Some(i.to_string());
179            }
180            if uuid.is_none() {
181                uuid = Uuid::parse_str(i).ok();
182            }
183        }
184        if uuid.is_none() {
185            uuid  = Uuid::parse_str(get_env_or_cfg(CMD_UUID, &cfg, "").as_str()).ok();
186        }
187
188        if let Some(a) = std::env::var_os(CMD_SECRET.to_uppercase()).map(|v| v) {
189            secret = Some(a.to_str().map(|v| v.to_string()).unwrap_or("".to_string()));
190            #[cfg(not(feature="keep_env_secret"))]
191            unsafe { 
192                std::env::remove_var(CMD_SECRET); 
193            }
194        }
195
196        Ok(ArgConfig {
197            uuid_gen: uuid.is_none(),
198            uuid: uuid.unwrap_or(Uuid::new_v4()),
199            db_url: link_db_user(db.unwrap_or(get_env_or_cfg(CMD_DB, &cfg, feature.default_url(&user).as_str())), user),
200            table: tbl.unwrap_or(get_env_or_cfg(CMD_TBL, &cfg, get_exec_name("public.", input[0].as_str()).as_str())),
201            token: token::Token::new(
202                token.unwrap_or(get_env_or_cfg(CMD_TOKEN, &cfg, "")),
203                ttl.unwrap_or(get_env_or_cfg(CMD_TTL, &cfg, "1")),
204                pwd
205            )?,
206            pk: KeyFile::new(
207                pk.unwrap_or(get_env_or_cfg(CMD_KEY, &cfg, &"")),
208                std::env::var_os(CMD_PASS).map(|p| p.to_string_lossy().to_string()),
209            )?,
210            cfg,
211            secret,
212        })
213    }
214
215    /// append with password if $PWD present in 'db'
216    pub fn db_url(&self) -> String {
217        let url = self.db_url.clone();
218        if let Some(i) = url.find(":$P") {
219            if let Some(y) = url.find("@") {
220                let pwd = url.as_str()[i+1..y].to_owned();
221                return url.replace(
222                    pwd.as_str(),
223                    self.token.value.clone().unwrap_or("".into()).as_str()).to_string();
224            }
225        }
226        url
227    }
228}
229
230#[inline]
231fn link_db_user(url: String, user: String) -> String {
232    url.replace("$USER", user.as_str())
233}
234
235#[inline]
236fn get_env_or_cfg(input: &str, cfg: &HashMap<String, String>, def: &str) -> String {
237    match std::env::var_os(input).map(|v| v) {
238        Some(a) => a.to_str().map(|v| v.to_string()).unwrap_or(def.to_string()),
239        None => cfg.get(input).unwrap_or(&def.to_string()).into()
240    }
241}
242
243/// Safe taking value
244#[inline]
245fn get_exec_name(schema: &str, input: &str) -> String {
246    if input.is_empty() {
247        return "".to_string();
248    }
249    let e = Path::new(input);
250    let name = e.file_name().map(|f|f.to_str().unwrap_or("")).unwrap_or("").to_string();
251    #[cfg(windows)]
252    let name = name.replace(".exe", "");
253    let name = if let Some(i) = name.rfind(std::path::MAIN_SEPARATOR_STR) {
254        name[i+1..].to_string()
255    } else {
256        name
257    };
258    format!("{}{}", schema, name)
259}
260
261#[inline]
262fn is_sound_schema_table(input: &str) -> bool {
263    let y = input.contains(".");
264    #[cfg(windows)]
265    let y = y && !input.ends_with(".exe");
266    y
267}
268
269#[inline]
270fn load(file: &str, cfg: &mut HashMap<String, String>) -> Result<(), String> {
271    let f = File::open(file).map_err(|e| e.to_string())?;
272    let reader = BufReader::new(f);
273    for line in reader.lines() {
274        if let Ok(l) = line {
275            if let Some(i) = l.chars().position(|c| c == '=' || c == ':' || c == '#' || c == ';' || c == '/' || c == '[') {
276                if l.as_bytes()[i] != b'#' && l.as_bytes()[i] != b';' && l.as_bytes()[i] != b'/' && l.as_bytes()[i] != b'[' {
277                    let key = l[..i].trim().to_lowercase();
278                    let value = l[i + 1..].trim().to_string();
279                    cfg.insert(key, value);
280                }
281            }
282        }
283    }
284    Ok(())
285}
286
287#[allow(warnings)]
288#[cfg(test)]
289mod tests {
290    use super::*;
291
292
293    #[test]
294    #[cfg(unix)]
295    fn test_file_name() {
296        assert_eq!(get_exec_name("","").as_str(), "");
297        assert_eq!(get_exec_name("","target/debug/marg").as_str(), "marg");
298        assert_eq!(get_exec_name("","marg").as_str(), "marg");
299    }
300
301    #[test]
302    #[cfg(windows)]
303    fn test_file_name() {
304        assert_eq!(get_exec_name("","").as_str(), "");
305        assert_eq!(get_exec_name("","target\\debug\\marg.exe").as_str(), "marg");
306        assert_eq!(get_exec_name("","target\\\\debug\\\\marg.exe").as_str(), "marg");
307    }
308
309    #[test]
310    fn config_args_file1_test() {
311        //
312        let cfg = ArgConfig::new(
313            vec!["".to_string()],
314                 SupportedDb::Postgres, "".to_string(), None).unwrap();
315        assert_eq!(0, cfg.cfg.len());
316
317    }
318
319    #[test]
320    fn config_args_file2_test() {
321        let url =  "postgresql://user:pwd@host/db".to_string();
322        let cfg = ArgConfig::new(
323            vec![url.clone()],
324                 SupportedDb::Postgres, "vk".to_string(), None).unwrap();
325        assert_eq!(url, cfg.db_url());
326    }
327
328    #[test]
329    fn config_args_file3_test() {
330        let url =  "postgresql://user:pwd@host/db".to_string();
331        let t = "public.table".to_string();
332        let cfg = ArgConfig::new(
333            vec![url.clone(), t.clone()],
334                 SupportedDb::Postgres, "".to_string(), None).unwrap();
335        assert_eq!(url, cfg.db_url());
336        assert_eq!(t, cfg.table);
337    }
338
339    #[test]
340    fn config_args_user_test() {
341        let user = match std::env::var_os("USER") {
342            Some(a) => a.to_str().unwrap_or("postgres").to_string(),
343            _ => "postgres".to_string(),
344        };
345
346        assert_eq!(format!("postgresql://{}:pwd@host/db", user), link_db_user("postgresql://$USER:pwd@host/db".to_string(), user.clone()));
347        assert_eq!(format!("postgresql://{}:$PWD@host/db", user), link_db_user("postgresql://$USER:$PWD@host/db".to_string(), user.clone()));
348        assert_eq!(format!("postgresql://{}@host/db", user), link_db_user("postgresql://$USER@host/db".to_string(), user.clone()));
349        assert_eq!(format!("postgresql://{}@host/db", ""), link_db_user("postgresql://@host/db".to_string(), user.clone()));
350
351    }
352
353
354}
marg/lib.rs

marg/
lib.rs
