use std::collections::BTreeMap;
use std::fs;
use std::path::PathBuf;
use std::time::{SystemTime, UNIX_EPOCH};
use serde::{Deserialize, Serialize};
use tirith_core::context_detect::{self, Provider};
use tirith_core::sudo_session;
const CACHE_TTL_SECS: u64 = 30;
#[derive(Debug, Serialize, Deserialize)]
struct CacheEnvelope {
#[serde(default = "default_schema_version")]
schema_version: u32,
captured_at: u64,
protection_mode: String,
contexts: BTreeMap<String, String>,
ssh_remote: bool,
sudo_active: bool,
#[serde(default)]
env_fingerprint: String,
}
fn default_schema_version() -> u32 {
1
}
fn current_env_fingerprint() -> String {
use sha2::{Digest, Sha256};
let mut h = Sha256::new();
for name in [
"TIRITH_STATUS",
"TIRITH_SSH_REMOTE",
"AWS_PROFILE",
"AWS_DEFAULT_PROFILE",
"KUBECONFIG",
] {
h.update(name.as_bytes());
h.update([0]);
if let Ok(value) = std::env::var(name) {
h.update(value.as_bytes());
}
h.update([0]);
}
format!("{:x}", h.finalize())
}
#[derive(Debug, Serialize)]
struct PublicEnvelope<'a> {
schema_version: u32,
protection_mode: &'a str,
contexts: &'a BTreeMap<String, String>,
ssh_remote: bool,
sudo_active: bool,
}
struct Status {
protection_mode: String,
contexts: BTreeMap<String, String>,
ssh_remote: bool,
sudo_active: bool,
}
pub fn run(short: bool, json: bool) -> i32 {
let status = match load_or_refresh() {
Ok(s) => s,
Err(_) => {
Status {
protection_mode: "off".into(),
contexts: BTreeMap::new(),
ssh_remote: false,
sudo_active: false,
}
}
};
if json {
let env = PublicEnvelope {
schema_version: 1,
protection_mode: &status.protection_mode,
contexts: &status.contexts,
ssh_remote: status.ssh_remote,
sudo_active: status.sudo_active,
};
match serde_json::to_string(&env) {
Ok(s) => {
println!("{s}");
0
}
Err(_) => {
println!(
"{{\"schema_version\":1,\"protection_mode\":\"off\",\"contexts\":{{}},\"ssh_remote\":false,\"sudo_active\":false}}"
);
0
}
}
} else if short {
println!("{}", format_short(&status));
0
} else {
println!("{}", format_long(&status));
0
}
}
fn format_short(s: &Status) -> String {
let mut out = format!(
"[tirith:{}]",
super::sanitize_for_human_output(&s.protection_mode, false)
);
for (k, v) in &s.contexts {
if v.is_empty() {
continue;
}
out.push_str(&format!(
"[{}:{}]",
super::sanitize_for_human_output(k, false),
super::sanitize_for_human_output(v, false)
));
}
if s.ssh_remote {
out.push_str("[ssh:remote]");
}
if s.sudo_active {
out.push_str("[sudo:active]");
}
out
}
fn format_long(s: &Status) -> String {
let mut parts = vec![format!(
"tirith: {}",
super::sanitize_for_human_output(&s.protection_mode, false)
)];
for (k, v) in &s.contexts {
if v.is_empty() {
continue;
}
parts.push(format!(
"{}: {}",
super::sanitize_for_human_output(k, false),
super::sanitize_for_human_output(v, false)
));
}
if s.ssh_remote {
parts.push("ssh: remote".into());
}
if s.sudo_active {
parts.push("sudo: session active".into());
}
parts.join("; ")
}
fn load_or_refresh() -> Result<Status, String> {
let cache_path = resolve_cache_path();
if let Some(path) = &cache_path {
if let Ok(bytes) = fs::read(path) {
if let Ok(env) = serde_json::from_slice::<CacheEnvelope>(&bytes) {
let now = unix_now();
if env.captured_at <= now
&& now - env.captured_at < CACHE_TTL_SECS
&& env.schema_version == 1
&& env.env_fingerprint == current_env_fingerprint()
{
return Ok(Status {
protection_mode: env.protection_mode,
contexts: env.contexts,
ssh_remote: env.ssh_remote,
sudo_active: env.sudo_active,
});
}
}
}
}
let status = refresh_status();
if let Some(path) = &cache_path {
let _ = write_cache(path, &status);
}
Ok(status)
}
fn refresh_status() -> Status {
let protection_mode = detect_protection_mode();
let ssh_remote = std::env::var("TIRITH_SSH_REMOTE")
.map(|v| {
let trimmed = v.trim();
!trimmed.is_empty() && trimmed != "0" && !trimmed.eq_ignore_ascii_case("false")
})
.unwrap_or(false);
let sudo_active = sudo_session::read_active_session().is_some();
let mut contexts = BTreeMap::new();
if let Ok(ctx) = context_detect::detect_single(Provider::Kube) {
contexts.insert(Provider::Kube.as_str().to_string(), ctx.context);
}
if let Ok(ctx) = context_detect::detect_single(Provider::Aws) {
contexts.insert(Provider::Aws.as_str().to_string(), ctx.context);
}
Status {
protection_mode,
contexts,
ssh_remote,
sudo_active,
}
}
fn detect_protection_mode() -> String {
protection_mode_from_status(std::env::var("TIRITH_STATUS").ok().as_deref())
}
pub(crate) fn protection_mode_from_status(status: Option<&str>) -> String {
match status {
Some("blocks") => "guarded".into(),
Some("warn-only") => "warn-only".into(),
Some("degraded") => "degraded".into(),
Some("off") | Some("") | None => "off".into(),
Some(other) => other.to_string(),
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum ProtectionHealth {
Guarded,
WarnOnly,
Degraded,
ConfiguredUnknown,
HookMissing,
Unknown,
}
impl ProtectionHealth {
pub(crate) fn classify(protection_mode: &str, hook_configured: bool) -> Self {
match protection_mode {
"guarded" if hook_configured => Self::Guarded,
"guarded" => Self::HookMissing,
"warn-only" => Self::WarnOnly,
"degraded" => Self::Degraded,
"off" if hook_configured => Self::ConfiguredUnknown,
"off" => Self::HookMissing,
_ => Self::Unknown,
}
}
pub(crate) fn exit_code(self) -> i32 {
match self {
Self::Guarded | Self::ConfiguredUnknown => 0,
_ => 1,
}
}
pub(crate) fn label(self) -> &'static str {
match self {
Self::Guarded => "guarded",
Self::WarnOnly => "warn-only",
Self::Degraded => "degraded",
Self::ConfiguredUnknown => "configured",
Self::HookMissing => "hook-missing",
Self::Unknown => "unknown",
}
}
}
#[cfg(test)]
pub(crate) fn protection_mode_for_test() -> String {
detect_protection_mode()
}
fn resolve_cache_path() -> Option<PathBuf> {
let uid = current_uid();
let file_name = format!("prompt-{uid}.cache");
if let Ok(rt_dir) = std::env::var("XDG_RUNTIME_DIR") {
let trimmed = rt_dir.trim();
if !trimmed.is_empty() {
let parent = PathBuf::from(trimmed).join("tirith");
if ensure_dir_0700(&parent).is_ok() {
return Some(parent.join(file_name));
}
}
}
if let Some(state) = tirith_core::policy::state_dir() {
if ensure_dir_0700(&state).is_ok() {
return Some(state.join(file_name));
}
}
None
}
fn current_uid() -> u32 {
#[cfg(unix)]
{
unsafe { libc::getuid() }
}
#[cfg(not(unix))]
{
0
}
}
fn ensure_dir_0700(dir: &std::path::Path) -> std::io::Result<()> {
fs::create_dir_all(dir)?;
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let perms = fs::Permissions::from_mode(0o700);
let _ = fs::set_permissions(dir, perms);
}
Ok(())
}
fn write_cache(path: &std::path::Path, status: &Status) -> std::io::Result<()> {
let envelope = CacheEnvelope {
schema_version: 1,
captured_at: unix_now(),
protection_mode: status.protection_mode.clone(),
contexts: status.contexts.clone(),
ssh_remote: status.ssh_remote,
sudo_active: status.sudo_active,
env_fingerprint: current_env_fingerprint(),
};
let body = serde_json::to_vec(&envelope).map_err(std::io::Error::other)?;
let parent = match path.parent() {
Some(p) => p,
None => return write_direct(path, &body),
};
let file_name = match path.file_name().and_then(|n| n.to_str()) {
Some(n) => n,
None => return write_direct(path, &body),
};
let tmp_name = format!(".{file_name}.{}.tmp", std::process::id());
let tmp_path = parent.join(tmp_name);
let mut opts = fs::OpenOptions::new();
opts.write(true).create(true).truncate(true);
#[cfg(unix)]
{
use std::os::unix::fs::OpenOptionsExt;
opts.mode(0o600);
}
let mut f = match opts.open(&tmp_path) {
Ok(f) => f,
Err(_) => {
return write_direct(path, &body);
}
};
use std::io::Write as _;
if let Err(e) = f.write_all(&body) {
let _ = fs::remove_file(&tmp_path);
return Err(e);
}
drop(f);
match fs::rename(&tmp_path, path) {
Ok(()) => Ok(()),
Err(e) => {
let _ = fs::remove_file(&tmp_path);
Err(e)
}
}
}
fn write_direct(path: &std::path::Path, body: &[u8]) -> std::io::Result<()> {
let mut opts = fs::OpenOptions::new();
opts.write(true).create(true).truncate(true);
#[cfg(unix)]
{
use std::os::unix::fs::OpenOptionsExt;
opts.mode(0o600);
}
let mut f = opts.open(path)?;
use std::io::Write as _;
f.write_all(body)
}
fn unix_now() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0)
}
#[cfg(test)]
fn render_short_for_test(
protection_mode: &str,
contexts: &[(&str, &str)],
ssh_remote: bool,
sudo_active: bool,
) -> String {
let mut map = BTreeMap::new();
for (k, v) in contexts {
map.insert((*k).to_string(), (*v).to_string());
}
format_short(&Status {
protection_mode: protection_mode.into(),
contexts: map,
ssh_remote,
sudo_active,
})
}
#[cfg(test)]
fn render_long_for_test(
protection_mode: &str,
contexts: &[(&str, &str)],
ssh_remote: bool,
sudo_active: bool,
) -> String {
let mut map = BTreeMap::new();
for (k, v) in contexts {
map.insert((*k).to_string(), (*v).to_string());
}
format_long(&Status {
protection_mode: protection_mode.into(),
contexts: map,
ssh_remote,
sudo_active,
})
}
#[cfg(test)]
mod tests {
use super::*;
use crate::cli::test_harness::{EnvGuard, ENV_LOCK};
#[test]
fn short_form_matches_spec_example() {
let line = render_short_for_test(
"guarded",
&[("aws", "prod"), ("kube", "payments-prod")],
false,
false,
);
assert_eq!(line, "[tirith:guarded][aws:prod][kube:payments-prod]");
}
#[test]
fn short_form_includes_ssh_and_sudo_when_active() {
let line = render_short_for_test("guarded", &[("aws", "prod")], true, true);
assert_eq!(line, "[tirith:guarded][aws:prod][ssh:remote][sudo:active]");
}
#[test]
fn short_form_no_contexts_is_just_tirith_segment() {
let line = render_short_for_test("off", &[], false, false);
assert_eq!(line, "[tirith:off]");
}
#[test]
fn short_form_skips_empty_context_values() {
let line = render_short_for_test("guarded", &[("kube", "")], false, false);
assert_eq!(line, "[tirith:guarded]");
}
#[test]
fn long_form_matches_spec_example() {
let line = render_long_for_test(
"guarded",
&[("aws", "prod"), ("kube", "payments-prod")],
false,
true,
);
assert_eq!(
line,
"tirith: guarded; aws: prod; kube: payments-prod; sudo: session active",
);
}
#[test]
fn long_form_no_contexts_only_tirith() {
let line = render_long_for_test("off", &[], false, false);
assert_eq!(line, "tirith: off");
}
#[test]
fn protection_mode_maps_known_values() {
let _lock = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
let _isolate = EnvGuard::remove("TIRITH_STATUS");
for (input, expected) in [
("blocks", "guarded"),
("warn-only", "warn-only"),
("degraded", "degraded"),
("off", "off"),
] {
let _value = EnvGuard::set("TIRITH_STATUS", std::path::Path::new(input));
assert_eq!(detect_protection_mode(), expected);
}
assert_eq!(detect_protection_mode(), "off");
}
#[test]
fn protection_mode_unknown_value_passes_through() {
let _lock = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
let _isolate = EnvGuard::remove("TIRITH_STATUS");
let _value = EnvGuard::set("TIRITH_STATUS", std::path::Path::new("futureValue"));
assert_eq!(detect_protection_mode(), "futureValue");
}
#[test]
fn protection_health_classify_and_exit_codes() {
let guarded = ProtectionHealth::classify("guarded", true);
assert_eq!(guarded, ProtectionHealth::Guarded);
assert_eq!(guarded.exit_code(), 0);
assert_eq!(guarded.label(), "guarded");
let warn_only = ProtectionHealth::classify("warn-only", true);
assert_eq!(warn_only, ProtectionHealth::WarnOnly);
assert_eq!(warn_only.exit_code(), 1);
let hook_missing = ProtectionHealth::classify("off", false);
assert_eq!(hook_missing, ProtectionHealth::HookMissing);
assert_eq!(hook_missing.exit_code(), 1);
assert_eq!(hook_missing.label(), "hook-missing");
let configured = ProtectionHealth::classify("off", true);
assert_eq!(configured, ProtectionHealth::ConfiguredUnknown);
assert_eq!(configured.exit_code(), 0);
assert_eq!(configured.label(), "configured");
let degraded = ProtectionHealth::classify("degraded", true);
assert_eq!(degraded, ProtectionHealth::Degraded);
assert_eq!(degraded.exit_code(), 1);
let unknown = ProtectionHealth::classify("futureValue", true);
assert_eq!(unknown, ProtectionHealth::Unknown);
assert_eq!(unknown.exit_code(), 1);
assert_eq!(unknown.label(), "unknown");
}
#[test]
fn cache_envelope_round_trips_via_serde() {
let env = CacheEnvelope {
schema_version: 1,
captured_at: 1_700_000_000,
protection_mode: "guarded".into(),
contexts: BTreeMap::from([
("aws".to_string(), "prod".to_string()),
("kube".to_string(), "payments-prod".to_string()),
]),
ssh_remote: true,
sudo_active: false,
env_fingerprint: String::new(),
};
let bytes = serde_json::to_vec(&env).unwrap();
let back: CacheEnvelope = serde_json::from_slice(&bytes).unwrap();
assert_eq!(back.protection_mode, "guarded");
assert_eq!(back.contexts.len(), 2);
assert!(back.ssh_remote);
assert!(!back.sudo_active);
}
}