diff options
Diffstat (limited to 'examples/rs-simple/src/main.rs')
| -rw-r--r-- | examples/rs-simple/src/main.rs | 590 |
1 files changed, 590 insertions, 0 deletions
diff --git a/examples/rs-simple/src/main.rs b/examples/rs-simple/src/main.rs new file mode 100644 index 000000000..647c351c2 --- /dev/null +++ b/examples/rs-simple/src/main.rs @@ -0,0 +1,590 @@ +use bytes::BytesMut; +use crossterm::{ + cursor, + event::{self, KeyCode, KeyEvent}, + ExecutableCommand, + terminal::{self, ClearType}, +}; +use moka::{future::Cache, Expiry}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use std::{ + fmt, + hash::{Hash, Hasher}, + io::self, + sync::Arc, + time::{Duration, Instant, SystemTime, UNIX_EPOCH}, +}; +use tokio::io::AsyncReadExt; +use tokio::sync::mpsc; +use tokio::sync::Mutex; +use tokio::sync::MutexGuard; +use tokio::net::TcpStream; +use tui::{ + backend::CrosstermBackend, + layout::{Layout, Constraint, Direction}, + style::{Style, Color, Modifier}, + Terminal, + widgets::{Block, Borders, List, ListItem, Row, Table, TableState}, +}; + +#[derive(Debug)] +enum ParseError { + Protocol(), + Json(), + Schema(), +} + +impl From<serde_json::Error> for ParseError { + fn from(_: serde_json::Error) -> Self { + ParseError::Json() + } +} + +#[derive(Serialize, Deserialize, Debug, PartialEq)] +#[serde(rename_all = "lowercase")] +enum EventName { + Invalid, New, End, Idle, Update, Analyse, + Guessed, Detected, + #[serde(rename = "detection-update")] + DetectionUpdate, + #[serde(rename = "not-detected")] + NotDetected, +} + +#[derive(Serialize, Deserialize, Copy, Clone, Debug)] +#[serde(rename_all = "lowercase")] +enum State { + Unknown, Info, Finished, +} + +#[derive(Serialize, Deserialize, Debug)] +struct FlowEvent { + #[serde(rename = "flow_event_name")] + name: EventName, + #[serde(rename = "flow_id")] + id: u64, + #[serde(rename = "flow_state")] + state: State, + #[serde(rename = "flow_first_seen")] + first_seen: u64, + #[serde(rename = "flow_src_last_pkt_time")] + src_last_pkt_time: u64, + #[serde(rename = "flow_dst_last_pkt_time")] + dst_last_pkt_time: u64, + #[serde(rename = "flow_idle_time")] + idle_time: u64, + #[serde(rename = "flow_src_packets_processed")] + src_packets_processed: u64, + #[serde(rename = "flow_dst_packets_processed")] + dst_packets_processed: u64, + #[serde(rename = "flow_src_tot_l4_payload_len")] + src_tot_l4_payload_len: u64, + #[serde(rename = "flow_dst_tot_l4_payload_len")] + dst_tot_l4_payload_len: u64, +} + +#[derive(Serialize, Deserialize, Debug)] +struct PacketEvent { + pkt_datalink: u16, + pkt_caplen: u64, + pkt_len: u64, + pkt_l4_len: u64, +} + +#[derive(Debug)] +enum EventType { + Flow(FlowEvent), + Packet(PacketEvent), + Other(), +} + +#[derive(Default)] +struct Stats { + ui_updates: u64, + flow_count: u64, + parse_errors: u64, + events: u64, + flow_events: u64, + packet_events: u64, + total_caplen: u64, + total_len: u64, + total_l4_len: u64, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum Expiration { + FlowIdleTime(u64), +} + +struct FlowExpiry; + +#[derive(Clone, Eq, Default, Debug)] +struct FlowKey { + id: u64, +} + +#[derive(Clone, Debug)] +struct FlowValue { + state: State, + total_src_packets: u64, + total_dst_packets: u64, + total_src_bytes: u64, + total_dst_bytes: u64, + first_seen: std::time::SystemTime, + last_seen: std::time::SystemTime, + timeout_in: std::time::SystemTime, +} + +impl Default for State { + fn default() -> State { + State::Unknown + } +} + +impl fmt::Display for State { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + State::Unknown => write!(f, "N/A"), + State::Info => write!(f, "Info"), + State::Finished => write!(f, "Finished"), + } + } +} + +impl Expiration { + fn as_duration(&self) -> Option<Duration> { + match self { + Expiration::FlowIdleTime(value) => Some(Duration::from_micros(*value)), + } + } +} + +impl fmt::Display for Expiration { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self.as_duration() { + Some(duration) => { + let secs = duration.as_secs(); + write!(f, "{} s", secs) + } + None => write!(f, "N/A"), + } + } +} + +impl Expiry<FlowKey, (Expiration, FlowValue)> for FlowExpiry { + fn expire_after_create( + &self, + _key: &FlowKey, + value: &(Expiration, FlowValue), + _current_time: Instant, + ) -> Option<Duration> { + value.0.as_duration() + } +} + +impl Hash for FlowKey { + fn hash<H: Hasher>(&self, state: &mut H) { + self.id.hash(state) + } +} + +impl PartialEq for FlowKey { + fn eq(&self, other: &Self) -> bool { + self.id == other.id + } +} + +#[tokio::main] +async fn main() { + let server_address = "127.0.0.1:7000"; + + let mut stream = match TcpStream::connect(server_address).await { + Ok(stream) => stream, + Err(e) => { + eprintln!("Connection to {} failed: {}", server_address, e); + return; + } + }; + if let Err(e) = terminal::enable_raw_mode() { + eprintln!("Could not enable terminal raw mode: {}", e); + return; + } + let mut stdout = io::stdout(); + if let Err(e) = stdout.execute(terminal::Clear(ClearType::All)) { + eprintln!("Could not clear your terminal: {}", e); + return; + } + if let Err(e) = stdout.execute(cursor::Hide) { + eprintln!("Could not hide your cursor: {}", e); + return; + } + let backend = CrosstermBackend::new(stdout); + let mut terminal = Terminal::new(backend); + + let mut buffer = BytesMut::with_capacity(33792usize); + let (tx, mut rx): (mpsc::Sender<String>, mpsc::Receiver<String>) = mpsc::channel(1024); + let data = Arc::new(Mutex::new(Stats::default())); + let data_tx = Arc::clone(&data); + let data_rx = Arc::clone(&data); + let flow_cache: Arc<Cache<FlowKey, (Expiration, FlowValue)>> = Arc::new(Cache::builder() + .expire_after(FlowExpiry) + .build()); + let flow_cache_rx = Arc::clone(&flow_cache); + + tokio::spawn(async move { + while let Some(msg) = rx.recv().await { + match parse_json(&msg) { + Ok(message) => { + let mut data_lock = data_tx.lock().await; + data_lock.events += 1; + update_stats(&message, &mut data_lock, &flow_cache).await; + } + Err(_message) => { + let mut data_lock = data_tx.lock().await; + data_lock.parse_errors += 1; + } + } + } + }); + tokio::spawn(async move { + loop { + let n = match stream.read_buf(&mut buffer).await { + Ok(len) => len, + Err(_) => { + continue; // Versuche es erneut, wenn ein Fehler auftritt + } + }; + if n == 0 { + break; + } + + while let Some(message) = parse_message(&mut buffer) { + match tx.send(message).await { + Ok(_) => (), + Err(_) => return + } + } + } + }); + + let mut table_state = TableState::default(); + + loop { + let flows: Vec<(FlowKey, (Expiration, FlowValue))> = flow_cache_rx.iter().map(|(k, v)| (k.as_ref().clone(), v.clone())).collect(); + let mut table_selected = match table_state.selected() { + Some(table_index) => { + if flows.len() > 0 && table_index >= flows.len() { + flows.len() - 1 + } else { + table_index + } + } + None => 0, + }; + + match read_keypress() { + Some(KeyCode::Esc) => break, + Some(KeyCode::Char('q')) => break, + Some(KeyCode::Up) => { + table_selected = match table_selected { + _ if flows.len() == 0 => 0, + i if i == 0 => flows.len() - 1, + i => i - 1, + }; + }, + Some(KeyCode::Down) => { + table_selected = match table_selected { + i if flows.len() == 0 || i >= flows.len() - 1 => 0, + i => i + 1, + }; + }, + Some(KeyCode::PageUp) => { + table_selected = match table_selected { + _ if flows.len() == 0 => 0, + i if i == 0 => flows.len() - 1, + i if i < 10 => 0, + i => i - 10, + }; + }, + Some(KeyCode::PageDown) => { + table_selected = match table_selected { + i if flows.len() == 0 || i >= flows.len() - 1 => 0, + i if flows.len() < 10 || i >= flows.len() - 10 => flows.len() - 1, + i => i + 10, + }; + }, + Some(KeyCode::Home) => { + table_selected = 0; + }, + Some(KeyCode::End) => { + table_selected = match table_selected { + _ if flows.len() == 0 => 0, + _ => flows.len() - 1, + }; + }, + Some(_) => (), + None => () + }; + + let mut data_lock = data_rx.lock().await; + data_lock.ui_updates += 1; + draw_ui(terminal.as_mut().unwrap(), &mut table_state, table_selected, &data_lock, &flows); + } + + if let Err(e) = terminal.unwrap().backend_mut().execute(cursor::Show) { + eprintln!("Could not show your cursor: {}", e); + return; + } + let mut stdout = io::stdout(); + if let Err(e) = stdout.execute(terminal::Clear(ClearType::All)) { + eprintln!("Could not clear your terminal: {}", e); + return; + } + if let Err(e) = terminal::disable_raw_mode() { + eprintln!("Could not disable raw mode: {}", e); + return; + } + println!("\nDone."); +} + +fn read_keypress() -> Option<KeyCode> { + if event::poll(Duration::from_millis(500)).unwrap() { + if let event::Event::Key(KeyEvent { code, .. }) = event::read().unwrap() { + return Some(code); + } + } + + None +} + +fn parse_message(buffer: &mut BytesMut) -> Option<String> { + if let Some(pos) = buffer.iter().position(|&b| b == b'\n') { + let message = buffer.split_to(pos + 1); + return Some(String::from_utf8_lossy(&message).to_string()); + } + + None +} + +fn parse_json(data: &str) -> Result<EventType, ParseError> { + let first_non_digit = data.find(|c: char| !c.is_ascii_digit()).unwrap_or(0); + let length_str = &data[0..first_non_digit]; + let length: usize = length_str.parse().unwrap_or(0); + if length == 0 { + return Err(ParseError::Protocol()); + } + + let json_str = &data[first_non_digit..first_non_digit + length]; + let value: Value = serde_json::from_str(json_str).map_err(|_| ParseError::Json()).unwrap(); + if value.get("flow_event_name").is_some() { + let flow_event: FlowEvent = serde_json::from_value(value)?; + return Ok(EventType::Flow(flow_event)); + } else if value.get("packet_event_name").is_some() { + let packet_event: PacketEvent = serde_json::from_value(value)?; + return Ok(EventType::Packet(packet_event)); + } else if value.get("daemon_event_name").is_some() || + value.get("error_event_name").is_some() { + return Ok(EventType::Other()); + } + + Err(ParseError::Schema()) +} + +async fn update_stats(event: &EventType, stats: &mut MutexGuard<'_, Stats>, cache: &Cache<FlowKey, (Expiration, FlowValue)>) { + match &event { + EventType::Flow(flow_event) => { + stats.flow_events += 1; + stats.flow_count = cache.entry_count(); + let key = FlowKey { id: flow_event.id }; + + if flow_event.name == EventName::End || + flow_event.name == EventName::Idle + { + cache.remove(&key).await; + return; + } + + let first_seen_seconds = flow_event.first_seen / 1_000_000; + let first_seen_nanos = (flow_event.first_seen % 1_000_000) * 1_000; + let first_seen_epoch = std::time::Duration::new(first_seen_seconds, first_seen_nanos as u32); + let first_seen_system = UNIX_EPOCH + first_seen_epoch; + + let last_seen = std::cmp::max(flow_event.src_last_pkt_time, + flow_event.dst_last_pkt_time); + let last_seen_seconds = last_seen / 1_000_000; + let last_seen_nanos = (last_seen % 1_000_000) * 1_000; + let last_seen_epoch = std::time::Duration::new(last_seen_seconds, last_seen_nanos as u32); + let last_seen_system = UNIX_EPOCH + last_seen_epoch; + + let timeout_seconds = (last_seen + flow_event.idle_time) / 1_000_000; + let timeout_nanos = ((last_seen + flow_event.idle_time) % 1_000_000) * 1_000; + let timeout_epoch = std::time::Duration::new(timeout_seconds, timeout_nanos as u32); + let timeout_system = UNIX_EPOCH + timeout_epoch; + + let value = FlowValue { + state: flow_event.state, + total_src_packets: flow_event.src_packets_processed, + total_dst_packets: flow_event.dst_packets_processed, + total_src_bytes: flow_event.src_tot_l4_payload_len, + total_dst_bytes: flow_event.dst_tot_l4_payload_len, + first_seen: first_seen_system, + last_seen: last_seen_system, + timeout_in: timeout_system, + }; + cache.insert(key, (Expiration::FlowIdleTime(flow_event.idle_time), value)).await; + } + EventType::Packet(packet_event) => { + stats.packet_events += 1; + stats.total_caplen += packet_event.pkt_caplen; + stats.total_len += packet_event.pkt_len; + stats.total_l4_len += packet_event.pkt_l4_len; + } + EventType::Other() => {} + } +} + +fn format_bytes(bytes: u64) -> String { + const KB: u64 = 1024; + const MB: u64 = KB * 1024; + const GB: u64 = MB * 1024; + + if bytes >= GB { + format!("{} GB", bytes / GB) + } else if bytes >= MB { + format!("{} MB", bytes / MB) + } else if bytes >= KB { + format!("{} kB", bytes / KB) + } else { + format!("{} B", bytes) + } +} + +fn draw_ui<B: tui::backend::Backend>(terminal: &mut Terminal<B>, table_state: &mut TableState, table_selected: usize, data: &MutexGuard<Stats>, flows: &Vec<(FlowKey, (Expiration, FlowValue))>) { + let general_items = vec![ + ListItem::new("TUI Updates..: ".to_owned() + &data.ui_updates.to_string()), + ListItem::new("Flows Cached.: ".to_owned() + &data.flow_count.to_string()), + ListItem::new("Total Events.: ".to_owned() + &data.events.to_string()), + ListItem::new("Parse Errors.: ".to_owned() + &data.parse_errors.to_string()), + ListItem::new("Flow Events..: ".to_owned() + &data.flow_events.to_string()), + ListItem::new("Packet Events: ".to_owned() + &data.packet_events.to_string()), + ]; + let packet_items = vec![ + ListItem::new("Total Capture Length: ".to_owned() + &format_bytes(data.total_caplen)), + ListItem::new("Total Length........: ".to_owned() + &format_bytes(data.total_len)), + ListItem::new("Total L4 Length.....: ".to_owned() + &format_bytes(data.total_l4_len)), + ]; + let table_rows: Vec<Row> = flows + .into_iter() + .map(|(key, (_exp, val))| { + let first_seen_display = match val.first_seen.elapsed() { + Ok(elapsed) => { + match elapsed.as_secs() { + t if t > (3_600 * 24) => format!("{} d ago", t / (3_600 * 24)), + t if t > 3_600 => format!("{} h ago", t / 3_600), + t if t > 60 => format!("{} min ago", t / 60), + t if t > 0 => format!("{} s ago", t), + t if t == 0 => "< 1 s ago".to_string(), + t => format!("INVALID: {}", t), + } + } + Err(err) => format!("ERROR: {}", err) + }; + + let last_seen_display = match val.last_seen.elapsed() { + Ok(elapsed) => { + match elapsed.as_secs() { + t if t > (3_600 * 24) => format!("{} d ago", t / (3_600 * 24)), + t if t > 3_600 => format!("{} h ago", t / 3_600), + t if t > 60 => format!("{} min ago", t / 60), + t if t > 0 => format!("{} s ago", t), + t if t == 0 => "< 1 s ago".to_string(), + t => format!("INVALID: {}", t), + } + } + Err(_err) => "ERROR".to_string() + }; + + let timeout_display = match val.timeout_in.duration_since(SystemTime::now()) { + Ok(elapsed) => { + match elapsed.as_secs() { + t if t > (3_600 * 24) => format!("in {} d", t / (3_600 * 24)), + t if t > 3_600 => format!("in {} h", t / 3_600), + t if t > 60 => format!("in {} min", t / 60), + t if t > 0 => format!("in {} s", t), + t if t == 0 => "in < 1 s".to_string(), + t => format!("INVALID: {}", t), + } + } + Err(_err) => "EXPIRED".to_string() + }; + + Row::new(vec![ + key.id.to_string(), + val.state.to_string(), + first_seen_display, + last_seen_display, + timeout_display, + (val.total_src_packets + val.total_dst_packets).to_string(), + format_bytes(val.total_src_bytes + val.total_dst_bytes), + ]) + }) + .collect(); + + terminal.draw(|f| { + let size = f.size(); + + let chunks = Layout::default() + .direction(Direction::Vertical) + .constraints( + [ + Constraint::Percentage(18), + Constraint::Percentage(82), + ].as_ref() + ) + .split(size); + + let top_chunks = Layout::default() + .direction(Direction::Horizontal) + .constraints( + [ + Constraint::Percentage(50), + Constraint::Percentage(50), + ].as_ref() + ) + .split(chunks[0]); + + let table_selected_abs = match table_selected { + _ if flows.len() == 0 => 0, + i => i + 1, + }; + let table = Table::new(table_rows) + .header(Row::new(vec!["Flow ID", "State", "First Seen", "Last Seen", "Timeout", "Total Packets", "Total Bytes"]) + .style(Style::default().fg(Color::Yellow).add_modifier(Modifier::BOLD))) + .block(Block::default().title("Flow Table (selected: ".to_string() + + &table_selected_abs.to_string() + + "): " + + &flows.len().to_string() + + " item(s)").borders(Borders::ALL)) + .highlight_style(Style::default().bg(Color::Blue)) + .widths(&[ + Constraint::Length(20), + Constraint::Length(20), + Constraint::Length(20), + Constraint::Length(20), + Constraint::Length(20), + Constraint::Length(20), + Constraint::Length(20), + ]); + + let general_list = List::new(general_items) + .block(Block::default().title("General").borders(Borders::ALL)); + let packet_list = List::new(packet_items) + .block(Block::default().title("Packet Events").borders(Borders::ALL)); + + table_state.select(Some(table_selected)); + f.render_widget(general_list, top_chunks[0]); + f.render_widget(packet_list, top_chunks[1]); + f.render_stateful_widget(table, chunks[1], table_state); + }).unwrap(); +} |