use crate::tui::theme::colors;
use ratatui::{
prelude::*,
widgets::{Block, Scrollbar, ScrollbarOrientation, ScrollbarState, StatefulWidget, Widget},
};
use std::collections::HashSet;
#[derive(Debug, Clone)]
pub enum TreeNode {
Group {
id: String,
label: String,
children: Vec<Self>,
item_count: usize,
vuln_count: usize,
},
Component {
id: String,
name: String,
version: Option<String>,
vuln_count: usize,
max_severity: Option<String>,
component_type: Option<String>,
ecosystem: Option<String>,
is_bookmarked: bool,
},
}
impl TreeNode {
pub(crate) fn id(&self) -> &str {
match self {
Self::Group { id, .. } | Self::Component { id, .. } => id,
}
}
pub(crate) fn label(&self) -> String {
match self {
Self::Group {
label, item_count, ..
} => format!("{label} ({item_count})"),
Self::Component {
name,
version,
ecosystem,
is_bookmarked,
..
} => {
let display_name = extract_display_name(name);
let mut result = if let Some(v) = version {
format!("{display_name}@{v}")
} else {
display_name
};
if let Some(eco) = ecosystem
&& eco != "Unknown"
{
use std::fmt::Write;
let _ = write!(result, " [{eco}]");
}
if *is_bookmarked {
result = format!("\u{2605} {result}");
}
result
}
}
}
pub(crate) const fn vuln_count(&self) -> usize {
match self {
Self::Group { vuln_count, .. } | Self::Component { vuln_count, .. } => *vuln_count,
}
}
pub(crate) fn max_severity(&self) -> Option<&str> {
match self {
Self::Component { max_severity, .. } => max_severity.as_deref(),
Self::Group { .. } => None,
}
}
pub(crate) const fn is_group(&self) -> bool {
matches!(self, Self::Group { .. })
}
pub(crate) fn children(&self) -> Option<&[Self]> {
match self {
Self::Group { children, .. } => Some(children),
Self::Component { .. } => None,
}
}
}
pub fn extract_display_name(name: &str) -> String {
if !name.contains('/') && !name.starts_with('.') && name.len() <= 40 {
return name.to_string();
}
if let Some(filename) = name.rsplit('/').next() {
let clean = filename
.trim_end_matches(".squashfs")
.trim_end_matches(".squ")
.trim_end_matches(".img")
.trim_end_matches(".bin")
.trim_end_matches(".unknown")
.trim_end_matches(".crt")
.trim_end_matches(".so")
.trim_end_matches(".a")
.trim_end_matches(".elf32");
if is_hash_like(clean) {
let parts: Vec<&str> = name.split('/').collect();
if parts.len() >= 2 {
for part in parts.iter().rev().skip(1) {
if !part.is_empty()
&& !part.starts_with('.')
&& !is_hash_like(part)
&& part.len() > 2
{
return format!("{part}/{filename}");
}
}
}
return filename.to_string();
}
return clean.to_string();
}
name.to_string()
}
fn is_hash_like(name: &str) -> bool {
if name.len() < 8 {
return false;
}
let clean = name.replace(['-', '_'], "");
clean.chars().all(|c| c.is_ascii_hexdigit())
|| (clean.chars().filter(char::is_ascii_digit).count() > clean.len() / 2)
}
pub fn detect_component_type(name: &str) -> &'static str {
let lower = name.to_lowercase();
let ext = std::path::Path::new(&lower)
.extension()
.and_then(|e| e.to_str())
.unwrap_or("");
if matches!(ext, "so") || lower.contains(".so.") {
return "lib";
}
if matches!(ext, "a") {
return "lib";
}
if matches!(ext, "crt" | "pem" | "key") {
return "cert";
}
if matches!(ext, "img" | "bin" | "elf" | "elf32") {
return "bin";
}
if matches!(ext, "squashfs" | "squ") {
return "fs";
}
if matches!(ext, "unknown") {
return "unk";
}
if lower.contains("lib") {
return "lib";
}
"file"
}
pub fn detect_component_label(name: &str) -> Option<&'static str> {
match detect_component_type(name) {
"lib" => Some("Shared library"),
"bin" => Some("Binary / ELF"),
"cert" => Some("Certificate"),
"fs" => Some("Filesystem image"),
"unk" => Some("Unknown format"),
_ => None,
}
}
#[derive(Debug, Clone, Default)]
pub struct TreeState {
pub selected: usize,
pub expanded: HashSet<String>,
pub offset: usize,
pub visible_count: usize,
}
impl TreeState {
pub(crate) fn new() -> Self {
Self::default()
}
pub(crate) fn toggle_expand(&mut self, node_id: &str) {
if self.expanded.contains(node_id) {
self.expanded.remove(node_id);
} else {
self.expanded.insert(node_id.to_string());
}
}
pub(crate) fn expand(&mut self, node_id: &str) {
self.expanded.insert(node_id.to_string());
}
pub(crate) fn collapse(&mut self, node_id: &str) {
self.expanded.remove(node_id);
}
pub(crate) fn is_expanded(&self, node_id: &str) -> bool {
self.expanded.contains(node_id)
}
pub(crate) const fn select_next(&mut self) {
if self.visible_count > 0 && self.selected < self.visible_count - 1 {
self.selected += 1;
}
}
pub(crate) const fn select_prev(&mut self) {
if self.selected > 0 {
self.selected -= 1;
}
}
pub(crate) const fn select_first(&mut self) {
self.selected = 0;
}
pub(crate) const fn select_last(&mut self) {
if self.visible_count > 0 {
self.selected = self.visible_count - 1;
}
}
}
#[derive(Debug, Clone)]
pub struct FlattenedItem {
pub label: String,
pub depth: usize,
pub is_group: bool,
pub is_expanded: bool,
pub is_last_sibling: bool,
pub vuln_count: usize,
pub ancestors_last: Vec<bool>,
pub max_severity: Option<String>,
}
pub struct Tree<'a> {
roots: &'a [TreeNode],
block: Option<Block<'a>>,
highlight_style: Style,
highlight_symbol: &'a str,
group_style: Style,
search_query: String,
}
impl<'a> Tree<'a> {
pub(crate) fn new(roots: &'a [TreeNode]) -> Self {
let scheme = colors();
Self {
roots,
block: None,
highlight_style: Style::default()
.bg(scheme.selection)
.add_modifier(Modifier::BOLD),
highlight_symbol: "â–¶ ",
group_style: Style::default().fg(scheme.primary).bold(),
search_query: String::new(),
}
}
pub(crate) fn block(mut self, block: Block<'a>) -> Self {
self.block = Some(block);
self
}
pub(crate) const fn highlight_style(mut self, style: Style) -> Self {
self.highlight_style = style;
self
}
pub(crate) fn search_query(mut self, query: &str) -> Self {
self.search_query = query.to_lowercase();
self
}
fn flatten(&self, state: &TreeState) -> Vec<FlattenedItem> {
let mut items = Vec::new();
self.flatten_nodes(self.roots, 0, state, &mut items, &[]);
items
}
#[allow(clippy::only_used_in_recursion)]
fn flatten_nodes(
&self,
nodes: &[TreeNode],
depth: usize,
state: &TreeState,
items: &mut Vec<FlattenedItem>,
ancestors_last: &[bool],
) {
for (i, node) in nodes.iter().enumerate() {
let is_last = i == nodes.len() - 1;
let is_expanded = state.is_expanded(node.id());
let mut current_ancestors = ancestors_last.to_vec();
current_ancestors.push(is_last);
items.push(FlattenedItem {
label: node.label(),
depth,
is_group: node.is_group(),
is_expanded,
is_last_sibling: is_last,
vuln_count: node.vuln_count(),
ancestors_last: current_ancestors.clone(),
max_severity: node.max_severity().map(std::string::ToString::to_string),
});
if is_expanded && let Some(children) = node.children() {
self.flatten_nodes(children, depth + 1, state, items, ¤t_ancestors);
}
}
}
}
impl StatefulWidget for Tree<'_> {
type State = TreeState;
fn render(self, area: Rect, buf: &mut Buffer, state: &mut Self::State) {
let inner_area = self.block.as_ref().map_or(area, |b| {
let inner = b.inner(area);
b.clone().render(area, buf);
inner
});
if inner_area.width < 4 || inner_area.height < 1 {
return;
}
let items = self.flatten(state);
let area = inner_area;
state.visible_count = items.len();
let visible_height = area.height as usize;
if state.selected >= state.offset + visible_height {
state.offset = state.selected - visible_height + 1;
} else if state.selected < state.offset {
state.offset = state.selected;
}
for (i, item) in items
.iter()
.skip(state.offset)
.take(visible_height)
.enumerate()
{
let y = area.y + i as u16;
let is_selected = state.offset + i == state.selected;
let mut prefix = String::new();
for (depth, is_last) in item.ancestors_last.iter().take(item.depth).enumerate() {
if depth < item.depth {
if *is_last {
prefix.push_str(" ");
} else {
prefix.push_str("│ ");
}
}
}
if item.depth > 0 {
if item.is_last_sibling {
prefix.push_str("└─ ");
} else {
prefix.push_str("├─ ");
}
}
let expand_indicator = if item.is_group {
if item.is_expanded { "â–¼ " } else { "â–¶ " }
} else {
"· "
};
let mut x = area.x;
let scheme = colors();
if is_selected {
let symbol = self.highlight_symbol;
for ch in symbol.chars() {
if x < area.x + area.width {
if let Some(cell) = buf.cell_mut((x, y)) {
cell.set_char(ch)
.set_style(Style::default().fg(scheme.accent));
}
x += 1;
}
}
} else {
x += self.highlight_symbol.len() as u16;
}
for ch in prefix.chars() {
if x < area.x + area.width {
if let Some(cell) = buf.cell_mut((x, y)) {
cell.set_char(ch)
.set_style(Style::default().fg(scheme.text_muted));
}
x += 1;
}
}
let indicator_style = if item.is_group {
Style::default().fg(scheme.accent)
} else {
Style::default().fg(scheme.muted)
};
for ch in expand_indicator.chars() {
if x < area.x + area.width {
if let Some(cell) = buf.cell_mut((x, y)) {
cell.set_char(ch).set_style(indicator_style);
}
x += 1;
}
}
let is_search_match = !self.search_query.is_empty()
&& item.label.to_lowercase().contains(&self.search_query);
let depth_color = if item.is_group {
scheme.primary
} else {
match item.depth {
0 | 1 => scheme.text,
2 => Color::Rgb(180, 180, 180),
_ => scheme.text_muted,
}
};
let label_style = if is_selected {
self.highlight_style
} else if is_search_match {
Style::default().fg(scheme.accent).bold()
} else if item.is_group {
self.group_style
} else if item.depth == 0 {
Style::default().fg(depth_color).bold()
} else {
Style::default().fg(depth_color)
};
let vuln_badge_width: u16 = if item.vuln_count > 0 {
3 + item.vuln_count.to_string().len() as u16 } else {
0
};
let remaining = (area.x + area.width).saturating_sub(x + vuln_badge_width) as usize;
let display_label = if item.label.len() > remaining && remaining > 3 {
let max_chars = remaining.saturating_sub(3);
let truncated: String = item.label.chars().take(max_chars).collect();
format!("{truncated}...")
} else {
item.label.clone()
};
for ch in display_label.chars() {
if x < area.x + area.width {
if let Some(cell) = buf.cell_mut((x, y)) {
cell.set_char(ch).set_style(label_style);
}
x += 1;
}
}
if item.vuln_count > 0 {
let (sev_char, sev_color) =
item.max_severity
.as_ref()
.map_or(('!', scheme.warning), |sev| {
match sev.to_lowercase().as_str() {
"critical" => ('C', scheme.critical),
"high" => ('H', scheme.high),
"medium" => ('M', scheme.medium),
"low" => ('L', scheme.low),
_ => ('!', scheme.warning),
}
});
if x < area.x + area.width {
if let Some(cell) = buf.cell_mut((x, y)) {
cell.set_char(' ');
}
x += 1;
}
let badge_style = Style::default()
.fg(scheme.badge_fg_dark)
.bg(sev_color)
.bold();
if x < area.x + area.width {
if let Some(cell) = buf.cell_mut((x, y)) {
cell.set_char(sev_char).set_style(badge_style);
}
x += 1;
}
if x < area.x + area.width {
x += 1;
}
let count_text = format!("{}", item.vuln_count);
let count_style = Style::default().fg(sev_color).bold();
for ch in count_text.chars() {
if x < area.x + area.width {
if let Some(cell) = buf.cell_mut((x, y)) {
cell.set_char(ch).set_style(count_style);
}
x += 1;
}
}
}
if is_selected {
while x < area.x + area.width {
if let Some(cell) = buf.cell_mut((x, y)) {
cell.set_style(self.highlight_style);
}
x += 1;
}
}
}
if items.len() > visible_height {
let scheme = colors();
let scrollbar = Scrollbar::new(ScrollbarOrientation::VerticalRight)
.thumb_style(Style::default().fg(scheme.accent))
.track_style(Style::default().fg(scheme.muted));
let mut scrollbar_state = ScrollbarState::new(items.len()).position(state.selected);
scrollbar.render(area, buf, &mut scrollbar_state);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_tree_state() {
let mut state = TreeState::new();
assert!(!state.is_expanded("test"));
state.toggle_expand("test");
assert!(state.is_expanded("test"));
state.toggle_expand("test");
assert!(!state.is_expanded("test"));
}
#[test]
fn test_tree_node() {
let node = TreeNode::Component {
id: "comp-1".to_string(),
name: "lodash".to_string(),
version: Some("4.17.21".to_string()),
vuln_count: 2,
max_severity: Some("high".to_string()),
component_type: Some("lib".to_string()),
ecosystem: Some("npm".to_string()),
is_bookmarked: false,
};
assert_eq!(node.label(), "lodash@4.17.21 [npm]");
assert_eq!(node.vuln_count(), 2);
assert_eq!(node.max_severity(), Some("high"));
assert!(!node.is_group());
}
#[test]
fn test_tree_node_bookmarked() {
let node = TreeNode::Component {
id: "comp-1".to_string(),
name: "lodash".to_string(),
version: Some("4.17.21".to_string()),
vuln_count: 0,
max_severity: None,
component_type: None,
ecosystem: None,
is_bookmarked: true,
};
assert_eq!(node.label(), "\u{2605} lodash@4.17.21");
}
#[test]
fn test_tree_node_unknown_ecosystem_hidden() {
let node = TreeNode::Component {
id: "comp-1".to_string(),
name: "lodash".to_string(),
version: Some("4.17.21".to_string()),
vuln_count: 0,
max_severity: None,
component_type: None,
ecosystem: Some("Unknown".to_string()),
is_bookmarked: false,
};
assert_eq!(node.label(), "lodash@4.17.21");
}
#[test]
fn test_extract_display_name() {
assert_eq!(
extract_display_name("./6488064-48136192.squashfs_v4_le_extract/SMASH/ShowProperty"),
"ShowProperty"
);
assert_eq!(extract_display_name("lodash"), "lodash");
assert_eq!(extract_display_name("openssl-1.1.1"), "openssl-1.1.1");
let hash_result = extract_display_name("./6488064-48136192.squashfs");
assert!(hash_result.len() <= 30);
}
#[test]
fn test_detect_component_type() {
assert_eq!(detect_component_type("libssl.so"), "lib");
assert_eq!(detect_component_type("libcrypto.so.1.1"), "lib");
assert_eq!(detect_component_type("server.crt"), "cert");
assert_eq!(detect_component_type("firmware.img"), "bin");
assert_eq!(detect_component_type("rootfs.squashfs"), "fs");
assert_eq!(detect_component_type("random.unknown"), "unk");
}
}