1use std::collections::HashMap;
8use std::sync::Arc;
9
10use crate::config::Config;
11use crate::profile::Profile;
12
13#[derive(Debug, Clone, Default)]
15pub struct SessionVariables {
16 pub hostname: String,
18 pub username: String,
20 pub path: String,
22 pub job: Option<String>,
24 pub last_command: Option<String>,
26 pub profile_name: String,
28 pub tty: String,
30 pub columns: usize,
32 pub rows: usize,
34 pub bell_count: usize,
36 pub selection: Option<String>,
38 pub tmux_pane_title: Option<String>,
40 pub exit_code: Option<i32>,
42 pub current_command: Option<String>,
44 pub custom: HashMap<String, String>,
46}
47
48impl SessionVariables {
49 pub fn new() -> Self {
51 let hostname = hostname::get()
52 .map(|h| h.to_string_lossy().to_string())
53 .unwrap_or_else(|_| "unknown".to_string());
54
55 let username = std::env::var("USER")
56 .or_else(|_| std::env::var("USERNAME"))
57 .unwrap_or_else(|_| "unknown".to_string());
58
59 let path = std::env::current_dir()
60 .map(|p| p.to_string_lossy().to_string())
61 .unwrap_or_else(|_| "~".to_string());
62
63 Self {
64 hostname,
65 username,
66 path,
67 profile_name: "Default".to_string(),
68 tty: std::env::var("TTY").unwrap_or_default(),
69 columns: 80,
70 rows: 24,
71 ..Default::default()
72 }
73 }
74
75 pub fn get(&self, name: &str) -> Option<String> {
77 match name {
78 "session.hostname" => Some(self.hostname.clone()),
79 "session.username" => Some(self.username.clone()),
80 "session.path" => Some(self.path.clone()),
81 "session.job" => self.job.clone(),
82 "session.last_command" => self.last_command.clone(),
83 "session.profile_name" => Some(self.profile_name.clone()),
84 "session.tty" => Some(self.tty.clone()),
85 "session.columns" => Some(self.columns.to_string()),
86 "session.rows" => Some(self.rows.to_string()),
87 "session.bell_count" => Some(self.bell_count.to_string()),
88 "session.selection" => self.selection.clone(),
89 "session.tmux_pane_title" => self.tmux_pane_title.clone(),
90 "session.exit_code" => self.exit_code.map(|c| c.to_string()),
91 "session.current_command" => self.current_command.clone(),
92 _ => {
93 if let Some(custom_name) = name.strip_prefix("session.") {
95 self.custom.get(custom_name).cloned()
96 } else {
97 None
98 }
99 }
100 }
101 }
102
103 pub fn set_path(&mut self, path: String) {
105 self.path = path;
106 }
107
108 pub fn set_dimensions(&mut self, cols: usize, rows: usize) {
110 self.columns = cols;
111 self.rows = rows;
112 }
113
114 pub fn increment_bell(&mut self) {
116 self.bell_count += 1;
117 }
118
119 pub fn set_custom(&mut self, name: &str, value: String) {
121 self.custom.insert(name.to_string(), value);
122 }
123
124 pub fn set_exit_code(&mut self, code: Option<i32>) {
126 self.exit_code = code;
127 }
128
129 pub fn set_current_command(&mut self, command: Option<String>) {
131 self.current_command = command;
132 }
133}
134
135#[derive(Clone)]
137pub struct BadgeState {
138 pub enabled: bool,
140 pub format: String,
142 pub rendered_text: String,
144 pub color: [u8; 3],
146 pub alpha: f32,
148 pub font: String,
150 pub font_bold: bool,
152 pub top_margin: f32,
154 pub right_margin: f32,
156 pub max_width: f32,
158 pub max_height: f32,
160 pub variables: Arc<parking_lot::RwLock<SessionVariables>>,
162 dirty: bool,
164}
165
166impl BadgeState {
167 pub fn new(config: &Config) -> Self {
169 Self {
170 enabled: config.badge_enabled,
171 format: config.badge_format.clone(),
172 rendered_text: String::new(),
173 color: config.badge_color,
174 alpha: config.badge_color_alpha,
175 font: config.badge_font.clone(),
176 font_bold: config.badge_font_bold,
177 top_margin: config.badge_top_margin,
178 right_margin: config.badge_right_margin,
179 max_width: config.badge_max_width,
180 max_height: config.badge_max_height,
181 variables: Arc::new(parking_lot::RwLock::new(SessionVariables::new())),
182 dirty: true,
183 }
184 }
185
186 pub fn update_config(&mut self, config: &Config) {
188 let format_changed = self.format != config.badge_format;
189
190 self.enabled = config.badge_enabled;
191 self.format = config.badge_format.clone();
192 self.color = config.badge_color;
193 self.alpha = config.badge_color_alpha;
194 self.font = config.badge_font.clone();
195 self.font_bold = config.badge_font_bold;
196 self.top_margin = config.badge_top_margin;
197 self.right_margin = config.badge_right_margin;
198 self.max_width = config.badge_max_width;
199 self.max_height = config.badge_max_height;
200
201 if format_changed {
202 self.dirty = true;
203 }
204 }
205
206 pub fn set_format(&mut self, format: String) {
208 if self.format != format {
209 self.format = format;
210 self.dirty = true;
211 }
212 }
213
214 pub fn mark_dirty(&mut self) {
216 self.dirty = true;
217 }
218
219 pub fn is_dirty(&self) -> bool {
221 self.dirty
222 }
223
224 pub fn clear_dirty(&mut self) {
226 self.dirty = false;
227 }
228
229 pub fn interpolate(&mut self) {
231 let variables = self.variables.read();
232 self.rendered_text = interpolate_badge_format(&self.format, &variables);
233 self.dirty = false;
234 }
235
236 pub fn text(&self) -> &str {
238 &self.rendered_text
239 }
240
241 pub fn variables_mut(&self) -> parking_lot::RwLockWriteGuard<'_, SessionVariables> {
243 self.variables.write()
244 }
245
246 pub fn apply_profile_settings(&mut self, profile: &Profile) {
251 let mut changed = false;
252
253 if let Some(ref text) = profile.badge_text
255 && self.format != *text
256 {
257 self.format = text.clone();
258 changed = true;
259 }
260
261 if let Some(color) = profile.badge_color {
263 self.color = color;
264 }
265
266 if let Some(alpha) = profile.badge_color_alpha {
268 self.alpha = alpha;
269 }
270
271 if let Some(ref font) = profile.badge_font {
273 self.font = font.clone();
274 }
275
276 if let Some(bold) = profile.badge_font_bold {
278 self.font_bold = bold;
279 }
280
281 if let Some(margin) = profile.badge_top_margin {
283 self.top_margin = margin;
284 }
285
286 if let Some(margin) = profile.badge_right_margin {
288 self.right_margin = margin;
289 }
290
291 if let Some(width) = profile.badge_max_width {
293 self.max_width = width;
294 }
295
296 if let Some(height) = profile.badge_max_height {
298 self.max_height = height;
299 }
300
301 if changed {
302 self.dirty = true;
303 }
304 }
305}
306
307pub fn interpolate_badge_format(format: &str, variables: &SessionVariables) -> String {
326 let mut result = String::with_capacity(format.len());
327 let mut chars = format.chars().peekable();
328
329 while let Some(ch) = chars.next() {
330 if ch == '\\' && chars.peek() == Some(&'(') {
331 chars.next();
333
334 let mut var_name = String::new();
336 for c in chars.by_ref() {
337 if c == ')' {
338 break;
339 }
340 var_name.push(c);
341 }
342
343 if let Some(value) = variables.get(&var_name) {
345 result.push_str(&value);
346 }
347 } else {
349 result.push(ch);
350 }
351 }
352
353 result
354}
355
356pub fn decode_badge_format(base64_format: &str) -> Option<String> {
360 use base64::Engine;
361 let engine = base64::engine::general_purpose::STANDARD;
362
363 let decoded = engine.decode(base64_format).ok()?;
364 let format = String::from_utf8(decoded).ok()?;
365
366 if format.contains("$(")
369 || format.contains("`")
370 || format.contains("eval")
371 || format.contains("exec")
372 {
373 log::warn!(
374 "Rejecting badge format with suspicious content: {:?}",
375 format
376 );
377 return None;
378 }
379
380 Some(format)
381}
382
383pub fn render_badge(
388 ctx: &egui::Context,
389 badge: &BadgeState,
390 window_width: f32,
391 _window_height: f32,
392) {
393 if !badge.enabled || badge.rendered_text.is_empty() {
394 return;
395 }
396
397 let color = egui::Color32::from_rgba_unmultiplied(
399 badge.color[0],
400 badge.color[1],
401 badge.color[2],
402 (badge.alpha * 255.0) as u8,
403 );
404
405 let font_id = egui::FontId::new(24.0, egui::FontFamily::Proportional);
407
408 egui::Area::new(egui::Id::new("badge_overlay"))
410 .fixed_pos(egui::pos2(0.0, badge.top_margin))
411 .order(egui::Order::Foreground)
412 .interactable(false)
413 .show(ctx, |ui| {
414 let text = &badge.rendered_text;
416
417 let text_rect = ui.painter().text(
419 egui::pos2(0.0, 0.0),
420 egui::Align2::LEFT_TOP,
421 text,
422 font_id.clone(),
423 egui::Color32::TRANSPARENT, );
425
426 let x = window_width - text_rect.width() - badge.right_margin;
428 let y = badge.top_margin;
429
430 ui.painter().text(
432 egui::pos2(x, y),
433 egui::Align2::LEFT_TOP,
434 text,
435 font_id,
436 color,
437 );
438 });
439}
440
441#[cfg(test)]
442mod tests {
443 use super::*;
444
445 #[test]
446 fn test_interpolate_basic() {
447 let vars = SessionVariables {
448 hostname: "myhost".to_string(),
449 username: "testuser".to_string(),
450 ..Default::default()
451 };
452
453 let result = interpolate_badge_format("\\(session.username)@\\(session.hostname)", &vars);
454 assert_eq!(result, "testuser@myhost");
455 }
456
457 #[test]
458 fn test_interpolate_missing_variable() {
459 let vars = SessionVariables::default();
460 let result = interpolate_badge_format("Hello \\(session.nonexistent) World", &vars);
461 assert_eq!(result, "Hello World");
462 }
463
464 #[test]
465 fn test_interpolate_no_variables() {
466 let vars = SessionVariables::default();
467 let result = interpolate_badge_format("Plain text", &vars);
468 assert_eq!(result, "Plain text");
469 }
470
471 #[test]
472 fn test_interpolate_escaped_backslash() {
473 let vars = SessionVariables::default();
474 let result = interpolate_badge_format("Path: C:\\Users", &vars);
476 assert_eq!(result, "Path: C:\\Users");
477 }
478
479 #[test]
480 fn test_decode_badge_format_valid() {
481 use base64::Engine;
482 let engine = base64::engine::general_purpose::STANDARD;
483 let encoded = engine.encode("Hello World");
484 let decoded = decode_badge_format(&encoded);
485 assert_eq!(decoded, Some("Hello World".to_string()));
486 }
487
488 #[test]
489 fn test_decode_badge_format_security_check() {
490 use base64::Engine;
491 let engine = base64::engine::general_purpose::STANDARD;
492
493 let encoded = engine.encode("$(whoami)");
495 assert!(decode_badge_format(&encoded).is_none());
496
497 let encoded = engine.encode("`whoami`");
499 assert!(decode_badge_format(&encoded).is_none());
500
501 let encoded = engine.encode("eval bad");
503 assert!(decode_badge_format(&encoded).is_none());
504 }
505
506 #[test]
507 fn test_session_variables_get() {
508 let vars = SessionVariables {
509 hostname: "test".to_string(),
510 columns: 120,
511 rows: 40,
512 ..Default::default()
513 };
514
515 assert_eq!(vars.get("session.hostname"), Some("test".to_string()));
516 assert_eq!(vars.get("session.columns"), Some("120".to_string()));
517 assert_eq!(vars.get("session.rows"), Some("40".to_string()));
518 assert_eq!(vars.get("session.nonexistent"), None);
519 }
520
521 #[test]
522 fn test_session_variables_custom() {
523 let mut vars = SessionVariables::default();
524 vars.set_custom("myvar", "myvalue".to_string());
525
526 assert_eq!(vars.get("session.myvar"), Some("myvalue".to_string()));
527 }
528
529 #[test]
530 fn test_interpolate_exit_code() {
531 let vars = SessionVariables {
532 exit_code: Some(1),
533 ..Default::default()
534 };
535
536 let result = interpolate_badge_format("Exit: \\(session.exit_code)", &vars);
537 assert_eq!(result, "Exit: 1");
538 }
539
540 #[test]
541 fn test_interpolate_current_command() {
542 let vars = SessionVariables {
543 current_command: Some("vim".to_string()),
544 ..Default::default()
545 };
546
547 let result = interpolate_badge_format("Running: \\(session.current_command)", &vars);
548 assert_eq!(result, "Running: vim");
549 }
550
551 #[test]
552 fn test_interpolate_exit_code_none() {
553 let vars = SessionVariables {
554 exit_code: None,
555 ..Default::default()
556 };
557
558 let result = interpolate_badge_format("Exit: \\(session.exit_code)", &vars);
559 assert_eq!(result, "Exit: ");
560 }
561}