use anyhow::{Context, Result};
use serde::Deserialize;
use std::process::Command;
use std::sync::Mutex;
use std::time::{SystemTime, UNIX_EPOCH};
static TOKEN_CACHE: Mutex<Option<CachedToken>> = Mutex::new(None);
static STORAGE_TOKEN_CACHE: Mutex<Option<CachedToken>> = Mutex::new(None);
const FABRIC_RESOURCE: &str = "https://analysis.windows.net/powerbi/api";
const STORAGE_RESOURCE: &str = "https://storage.azure.com";
#[derive(Debug, Deserialize, Clone)]
#[serde(rename_all = "camelCase")]
pub struct AzToken {
pub access_token: String,
pub expires_on: String,
pub tenant: Option<String>,
pub subscription: Option<String>,
}
#[derive(Clone)]
struct CachedToken {
token: String,
expires_at: u64,
}
#[derive(Debug)]
pub struct AuthStatus {
pub tenant: String,
pub subscription: String,
pub token_valid: bool,
pub expires_on: String,
}
fn az_get_token(resource: &str) -> Result<AzToken> {
let output = Command::new("az")
.args([
"account",
"get-access-token",
"--resource",
resource,
"--output",
"json",
])
.output()
.context("Failed to run 'az' CLI; is Azure CLI installed?")?;
if !output.status.success() {
let stderr = String::from_utf8_lossy(&output.stderr);
anyhow::bail!("az account get-access-token failed: {}", stderr.trim());
}
let token: AzToken =
serde_json::from_slice(&output.stdout).context("Failed to parse az CLI token output")?;
Ok(token)
}
fn parse_expires_on(expires_on: &str) -> u64 {
expires_on.trim().parse::<u64>().unwrap_or(0)
}
fn now_epoch() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs()
}
pub fn get_fabric_token() -> Result<String> {
get_cached_token(&TOKEN_CACHE, FABRIC_RESOURCE)
}
pub fn get_storage_token() -> Result<String> {
get_cached_token(&STORAGE_TOKEN_CACHE, STORAGE_RESOURCE)
}
fn get_cached_token(cache: &Mutex<Option<CachedToken>>, resource: &str) -> Result<String> {
let mut lock = cache.lock().unwrap();
if let Some(ref cached) = *lock {
if now_epoch() + 60 < cached.expires_at {
return Ok(cached.token.clone());
}
}
let az = az_get_token(resource)?;
let expires_at = parse_expires_on(&az.expires_on);
let token = az.access_token.clone();
*lock = Some(CachedToken {
token: token.clone(),
expires_at,
});
Ok(token)
}
pub fn check_auth_status() -> Result<AuthStatus> {
let az = az_get_token(FABRIC_RESOURCE)?;
Ok(AuthStatus {
tenant: az.tenant.unwrap_or_else(|| "unknown".into()),
subscription: az.subscription.unwrap_or_else(|| "unknown".into()),
token_valid: parse_expires_on(&az.expires_on) > now_epoch(),
expires_on: az.expires_on,
})
}