From 069f72f7748756be4c71392211693b911165f86d Mon Sep 17 00:00:00 2001
From: Joe Prosser <joe.prosser@uipath.com>
Date: Wed, 3 Apr 2024 18:47:57 +0100
Subject: [PATCH] feat(cli): add feature to make context a required field

---
 CHANGELOG.md               |  4 ++++
 cli/src/commands/config.rs | 14 +++++++++++++-
 cli/src/config.rs          |  2 ++
 cli/src/main.rs            | 36 ++++++++++++++++++++++++++++++++++--
 4 files changed, 53 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3360b94f..843a19cc 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,3 +1,7 @@
+# Unreleased
+- Add option to make context option a required field
+- Make context option a required field for internal users
+
 # v0.23.0
 
 - Add `get emails`
diff --git a/cli/src/commands/config.rs b/cli/src/commands/config.rs
index 44f52c4c..6a11bce3 100644
--- a/cli/src/commands/config.rs
+++ b/cli/src/commands/config.rs
@@ -7,7 +7,7 @@ use std::path::Path;
 use structopt::StructOpt;
 
 use crate::{
-    config::{self, ContextConfig, ReinferConfig},
+    config::{self, write_reinfer_config, ContextConfig, ReinferConfig},
     utils,
 };
 use anyhow::{anyhow, Result};
@@ -68,6 +68,14 @@ pub enum ConfigArgs {
         /// The name of the context.
         name: String,
     },
+
+    #[structopt(name = "set-context-required")]
+    /// Set whether  context is a required field
+    SetContextRequired {
+        // Whether the context is a required field
+        #[structopt(name = "is-required", parse(try_from_str))]
+        is_required: bool,
+    },
 }
 
 pub fn run(
@@ -76,6 +84,10 @@ pub fn run(
     config_path: impl AsRef<Path>,
 ) -> Result<ReinferConfig> {
     match args {
+        ConfigArgs::SetContextRequired { is_required } => {
+            config.context_is_required = *is_required;
+            write_reinfer_config(config_path, &config)?
+        }
         ConfigArgs::ListContexts { tokens } if config.num_contexts() > 0 => {
             let mut contexts = config.get_all_contexts().clone();
             contexts.sort_unstable_by(|lhs, rhs| lhs.name.cmp(&rhs.name));
diff --git a/cli/src/config.rs b/cli/src/config.rs
index 2205e8c2..5cbaa6c5 100644
--- a/cli/src/config.rs
+++ b/cli/src/config.rs
@@ -12,6 +12,8 @@ use std::{
 pub struct ReinferConfig {
     current_context: Option<String>,
     contexts: Vec<ContextConfig>,
+    #[serde(default)]
+    pub context_is_required: bool,
 }
 
 impl ReinferConfig {
diff --git a/cli/src/main.rs b/cli/src/main.rs
index 5de89fff..0f321b8c 100644
--- a/cli/src/main.rs
+++ b/cli/src/main.rs
@@ -134,14 +134,46 @@ fn client_from_args(args: &Args, config: &ReinferConfig) -> Result<Client> {
         backoff_factor: 2.0,
     };
 
-    Client::new(ClientConfig {
+    let client = Client::new(ClientConfig {
         endpoint,
         token,
         accept_invalid_certificates,
         proxy,
         retry_config: Some(retry_config),
     })
-    .context("Failed to initialise the HTTP client.")
+    .context("Failed to initialise the HTTP client.")?;
+
+    check_if_context_is_a_required_field(config, &client, &args)?;
+
+    Ok(client)
+}
+
+const DOMAINS_THAT_REQUIRE_CONTEXT: [&str; 2] = ["uipath.com", "reinfer.io"];
+
+fn check_if_context_is_a_required_field(
+    config: &ReinferConfig,
+    client: &Client,
+    args: &Args,
+) -> Result<()> {
+    if config.context_is_required {
+        return Err(anyhow!(
+            "Please provide a context with the `re -c <context>` option"
+        ));
+    }
+
+    let current_user = client.get_current_user()?;
+
+    if DOMAINS_THAT_REQUIRE_CONTEXT
+        .iter()
+        .any(|domain| current_user.email.0.ends_with(domain))
+        && args.context.is_none()
+    {
+        return Err(anyhow!(
+            "As a UiPath user, please provide a context with the `re -c <context>` option"
+        ));
+    };
+
+    Ok(())
 }
 
 fn find_configuration(args: &Args) -> Result<PathBuf> {