Support falling back to environment variables if command line arguments don't exist

This commit is contained in:
Travis Burtrum 2019-12-21 01:48:48 -05:00
parent 327a23d82d
commit 1650cefa0d
4 changed files with 86 additions and 13 deletions

View File

@ -49,6 +49,18 @@ usage: wireguard-proxy [options...]
-h, --help print this usage text
-st, --socket-timeout <seconds> Socket timeout (time to wait for data)
before terminating, default: 0
Environment variable support:
For every command line option, short and long, if you replace all
leading - with WGP_, and replace all remaining - with _, and uppercase
the whole thing, if you don't specify that command line option we will
read that environment variable for the argument. boolean arguments are
true if anything but unset, empty, 0, or false.
Examples:
--tcp-target ARG is WGP_TCP_TARGET=ARG
--socket-timeout 5 is WGP_SOCKET_TIMEOUT=5
--tls is WGP_TLS=1 or WGP_TLS=true
WGP_TLS=0 or WGP_TLS=false would be like not sending --tls
```
Binaries:

View File

@ -124,6 +124,18 @@ fn main() {
sha256 hashes preceded by "sha256//"
and separated by ";". Identical to curl's
--pinnedpubkey and CURLOPT_PINNEDPUBLICKEY
Environment variable support:
For every command line option, short and long, if you replace all
leading - with WGP_, and replace all remaining - with _, and uppercase
the whole thing, if you don't specify that command line option we will
read that environment variable for the argument. boolean arguments are
true if anything but unset, empty, 0, or false.
Examples:
--tcp-target ARG is WGP_TCP_TARGET=ARG
--socket-timeout 5 is WGP_SOCKET_TIMEOUT=5
--tls is WGP_TLS=1 or WGP_TLS=true
WGP_TLS=0 or WGP_TLS=false would be like not sending --tls
"#, default_udp_host_target, default_udp_host_target, default_socket_timeout);
return;
} else if args.flag("-s") || args.flag("--self-test") {
@ -137,10 +149,12 @@ fn main() {
let mut proxyd_args = vec!["-th", tcp_host, "-ut", host];
let tls_key = tls_key.as_ref().map(String::as_str);
let tls_cert = tls_cert.as_ref().map(String::as_str);
if tls {
let tls_key = tls_key.unwrap();
let tls_cert = tls_cert.unwrap();
proxyd_args.extend(["-tk", tls_key, "-tc", tls_cert].iter().cloned());
proxyd_args.extend(["-tk", &tls_key, "-tc", &tls_cert].iter().cloned());
}
println!("executing: {} {}", proxy, proxyd_args.join(" "));
@ -153,11 +167,12 @@ fn main() {
let mut proxy_args = vec!["-tt", tcp_host];
let pinnedpubkey = pinnedpubkey.as_ref().map(String::as_str);
if tls {
proxy_args.push("--tls");
if pinnedpubkey.is_some() {
proxy_args.push("--pinnedpubkey");
proxy_args.push(pinnedpubkey.unwrap());
proxy_args.push(&pinnedpubkey.unwrap());
}
}
@ -236,6 +251,7 @@ fn main() {
if tls {
let hostname = tcp_host.split(":").next();
let pinnedpubkey = pinnedpubkey.as_ref().map(String::as_str);
match pinnedpubkey {
Some(pinnedpubkey) =>
println!("executing: wireguard-proxy -tt {} --tls --pinnedpubkey {}", tcp_host, pinnedpubkey),
@ -253,7 +269,7 @@ fn main() {
println!("waiting: {:?} for wireguard-proxy client to come up.....", sleep);
thread::sleep(sleep);
first_arg = host;
first_arg = host.to_owned();
}
let server = Server::new(

View File

@ -54,6 +54,18 @@ fn main() {
-h, --help print this usage text
-st, --socket-timeout <seconds> Socket timeout (time to wait for data)
before terminating, default: {}
Environment variable support:
For every command line option, short and long, if you replace all
leading - with WGP_, and replace all remaining - with _, and uppercase
the whole thing, if you don't specify that command line option we will
read that environment variable for the argument. boolean arguments are
true if anything but unset, empty, 0, or false.
Examples:
--tcp-target ARG is WGP_TCP_TARGET=ARG
--socket-timeout 5 is WGP_SOCKET_TIMEOUT=5
--tls is WGP_TLS=1 or WGP_TLS=true
WGP_TLS=0 or WGP_TLS=false would be like not sending --tls
"#, default_udp_host_target, default_udp_host_target, default_socket_timeout);
return;
}
@ -61,9 +73,9 @@ fn main() {
let socket_timeout = args.get(&["-st", "--socket-timeout"], default_socket_timeout);
if tcp_target.is_some() {
client(tcp_target.unwrap(), socket_timeout, args);
client(&tcp_target.unwrap(), socket_timeout, args);
} else {
server(tcp_host.unwrap(), socket_timeout, args);
server(&tcp_host.unwrap(), socket_timeout, args);
}
}
@ -85,9 +97,9 @@ fn client(tcp_target: &str, socket_timeout: u64, args: Args) {
);
if tls {
let hostname = args.get_option(&["--tls-hostname"]).or_else(|| tcp_target.split(":").next());
let hostname = args.get_option(&["--tls-hostname"]).or_else(|| tcp_target.split(":").next().map(&str::to_owned));
let pinnedpubkey = args.get_option(&["--pinnedpubkey"]);
proxy_client.start_tls(hostname, pinnedpubkey).expect("error running tls proxy_client");
proxy_client.start_tls(hostname.as_ref().map(String::as_str), pinnedpubkey.as_ref().map(String::as_str)).expect("error running tls proxy_client");
} else {
proxy_client.start().expect("error running proxy_client");
}
@ -138,7 +150,7 @@ fn server(tcp_host: &str, socket_timeout: u64, args: Args) {
);
if tls_key.is_some() && tls_cert.is_some() {
proxy_server.start_tls(tls_key.unwrap(), tls_cert.unwrap()).expect("error running TLS proxy_server");
proxy_server.start_tls(&tls_key.unwrap(), &tls_cert.unwrap()).expect("error running TLS proxy_server");
} else if tls_key.is_none() && tls_cert.is_none() {
proxy_server.start().expect("error running proxy_server");
} else {

View File

@ -24,6 +24,25 @@ mod tls {
use tls::{TlsStream, TlsListener};
fn arg_to_env<'a>(arg: &'a str) -> String {
let env = "WGP_".to_owned();
let mut env = env + &arg.trim_matches('-').replace("-", "_");
env.make_ascii_uppercase();
env
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_arg_to_env() {
assert_eq!(arg_to_env("--tcp-host"), "WGP_TCP_HOST");
assert_eq!(arg_to_env("--tls"), "WGP_TLS");
assert_eq!(arg_to_env("-h"), "WGP_H");
}
}
pub struct Args<'a> {
args: &'a Vec<String>,
}
@ -33,26 +52,40 @@ impl<'a> Args<'a> {
Args { args }
}
pub fn flag(&self, flag: &'a str) -> bool {
self.args.contains(&flag.to_owned())
if self.args.contains(&flag.to_owned()) {
return true;
}
// because env we want slightly special handling of empty/0/false
match std::env::var(arg_to_env(flag)) {
Ok(env) => &env != "" && &env != "0" && &env != "false",
Err(_) => false,
}
}
pub fn get_option(&self, flags: &[&'a str]) -> Option<&'a str> {
pub fn get_option(&self, flags: &[&'a str]) -> Option<String> {
for flag in flags.iter() {
let mut found = false;
for arg in self.args.iter() {
if found {
return Some(arg);
return Some(arg.to_owned());
}
if arg == flag {
found = true;
}
}
}
// no matching arguments are found, so check env variables as a fallback
for flag in flags.iter() {
let env = std::env::var(arg_to_env(flag)).ok();
if env.is_some() {
return env;
}
}
return None;
}
pub fn get_str(&self, flags: &[&'a str], def: &'a str) -> &'a str {
pub fn get_str(&self, flags: &[&'a str], def: &'a str) -> String {
match self.get_option(flags) {
Some(ret) => ret,
None => def,
None => def.to_owned(),
}
}
pub fn get<T: FromStr>(&self, flags: &[&'a str], def: T) -> T {