use std::collections::HashMap;
use std::sync::Arc;
use crate::config::Config;
use crate::profile::Profile;
#[derive(Debug, Clone, Default)]
pub struct SessionVariables {
pub hostname: String,
pub username: String,
pub path: String,
pub job: Option<String>,
pub last_command: Option<String>,
pub profile_name: String,
pub tty: String,
pub columns: usize,
pub rows: usize,
pub bell_count: usize,
pub selection: Option<String>,
pub tmux_pane_title: Option<String>,
pub exit_code: Option<i32>,
pub current_command: Option<String>,
pub custom: HashMap<String, String>,
}
impl SessionVariables {
pub fn new() -> Self {
let hostname = hostname::get()
.map(|h| h.to_string_lossy().to_string())
.unwrap_or_else(|_| "unknown".to_string());
let username = std::env::var("USER")
.or_else(|_| std::env::var("USERNAME"))
.unwrap_or_else(|_| "unknown".to_string());
let path = std::env::current_dir()
.map(|p| p.to_string_lossy().to_string())
.unwrap_or_else(|_| "~".to_string());
Self {
hostname,
username,
path,
profile_name: "Default".to_string(),
tty: std::env::var("TTY").unwrap_or_default(),
columns: 80,
rows: 24,
..Default::default()
}
}
pub fn get(&self, name: &str) -> Option<String> {
match name {
"session.hostname" => Some(self.hostname.clone()),
"session.username" => Some(self.username.clone()),
"session.path" => Some(self.path.clone()),
"session.job" => self.job.clone(),
"session.last_command" => self.last_command.clone(),
"session.profile_name" => Some(self.profile_name.clone()),
"session.tty" => Some(self.tty.clone()),
"session.columns" => Some(self.columns.to_string()),
"session.rows" => Some(self.rows.to_string()),
"session.bell_count" => Some(self.bell_count.to_string()),
"session.selection" => self.selection.clone(),
"session.tmux_pane_title" => self.tmux_pane_title.clone(),
"session.exit_code" => self.exit_code.map(|c| c.to_string()),
"session.current_command" => self.current_command.clone(),
_ => {
if let Some(custom_name) = name.strip_prefix("session.") {
self.custom.get(custom_name).cloned()
} else {
None
}
}
}
}
pub fn set_path(&mut self, path: String) {
self.path = path;
}
pub fn set_dimensions(&mut self, cols: usize, rows: usize) {
self.columns = cols;
self.rows = rows;
}
pub fn increment_bell(&mut self) {
self.bell_count += 1;
}
pub fn set_custom(&mut self, name: &str, value: String) {
self.custom.insert(name.to_string(), value);
}
pub fn set_exit_code(&mut self, code: Option<i32>) {
self.exit_code = code;
}
pub fn set_current_command(&mut self, command: Option<String>) {
self.current_command = command;
}
}
#[derive(Clone)]
pub struct BadgeState {
pub enabled: bool,
pub format: String,
pub rendered_text: String,
pub color: [u8; 3],
pub alpha: f32,
pub font: String,
pub font_bold: bool,
pub top_margin: f32,
pub right_margin: f32,
pub max_width: f32,
pub max_height: f32,
pub variables: Arc<parking_lot::RwLock<SessionVariables>>,
dirty: bool,
}
impl BadgeState {
pub fn new(config: &Config) -> Self {
Self {
enabled: config.badge_enabled,
format: config.badge_format.clone(),
rendered_text: String::new(),
color: config.badge_color,
alpha: config.badge_color_alpha,
font: config.badge_font.clone(),
font_bold: config.badge_font_bold,
top_margin: config.badge_top_margin,
right_margin: config.badge_right_margin,
max_width: config.badge_max_width,
max_height: config.badge_max_height,
variables: Arc::new(parking_lot::RwLock::new(SessionVariables::new())),
dirty: true,
}
}
pub fn update_config(&mut self, config: &Config) {
let format_changed = self.format != config.badge_format;
self.enabled = config.badge_enabled;
self.format = config.badge_format.clone();
self.color = config.badge_color;
self.alpha = config.badge_color_alpha;
self.font = config.badge_font.clone();
self.font_bold = config.badge_font_bold;
self.top_margin = config.badge_top_margin;
self.right_margin = config.badge_right_margin;
self.max_width = config.badge_max_width;
self.max_height = config.badge_max_height;
if format_changed {
self.dirty = true;
}
}
pub fn set_format(&mut self, format: String) {
if self.format != format {
self.format = format;
self.dirty = true;
}
}
pub fn mark_dirty(&mut self) {
self.dirty = true;
}
pub fn is_dirty(&self) -> bool {
self.dirty
}
pub fn clear_dirty(&mut self) {
self.dirty = false;
}
pub fn interpolate(&mut self) {
let variables = self.variables.read();
self.rendered_text = interpolate_badge_format(&self.format, &variables);
self.dirty = false;
}
pub fn text(&self) -> &str {
&self.rendered_text
}
pub fn variables_mut(&self) -> parking_lot::RwLockWriteGuard<'_, SessionVariables> {
self.variables.write()
}
pub fn apply_profile_settings(&mut self, profile: &Profile) {
let mut changed = false;
if let Some(ref text) = profile.badge_text
&& self.format != *text
{
self.format = text.clone();
changed = true;
}
if let Some(color) = profile.badge_color {
self.color = color;
}
if let Some(alpha) = profile.badge_color_alpha {
self.alpha = alpha;
}
if let Some(ref font) = profile.badge_font {
self.font = font.clone();
}
if let Some(bold) = profile.badge_font_bold {
self.font_bold = bold;
}
if let Some(margin) = profile.badge_top_margin {
self.top_margin = margin;
}
if let Some(margin) = profile.badge_right_margin {
self.right_margin = margin;
}
if let Some(width) = profile.badge_max_width {
self.max_width = width;
}
if let Some(height) = profile.badge_max_height {
self.max_height = height;
}
if changed {
self.dirty = true;
}
}
}
pub fn interpolate_badge_format(format: &str, variables: &SessionVariables) -> String {
let mut result = String::with_capacity(format.len());
let mut chars = format.chars().peekable();
while let Some(ch) = chars.next() {
if ch == '\\' && chars.peek() == Some(&'(') {
chars.next();
let mut var_name = String::new();
for c in chars.by_ref() {
if c == ')' {
break;
}
var_name.push(c);
}
if let Some(value) = variables.get(&var_name) {
result.push_str(&value);
}
} else {
result.push(ch);
}
}
result
}
pub fn decode_badge_format(base64_format: &str) -> Option<String> {
use base64::Engine;
let engine = base64::engine::general_purpose::STANDARD;
let decoded = engine.decode(base64_format).ok()?;
let format = String::from_utf8(decoded).ok()?;
if format.contains("$(")
|| format.contains("`")
|| format.contains("eval")
|| format.contains("exec")
{
log::warn!(
"Rejecting badge format with suspicious content: {:?}",
format
);
return None;
}
Some(format)
}
pub struct BadgeInsets {
pub top: f32,
pub bottom: f32,
pub right: f32,
}
pub fn render_badge(
ctx: &egui::Context,
badge: &BadgeState,
window_width: f32,
_window_height: f32,
insets: &BadgeInsets,
) {
if !badge.enabled || badge.rendered_text.is_empty() {
return;
}
let color = egui::Color32::from_rgba_unmultiplied(
badge.color[0],
badge.color[1],
badge.color[2],
(badge.alpha * 255.0) as u8,
);
let font_id = egui::FontId::new(24.0, egui::FontFamily::Proportional);
let top_offset = badge.top_margin + insets.top;
egui::Area::new(egui::Id::new("badge_overlay"))
.fixed_pos(egui::pos2(0.0, top_offset))
.order(egui::Order::Foreground)
.interactable(false)
.show(ctx, |ui| {
let text = &badge.rendered_text;
let text_rect = ui.painter().text(
egui::pos2(0.0, 0.0),
egui::Align2::LEFT_TOP,
text,
font_id.clone(),
egui::Color32::TRANSPARENT, );
let x = window_width - text_rect.width() - badge.right_margin - insets.right;
let y = top_offset;
ui.painter().text(
egui::pos2(x, y),
egui::Align2::LEFT_TOP,
text,
font_id,
color,
);
});
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_interpolate_basic() {
let vars = SessionVariables {
hostname: "myhost".to_string(),
username: "testuser".to_string(),
..Default::default()
};
let result = interpolate_badge_format("\\(session.username)@\\(session.hostname)", &vars);
assert_eq!(result, "testuser@myhost");
}
#[test]
fn test_interpolate_missing_variable() {
let vars = SessionVariables::default();
let result = interpolate_badge_format("Hello \\(session.nonexistent) World", &vars);
assert_eq!(result, "Hello World");
}
#[test]
fn test_interpolate_no_variables() {
let vars = SessionVariables::default();
let result = interpolate_badge_format("Plain text", &vars);
assert_eq!(result, "Plain text");
}
#[test]
fn test_interpolate_escaped_backslash() {
let vars = SessionVariables::default();
let result = interpolate_badge_format("Path: C:\\Users", &vars);
assert_eq!(result, "Path: C:\\Users");
}
#[test]
fn test_decode_badge_format_valid() {
use base64::Engine;
let engine = base64::engine::general_purpose::STANDARD;
let encoded = engine.encode("Hello World");
let decoded = decode_badge_format(&encoded);
assert_eq!(decoded, Some("Hello World".to_string()));
}
#[test]
fn test_decode_badge_format_security_check() {
use base64::Engine;
let engine = base64::engine::general_purpose::STANDARD;
let encoded = engine.encode("$(whoami)");
assert!(decode_badge_format(&encoded).is_none());
let encoded = engine.encode("`whoami`");
assert!(decode_badge_format(&encoded).is_none());
let encoded = engine.encode("eval bad");
assert!(decode_badge_format(&encoded).is_none());
}
#[test]
fn test_session_variables_get() {
let vars = SessionVariables {
hostname: "test".to_string(),
columns: 120,
rows: 40,
..Default::default()
};
assert_eq!(vars.get("session.hostname"), Some("test".to_string()));
assert_eq!(vars.get("session.columns"), Some("120".to_string()));
assert_eq!(vars.get("session.rows"), Some("40".to_string()));
assert_eq!(vars.get("session.nonexistent"), None);
}
#[test]
fn test_session_variables_custom() {
let mut vars = SessionVariables::default();
vars.set_custom("myvar", "myvalue".to_string());
assert_eq!(vars.get("session.myvar"), Some("myvalue".to_string()));
}
#[test]
fn test_interpolate_exit_code() {
let vars = SessionVariables {
exit_code: Some(1),
..Default::default()
};
let result = interpolate_badge_format("Exit: \\(session.exit_code)", &vars);
assert_eq!(result, "Exit: 1");
}
#[test]
fn test_interpolate_current_command() {
let vars = SessionVariables {
current_command: Some("vim".to_string()),
..Default::default()
};
let result = interpolate_badge_format("Running: \\(session.current_command)", &vars);
assert_eq!(result, "Running: vim");
}
#[test]
fn test_interpolate_exit_code_none() {
let vars = SessionVariables {
exit_code: None,
..Default::default()
};
let result = interpolate_badge_format("Exit: \\(session.exit_code)", &vars);
assert_eq!(result, "Exit: ");
}
}