Skip to main content

cosq_client/
auth.rs

1//! Azure authentication via the Azure CLI
2//!
3//! Uses `az account get-access-token` to acquire tokens for Azure Resource Manager
4//! and Cosmos DB data plane access.
5
6use serde::Deserialize;
7use tokio::process::Command;
8
9use crate::error::ClientError;
10
11/// Cosmos DB data plane resource scope
12pub const COSMOS_RESOURCE: &str = "https://cosmos.azure.com";
13
14/// Azure Resource Manager resource scope
15pub const ARM_RESOURCE: &str = "https://management.azure.com";
16
17/// Status of the current Azure CLI authentication session
18#[derive(Debug, Clone)]
19pub struct AuthStatus {
20    pub logged_in: bool,
21    pub user: Option<String>,
22    pub subscription_name: Option<String>,
23    pub subscription_id: Option<String>,
24    pub tenant_id: Option<String>,
25}
26
27/// Azure CLI account info
28#[derive(Debug, Deserialize)]
29#[serde(rename_all = "camelCase")]
30struct AzAccountInfo {
31    user: AzUser,
32    name: String,
33    id: String,
34    tenant_id: String,
35}
36
37#[derive(Debug, Deserialize)]
38struct AzUser {
39    name: String,
40}
41
42/// Azure CLI-based authentication provider.
43pub struct AzCliAuth;
44
45impl AzCliAuth {
46    /// Check the current Azure CLI login status.
47    pub async fn check_status() -> Result<AuthStatus, ClientError> {
48        let output = Command::new("az")
49            .args(["account", "show", "--output", "json"])
50            .output()
51            .await
52            .map_err(|e| {
53                ClientError::az_cli(
54                    format!("failed to run `az` command: {e}"),
55                    "Install the Azure CLI: https://aka.ms/install-azure-cli",
56                )
57            })?;
58
59        if !output.status.success() {
60            let stderr = String::from_utf8_lossy(&output.stderr);
61            if stderr.contains("az login") || stderr.contains("not logged in") {
62                return Ok(AuthStatus {
63                    logged_in: false,
64                    user: None,
65                    subscription_name: None,
66                    subscription_id: None,
67                    tenant_id: None,
68                });
69            }
70            return Err(ClientError::az_cli(
71                stderr.trim().to_string(),
72                "Try running `az login` first",
73            ));
74        }
75
76        let info: AzAccountInfo =
77            serde_json::from_slice(&output.stdout).map_err(|e| ClientError::auth(e.to_string()))?;
78
79        Ok(AuthStatus {
80            logged_in: true,
81            user: Some(info.user.name),
82            subscription_name: Some(info.name),
83            subscription_id: Some(info.id),
84            tenant_id: Some(info.tenant_id),
85        })
86    }
87
88    /// Get an access token for the specified resource.
89    pub async fn get_token(resource: &str) -> Result<String, ClientError> {
90        let output = Command::new("az")
91            .args([
92                "account",
93                "get-access-token",
94                "--resource",
95                resource,
96                "--query",
97                "accessToken",
98                "--output",
99                "tsv",
100            ])
101            .output()
102            .await
103            .map_err(|e| {
104                ClientError::az_cli(
105                    format!("failed to run `az` command: {e}"),
106                    "Install the Azure CLI: https://aka.ms/install-azure-cli",
107                )
108            })?;
109
110        if !output.status.success() {
111            let stderr = String::from_utf8_lossy(&output.stderr);
112            return Err(ClientError::az_cli(
113                format!("failed to get access token: {}", stderr.trim()),
114                "Try running `az login` to refresh your credentials",
115            ));
116        }
117
118        let token = String::from_utf8_lossy(&output.stdout).trim().to_string();
119        if token.is_empty() {
120            return Err(ClientError::auth("received empty access token"));
121        }
122
123        Ok(token)
124    }
125
126    /// Run `az login` interactively.
127    pub async fn login() -> Result<(), ClientError> {
128        let status = Command::new("az")
129            .args(["login"])
130            .status()
131            .await
132            .map_err(|e| {
133                ClientError::az_cli(
134                    format!("failed to run `az login`: {e}"),
135                    "Install the Azure CLI: https://aka.ms/install-azure-cli",
136                )
137            })?;
138
139        if !status.success() {
140            return Err(ClientError::auth("az login failed"));
141        }
142
143        Ok(())
144    }
145
146    /// Get the signed-in user's principal (object) ID from Azure AD.
147    pub async fn get_principal_id() -> Result<String, ClientError> {
148        let output = Command::new("az")
149            .args([
150                "ad",
151                "signed-in-user",
152                "show",
153                "--query",
154                "id",
155                "--output",
156                "tsv",
157            ])
158            .output()
159            .await
160            .map_err(|e| {
161                ClientError::az_cli(
162                    format!("failed to run `az` command: {e}"),
163                    "Install the Azure CLI: https://aka.ms/install-azure-cli",
164                )
165            })?;
166
167        if !output.status.success() {
168            let stderr = String::from_utf8_lossy(&output.stderr);
169            return Err(ClientError::az_cli(
170                format!("failed to get principal ID: {}", stderr.trim()),
171                "Try running `az login` to refresh your credentials",
172            ));
173        }
174
175        let id = String::from_utf8_lossy(&output.stdout).trim().to_string();
176        if id.is_empty() {
177            return Err(ClientError::auth("received empty principal ID"));
178        }
179
180        Ok(id)
181    }
182
183    /// Run `az logout`.
184    pub async fn logout() -> Result<(), ClientError> {
185        let status = Command::new("az")
186            .args(["logout"])
187            .status()
188            .await
189            .map_err(|e| {
190                ClientError::az_cli(
191                    format!("failed to run `az logout`: {e}"),
192                    "Install the Azure CLI: https://aka.ms/install-azure-cli",
193                )
194            })?;
195
196        if !status.success() {
197            return Err(ClientError::auth("az logout failed"));
198        }
199
200        Ok(())
201    }
202}