1 files changed, 43 insertions, 14 deletions

diff --git a/mullvad-cli/src/cmds/account.rs b/mullvad-cli/src/cmds/account.rs
index 4b5a4f1fbd..4507b555c0 100644
--- a/mullvad-cli/src/cmds/account.rs
+++ b/mullvad-cli/src/cmds/account.rs
@@ -2,7 +2,7 @@ use crate::{new_rpc_client, Command, Error, Result};
 use itertools::Itertools;
 use mullvad_management_interface::{
     types::{self, Timestamp},
-    Code,
+    Code, ManagementServiceClient,
 };
 use mullvad_types::account::AccountToken;
 use std::io::{self, Write};
@@ -22,7 +22,7 @@ impl Command for Account {
             .subcommand(clap::App::new("create").about("Create and log in to a new account"))
             .subcommand(
                 clap::App::new("login").about("Log in to an account").arg(
-                    clap::Arg::new("token")
+                    clap::Arg::new("account")
                         .help("The Mullvad account token to configure the client with")
                         .required(false),
                 ),
@@ -35,18 +35,20 @@ impl Command for Account {
                 clap::App::new("list-devices")
                     .about("List devices associated with an account")
                     .arg(
-                        clap::Arg::new("token")
+                        clap::Arg::new("account")
                             .help("Mullvad account number")
-                            .required(true),
+                            .long("account")
+                            .takes_value(true),
                     ),
             )
             .subcommand(
                 clap::App::new("revoke-device")
                     .about("Revoke a device associated with an account")
                     .arg(
-                        clap::Arg::new("token")
+                        clap::Arg::new("account")
                             .help("Mullvad account number")
-                            .required(true),
+                            .long("account")
+                            .takes_value(true),
                     )
                     .arg(
                         clap::Arg::new("device")
@@ -67,16 +69,15 @@ impl Command for Account {
         if let Some(_matches) = matches.subcommand_matches("create") {
             self.create().await
         } else if let Some(set_matches) = matches.subcommand_matches("login") {
-            self.login(parse_token(set_matches)).await
+            self.login(parse_token_else_stdin(set_matches)).await
         } else if let Some(_matches) = matches.subcommand_matches("logout") {
             self.logout().await
         } else if let Some(_matches) = matches.subcommand_matches("get") {
             self.get().await
         } else if let Some(set_matches) = matches.subcommand_matches("list-devices") {
-            self.list_devices(parse_token(set_matches)).await
+            self.list_devices(set_matches).await
         } else if let Some(set_matches) = matches.subcommand_matches("revoke-device") {
-            self.revoke_device(parse_token(set_matches), parse_device_id(set_matches))
-                .await
+            self.revoke_device(set_matches).await
         } else if let Some(matches) = matches.subcommand_matches("redeem") {
             let voucher = matches.value_of_t_or_exit("voucher");
             self.redeem_voucher(voucher).await
@@ -129,8 +130,9 @@ impl Account {
         Ok(())
     }
 
-    async fn list_devices(&self, token: String) -> Result<()> {
+    async fn list_devices(&self, matches: &clap::ArgMatches) -> Result<()> {
         let mut rpc = new_rpc_client().await?;
+        let token = self.parse_account_else_current(&mut rpc, matches).await?;
         let devices = rpc.list_devices(token).await?.into_inner();
 
         println!("{:?}", devices);
@@ -138,8 +140,12 @@ impl Account {
         Ok(())
     }
 
-    async fn revoke_device(&self, token: String, device_id: String) -> Result<()> {
+    async fn revoke_device(&self, matches: &clap::ArgMatches) -> Result<()> {
         let mut rpc = new_rpc_client().await?;
+
+        let token = self.parse_account_else_current(&mut rpc, matches).await?;
+        let device_id = parse_device_id(matches);
+
         rpc.remove_device(types::DeviceRemoval {
             account_token: token,
             device_id,
@@ -149,6 +155,29 @@ impl Account {
         Ok(())
     }
 
+    async fn parse_account_else_current(
+        &self,
+        rpc: &mut ManagementServiceClient,
+        matches: &clap::ArgMatches,
+    ) -> Result<String> {
+        match matches.value_of("account").map(str::to_string) {
+            Some(token) => Ok(token),
+            None => {
+                let device = rpc
+                    .get_device(())
+                    .await
+                    .map_err(|error| match error.code() {
+                        mullvad_management_interface::Code::NotFound => {
+                            Error::CommandFailed("Log in or specify an account")
+                        }
+                        _ => Error::RpcFailedExt("Failed to obtain device", error),
+                    })?
+                    .into_inner();
+                Ok(device.account_token)
+            }
+        }
+    }
+
     async fn redeem_voucher(&self, mut voucher: String) -> Result<()> {
         let mut rpc = new_rpc_client().await?;
         voucher.retain(|c| c.is_alphanumeric());
@@ -198,8 +227,8 @@ impl Account {
     }
 }
 
-fn parse_token(matches: &clap::ArgMatches) -> String {
-    parse_from_match_else_stdin("Enter account token: ", "token", matches)
+fn parse_token_else_stdin(matches: &clap::ArgMatches) -> String {
+    parse_from_match_else_stdin("Enter account number: ", "account", matches)
 }
 
 fn parse_device_id(matches: &clap::ArgMatches) -> String {