Merge branch 'add-wait-flag-to-cli' into master

4 files changed, 99 insertions, 180 deletions

diff --git a/mullvad-cli/src/cmds/connect.rs b/mullvad-cli/src/cmds/connect.rs
index 6fc7072580..c93ed8c2e2 100644
--- a/mullvad-cli/src/cmds/connect.rs
+++ b/mullvad-cli/src/cmds/connect.rs
@@ -1,5 +1,6 @@
-use crate::{new_rpc_client, Command, Result};
-use talpid_types::ErrorExt;
+use crate::{format, new_rpc_client, state, Command, Error, Result};
+use futures::StreamExt;
+use mullvad_management_interface::types::tunnel_state::State;
 
 pub struct Connect;
 
@@ -12,13 +13,38 @@ impl Command for Connect {
     fn clap_subcommand(&self) -> clap::App<'static, 'static> {
         clap::SubCommand::with_name(self.name())
             .about("Command the client to start establishing a VPN tunnel")
+            .arg(
+                clap::Arg::with_name("wait")
+                    .long("wait")
+                    .short("w")
+                    .help("Wait until connected before exiting"),
+            )
     }
 
-    async fn run(&self, _: &clap::ArgMatches<'_>) -> Result<()> {
+    async fn run(&self, matches: &clap::ArgMatches<'_>) -> Result<()> {
         let mut rpc = new_rpc_client().await?;
-        if let Err(e) = rpc.connect_tunnel(()).await {
-            eprintln!("{}", e.display_chain());
+
+        let receiver_option = if matches.is_present("wait") {
+            Some(state::state_listen(rpc.clone()))
+        } else {
+            None
+        };
+
+        if rpc.connect_tunnel(()).await?.into_inner() {
+            if let Some(mut receiver) = receiver_option {
+                while let Some(state) = receiver.next().await {
+                    let state = state?;
+                    format::print_state(&state);
+                    match state.state.unwrap() {
+                        State::Connected(_) => return Ok(()),
+                        State::Error(_) => return Err(Error::CommandFailed("connect")),
+                        _ => {}
+                    }
+                }
+                return Err(Error::StatusListenerFailed);
+            }
         }
+
         Ok(())
     }
 }
diff --git a/mullvad-cli/src/cmds/disconnect.rs b/mullvad-cli/src/cmds/disconnect.rs
index f976b96fab..96d5843d20 100644
--- a/mullvad-cli/src/cmds/disconnect.rs
+++ b/mullvad-cli/src/cmds/disconnect.rs
@@ -1,4 +1,6 @@
-use crate::{new_rpc_client, Command, Result};
+use crate::{format, new_rpc_client, state, Command, Error, Result};
+use futures::StreamExt;
+use mullvad_management_interface::types::tunnel_state::State::Disconnected;
 
 pub struct Disconnect;
 
@@ -11,11 +13,37 @@ impl Command for Disconnect {
     fn clap_subcommand(&self) -> clap::App<'static, 'static> {
         clap::SubCommand::with_name(self.name())
             .about("Command the client to disconnect the VPN tunnel")
+            .arg(
+                clap::Arg::with_name("wait")
+                    .long("wait")
+                    .short("w")
+                    .help("Wait until disconnected before exiting"),
+            )
     }
 
-    async fn run(&self, _: &clap::ArgMatches<'_>) -> Result<()> {
+    async fn run(&self, matches: &clap::ArgMatches<'_>) -> Result<()> {
         let mut rpc = new_rpc_client().await?;
-        rpc.disconnect_tunnel(()).await?;
+
+        let receiver_option = if matches.is_present("wait") {
+            Some(state::state_listen(rpc.clone()))
+        } else {
+            None
+        };
+
+        if rpc.disconnect_tunnel(()).await?.into_inner() {
+            if let Some(mut receiver) = receiver_option {
+                while let Some(state) = receiver.next().await {
+                    let state = state?;
+                    format::print_state(&state);
+                    match state.state.unwrap() {
+                        Disconnected(_) => return Ok(()),
+                        _ => {}
+                    }
+                }
+                return Err(Error::StatusListenerFailed);
+            }
+        }
+
         Ok(())
     }
 }
diff --git a/mullvad-cli/src/cmds/reconnect.rs b/mullvad-cli/src/cmds/reconnect.rs
index ecc0b089c0..87de63144f 100644
--- a/mullvad-cli/src/cmds/reconnect.rs
+++ b/mullvad-cli/src/cmds/reconnect.rs
@@ -1,5 +1,6 @@
-use crate::{new_rpc_client, Command, Result};
-use talpid_types::ErrorExt;
+use crate::{format, new_rpc_client, state, Command, Error, Result};
+use futures::StreamExt;
+use mullvad_management_interface::types::tunnel_state::State;
 
 pub struct Reconnect;
 
@@ -10,14 +11,40 @@ impl Command for Reconnect {
     }
 
     fn clap_subcommand(&self) -> clap::App<'static, 'static> {
-        clap::SubCommand::with_name(self.name()).about("Command the client to reconnect")
+        clap::SubCommand::with_name(self.name())
+            .about("Command the client to reconnect")
+            .arg(
+                clap::Arg::with_name("wait")
+                    .long("wait")
+                    .short("w")
+                    .help("Wait until reconnected before exiting"),
+            )
     }
 
-    async fn run(&self, _: &clap::ArgMatches<'_>) -> Result<()> {
+    async fn run(&self, matches: &clap::ArgMatches<'_>) -> Result<()> {
         let mut rpc = new_rpc_client().await?;
-        if let Err(e) = rpc.reconnect_tunnel(()).await {
-            eprintln!("{}", e.display_chain());
+
+        let receiver_option = if matches.is_present("wait") {
+            Some(state::state_listen(rpc.clone()))
+        } else {
+            None
+        };
+
+        if rpc.reconnect_tunnel(()).await?.into_inner() {
+            if let Some(mut receiver) = receiver_option {
+                while let Some(state) = receiver.next().await {
+                    let state = state?;
+                    format::print_state(&state);
+                    match state.state.unwrap() {
+                        State::Connected(_) => return Ok(()),
+                        State::Error(_) => return Err(Error::CommandFailed("reconnect")),
+                        _ => {}
+                    }
+                }
+                return Err(Error::StatusListenerFailed);
+            }
         }
+
         Ok(())
     }
 }
diff --git a/mullvad-cli/src/cmds/status.rs b/mullvad-cli/src/cmds/status.rs
index b0c4d40ad0..76a4ed2e67 100644
--- a/mullvad-cli/src/cmds/status.rs
+++ b/mullvad-cli/src/cmds/status.rs
@@ -1,17 +1,7 @@
-use crate::{format::print_keygen_event, new_rpc_client, Command, Error, Result};
+use crate::{format, format::print_keygen_event, new_rpc_client, Command, Error, Result};
 use mullvad_management_interface::{
-    types::{
-        daemon_event::Event as EventType,
-        error_state::{
-            firewall_policy_error::ErrorType as FirewallPolicyErrorType, Cause as ErrorStateCause,
-            FirewallPolicyError, GenerationError,
-        },
-        ErrorState, ProxyType, TransportProtocol, TunnelEndpoint, TunnelState, TunnelType,
-    },
-    ManagementServiceClient,
+    types::daemon_event::Event as EventType, ManagementServiceClient,
 };
-use mullvad_types::auth_failed::AuthFailed;
-use std::fmt::Write;
 
 pub struct Status;
 
@@ -45,7 +35,7 @@ impl Command for Status {
         let mut rpc = new_rpc_client().await?;
         let state = rpc.get_tunnel_state(()).await?.into_inner();
 
-        print_state(&state);
+        format::print_state(&state);
         if matches.is_present("location") {
             print_location(&mut rpc).await?;
         }
@@ -58,7 +48,7 @@ impl Command for Status {
             while let Some(event) = events.message().await? {
                 match event.event.unwrap() {
                     EventType::TunnelState(new_state) => {
-                        print_state(&new_state);
+                        format::print_state(&new_state);
                         use mullvad_management_interface::types::tunnel_state::State::*;
                         match new_state.state.unwrap() {
                             Connected(..) | Disconnected(..) => {
@@ -98,151 +88,6 @@ impl Command for Status {
     }
 }
 
-fn print_state(state: &TunnelState) {
-    use mullvad_management_interface::types::{tunnel_state, tunnel_state::State::*};
-
-    print!("Tunnel status: ");
-    match state.state.as_ref().unwrap() {
-        Error(error) => print_error_state(error.error_state.as_ref().unwrap()),
-        Connected(tunnel_state::Connected { relay_info }) => {
-            let endpoint = relay_info
-                .as_ref()
-                .unwrap()
-                .tunnel_endpoint
-                .as_ref()
-                .unwrap();
-            println!("Connected to {}", format_endpoint(&endpoint));
-        }
-        Connecting(tunnel_state::Connecting { relay_info }) => {
-            let endpoint = relay_info
-                .as_ref()
-                .unwrap()
-                .tunnel_endpoint
-                .as_ref()
-                .unwrap();
-            println!("Connecting to {}...", format_endpoint(&endpoint));
-        }
-        Disconnected(_) => println!("Disconnected"),
-        Disconnecting(_) => println!("Disconnecting..."),
-    }
-}
-
-fn format_endpoint(endpoint: &TunnelEndpoint) -> String {
-    let mut out = format!(
-        "{} {} over {}",
-        match TunnelType::from_i32(endpoint.tunnel_type).expect("unknown tunnel protocol") {
-            TunnelType::Wireguard => "WireGuard",
-            TunnelType::Openvpn => "OpenVPN",
-        },
-        endpoint.address,
-        format_protocol(
-            TransportProtocol::from_i32(endpoint.protocol).expect("unknown transport protocol")
-        ),
-    );
-
-    if let Some(ref proxy) = endpoint.proxy {
-        write!(
-            &mut out,
-            " via {} {} over {}",
-            match ProxyType::from_i32(proxy.proxy_type).expect("unknown proxy type") {
-                ProxyType::Shadowsocks => "Shadowsocks",
-                ProxyType::Custom => "custom bridge",
-            },
-            proxy.address,
-            format_protocol(
-                TransportProtocol::from_i32(proxy.protocol).expect("unknown transport protocol")
-            ),
-        )
-        .unwrap();
-    }
-
-    out
-}
-
-fn print_error_state(error_state: &ErrorState) {
-    if error_state.blocking_error.is_some() {
-        eprintln!("Mullvad daemon failed to setup firewall rules!");
-        eprintln!("Deamon cannot block traffic from flowing, non-local traffic will leak");
-    }
-
-    match ErrorStateCause::from_i32(error_state.cause) {
-        Some(ErrorStateCause::AuthFailed) => {
-            println!(
-                "Blocked: {}",
-                AuthFailed::from(error_state.auth_fail_reason.as_ref())
-            );
-        }
-        #[cfg(target_os = "linux")]
-        Some(ErrorStateCause::SetFirewallPolicyError) => {
-            println!("Blocked: {}", error_state_to_string(error_state));
-            println!("Your kernel might be terribly out of date or missing nftables");
-        }
-        _ => println!("Blocked: {}", error_state_to_string(error_state)),
-    }
-}
-
-fn error_state_to_string(error_state: &ErrorState) -> String {
-    use ErrorStateCause::*;
-
-    let error_str = match ErrorStateCause::from_i32(error_state.cause).expect("unknown error cause")
-    {
-        AuthFailed => {
-            return if error_state.auth_fail_reason.is_empty() {
-                "Authentication with remote server failed".to_string()
-            } else {
-                format!(
-                    "Authentication with remote server failed: {}",
-                    error_state.auth_fail_reason
-                )
-            };
-        }
-        Ipv6Unavailable => "Failed to configure IPv6 because it's disabled in the platform",
-        SetFirewallPolicyError => {
-            return policy_error_to_string(error_state.policy_error.as_ref().unwrap())
-        }
-        SetDnsError => "Failed to set system DNS server",
-        StartTunnelError => "Failed to start connection to remote server",
-        TunnelParameterError => {
-            return format!(
-                "Failure to generate tunnel parameters: {}",
-                tunnel_parameter_error_to_string(error_state.parameter_error)
-            );
-        }
-        IsOffline => "This device is offline, no tunnels can be established",
-        TapAdapterProblem => "A problem with the TAP adapter has been detected",
-        #[cfg(target_os = "android")]
-        VpnPermissionDenied => "The Android VPN permission was denied when creating the tunnel",
-        #[cfg(not(target_os = "android"))]
-        _ => unreachable!("unknown error cause"),
-    };
-
-    error_str.to_string()
-}
-
-fn tunnel_parameter_error_to_string(parameter_error: i32) -> &'static str {
-    match GenerationError::from_i32(parameter_error).expect("unknown generation error") {
-        GenerationError::NoMatchingRelay => "Failure to select a matching tunnel relay",
-        GenerationError::NoMatchingBridgeRelay => "Failure to select a matching bridge relay",
-        GenerationError::NoWireguardKey => "No wireguard key available",
-        GenerationError::CustomTunnelHostResolutionError => {
-            "Can't resolve hostname for custom tunnel host"
-        }
-    }
-}
-
-fn policy_error_to_string(policy_error: &FirewallPolicyError) -> String {
-    let cause = match FirewallPolicyErrorType::from_i32(policy_error.r#type)
-        .expect("unknown policy error")
-    {
-        FirewallPolicyErrorType::Generic => return "Failed to set firewall policy".to_string(),
-        FirewallPolicyErrorType::Locked => format!(
-            "An application prevented the firewall policy from being set: {} (pid {})",
-            policy_error.lock_name, policy_error.lock_pid
-        ),
-    };
-    format!("Failed to set firewall policy: {}", cause)
-}
-
 async fn print_location(rpc: &mut ManagementServiceClient) -> Result<()> {
     let location = rpc.get_current_location(()).await;
     let location = match location {
@@ -278,10 +123,3 @@ async fn print_location(rpc: &mut ManagementServiceClient) -> Result<()> {
     );
     Ok(())
 }
-
-fn format_protocol(protocol: TransportProtocol) -> &'static str {
-    match protocol {
-        TransportProtocol::Udp => "UDP",
-        TransportProtocol::Tcp => "TCP",
-    }
-}