1 files changed, 31 insertions, 5 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(())
     }
 }