mod erased;
mod event;
mod heuristic;
mod pipeline;
pub use event::Event;
pub use heuristic::{DEFAULT_PROBE_PACKETS, PROBE_BUFFER_CAP};
pub use pipeline::{Pipeline, PipelineBuilder, PipelineIter};
use std::hash::Hash;
use std::marker::PhantomData;
use std::time::Duration;
use crate::PacketView;
use crate::Timestamp;
use crate::dedup::Dedup;
use crate::detect::signatures::SignatureFn;
use crate::driver::FlowDriver;
use crate::event::FlowEvent;
use crate::extractor::{FlowExtractor, L4Proto, TcpInfo};
use crate::reassembler::NoopReassemblerFactory;
use crate::session::{DatagramParser, SessionParser};
use crate::tracker::{FlowTracker, FlowTrackerConfig};
use erased::{ConcreteDatagramSlot, ConcreteSlot, DriverSlot};
use heuristic::{HeuristicDatagramSlot, HeuristicSessionSlot};
type IdleTimeoutFn<K> = Box<dyn Fn(&K, Option<L4Proto>) -> Option<Duration> + Send + 'static>;
pub struct Driver<E, M>
where
E: FlowExtractor,
E::Key: Hash + Eq + Clone + Send + 'static,
M: Send + 'static,
{
central: FlowDriver<E, NoopReassemblerFactory, ()>,
extractor: E,
emit_packet_details: bool,
slots: Vec<Box<dyn DriverSlot<E::Key, M>>>,
_marker: PhantomData<M>,
}
impl<E, M> Driver<E, M>
where
E: FlowExtractor + Clone + Send + 'static,
E::Key: Hash + Eq + Clone + Send + 'static,
M: Send + 'static,
{
pub fn builder(extractor: E) -> DriverBuilder<E, M> {
DriverBuilder {
extractor,
config: FlowTrackerConfig::default(),
monotonic_timestamps: false,
emit_anomalies: false,
emit_packet_details: false,
dedup: None,
idle_timeout_fn: None,
slots: Vec::new(),
_marker: PhantomData,
}
}
pub fn track<'v>(&mut self, view: impl Into<PacketView<'v>>) -> Vec<Event<E::Key, M>> {
let view: PacketView<'v> = view.into();
let ts = view.timestamp;
let mut out: Vec<Event<E::Key, M>> = Vec::new();
let (tcp_for_packet, frame_for_packet): (Option<TcpInfo>, Option<Vec<u8>>) =
if self.emit_packet_details {
let tcp = self.extractor.extract(view).and_then(|e| e.tcp);
(tcp, Some(view.frame.to_vec()))
} else {
(None, None)
};
let mut tcp_slot = tcp_for_packet;
let mut frame_slot = frame_for_packet;
for flow_ev in self.central.track(view).into_iter() {
let (this_tcp, this_frame) = if matches!(flow_ev, FlowEvent::Packet { .. }) {
let pair = (tcp_slot, frame_slot.take());
tcp_slot = None; pair
} else {
(None, None)
};
out.extend(map_flow_event_with_details::<E::Key, M>(
flow_ev, this_tcp, this_frame,
));
}
for slot in &mut self.slots {
out.extend(slot.track(view, ts));
}
out
}
pub fn sweep(&mut self, now: Timestamp) -> Vec<Event<E::Key, M>> {
let mut out: Vec<Event<E::Key, M>> = Vec::new();
for flow_ev in self.central.sweep(now) {
out.extend(map_flow_event_with_details::<E::Key, M>(
flow_ev, None, None,
));
}
for slot in &mut self.slots {
out.extend(slot.sweep(now));
}
out
}
pub fn finish(&mut self) -> Vec<Event<E::Key, M>> {
let mut out: Vec<Event<E::Key, M>> = Vec::new();
for flow_ev in self.central.finish() {
out.extend(map_flow_event_with_details::<E::Key, M>(
flow_ev, None, None,
));
}
for slot in &mut self.slots {
out.extend(slot.finish());
}
out
}
pub fn tracker(&self) -> &FlowTracker<E, ()> {
self.central.tracker()
}
pub fn tracker_mut(&mut self) -> &mut FlowTracker<E, ()> {
self.central.tracker_mut()
}
}
pub struct DriverBuilder<E, M>
where
E: FlowExtractor,
M: Send + 'static,
{
extractor: E,
config: FlowTrackerConfig,
monotonic_timestamps: bool,
emit_anomalies: bool,
emit_packet_details: bool,
dedup: Option<Dedup>,
idle_timeout_fn: Option<IdleTimeoutFn<E::Key>>,
slots: Vec<Box<dyn DriverSlot<E::Key, M>>>,
_marker: PhantomData<M>,
}
impl<E, M> DriverBuilder<E, M>
where
E: FlowExtractor + Clone + Send + 'static,
E::Key: Hash + Eq + Clone + Send + 'static,
M: Send + 'static,
{
pub fn config(mut self, c: FlowTrackerConfig) -> Self {
self.config = c;
self
}
pub fn monotonic_timestamps(mut self, on: bool) -> Self {
self.monotonic_timestamps = on;
self
}
pub fn emit_packet_details(mut self, on: bool) -> Self {
self.emit_packet_details = on;
self
}
pub fn emit_anomalies(mut self, on: bool) -> Self {
self.emit_anomalies = on;
self
}
pub fn dedup(mut self, dedup: Dedup) -> Self {
self.dedup = Some(dedup);
self
}
pub fn idle_timeout_fn<F>(mut self, f: F) -> Self
where
F: Fn(&E::Key, Option<L4Proto>) -> Option<Duration> + Send + 'static,
{
self.idle_timeout_fn = Some(Box::new(f));
self
}
pub fn session_on_ports<P, I, F>(mut self, parser: P, ports: I, lift: F) -> Self
where
P: SessionParser + Clone + Send + 'static,
P::Message: Send + 'static,
I: IntoIterator<Item = u16>,
F: Fn(P::Message) -> M + Send + 'static,
{
let port_set: smallvec::SmallVec<[u16; 4]> = ports.into_iter().collect();
let slot = ConcreteSlot::new(
self.extractor.clone(),
parser,
self.config.clone(),
Some(port_set),
self.monotonic_timestamps,
lift,
);
self.slots.push(Box::new(slot));
self
}
pub fn session_broadcast<P, F>(mut self, parser: P, lift: F) -> Self
where
P: SessionParser + Clone + Send + 'static,
P::Message: Send + 'static,
F: Fn(P::Message) -> M + Send + 'static,
{
let slot = ConcreteSlot::new(
self.extractor.clone(),
parser,
self.config.clone(),
None,
self.monotonic_timestamps,
lift,
);
self.slots.push(Box::new(slot));
self
}
pub fn datagram_on_ports<D, I, F>(mut self, parser: D, ports: I, lift: F) -> Self
where
D: DatagramParser + Clone + Send + 'static,
D::Message: Send + 'static,
I: IntoIterator<Item = u16>,
F: Fn(D::Message) -> M + Send + 'static,
{
let port_set: smallvec::SmallVec<[u16; 4]> = ports.into_iter().collect();
let slot = ConcreteDatagramSlot::new(
self.extractor.clone(),
parser,
self.config.clone(),
Some(port_set),
self.monotonic_timestamps,
lift,
);
self.slots.push(Box::new(slot));
self
}
pub fn datagram_broadcast<D, F>(mut self, parser: D, lift: F) -> Self
where
D: DatagramParser + Clone + Send + 'static,
D::Message: Send + 'static,
F: Fn(D::Message) -> M + Send + 'static,
{
let slot = ConcreteDatagramSlot::new(
self.extractor.clone(),
parser,
self.config.clone(),
None,
self.monotonic_timestamps,
lift,
);
self.slots.push(Box::new(slot));
self
}
pub fn session_heuristic<P, F>(self, parser: P, signature: SignatureFn, lift: F) -> Self
where
P: SessionParser + Clone + Send + 'static,
P::Message: Send + 'static,
F: Fn(P::Message) -> M + Send + 'static,
{
self.session_heuristic_with_budget(parser, signature, DEFAULT_PROBE_PACKETS, lift)
}
pub fn session_heuristic_with_budget<P, F>(
mut self,
parser: P,
signature: SignatureFn,
max_probe_packets: u8,
lift: F,
) -> Self
where
P: SessionParser + Clone + Send + 'static,
P::Message: Send + 'static,
F: Fn(P::Message) -> M + Send + 'static,
{
let slot = HeuristicSessionSlot::new(
self.extractor.clone(),
parser,
self.config.clone(),
signature,
max_probe_packets,
self.monotonic_timestamps,
lift,
);
self.slots.push(Box::new(slot));
self
}
pub fn datagram_heuristic<D, F>(self, parser: D, signature: SignatureFn, lift: F) -> Self
where
D: DatagramParser + Clone + Send + 'static,
D::Message: Send + 'static,
F: Fn(D::Message) -> M + Send + 'static,
{
self.datagram_heuristic_with_budget(parser, signature, DEFAULT_PROBE_PACKETS, lift)
}
pub fn datagram_heuristic_with_budget<D, F>(
mut self,
parser: D,
signature: SignatureFn,
max_probe_packets: u8,
lift: F,
) -> Self
where
D: DatagramParser + Clone + Send + 'static,
D::Message: Send + 'static,
F: Fn(D::Message) -> M + Send + 'static,
{
let slot = HeuristicDatagramSlot::new(
self.extractor.clone(),
parser,
self.config.clone(),
signature,
max_probe_packets,
self.monotonic_timestamps,
lift,
);
self.slots.push(Box::new(slot));
self
}
pub fn build(self) -> Driver<E, M> {
let mut central =
FlowDriver::with_config(self.extractor.clone(), NoopReassemblerFactory, self.config)
.with_emit_anomalies(self.emit_anomalies)
.with_monotonic_timestamps(self.monotonic_timestamps);
if let Some(d) = self.dedup {
central = central.with_dedup(d);
}
if let Some(f) = self.idle_timeout_fn {
central
.tracker_mut()
.set_idle_timeout_fn(move |k, l4| f(k, l4));
}
Driver {
central,
extractor: self.extractor,
emit_packet_details: self.emit_packet_details,
slots: self.slots,
_marker: PhantomData,
}
}
}
fn map_flow_event_with_details<K, M>(
ev: FlowEvent<K>,
tcp: Option<TcpInfo>,
frame: Option<Vec<u8>>,
) -> Option<Event<K, M>> {
match ev {
FlowEvent::Started { key, ts, l4, .. } => Some(Event::FlowStarted { key, ts, l4 }),
FlowEvent::Established { key, ts, l4 } => Some(Event::FlowEstablished { key, ts, l4 }),
FlowEvent::Packet { key, side, len, ts } => Some(Event::FlowPacket {
key,
side,
len,
ts,
tcp,
frame,
}),
FlowEvent::Ended {
key,
reason,
stats,
history,
l4,
} => {
let ts = stats.last_seen;
Some(Event::FlowEnded {
key,
reason,
stats,
history,
l4,
ts,
})
}
FlowEvent::Tick { key, stats, ts } => Some(Event::FlowTick { key, stats, ts }),
FlowEvent::FlowAnomaly { key, kind, ts } => Some(Event::FlowAnomaly { key, kind, ts }),
FlowEvent::TrackerAnomaly { kind, ts } => Some(Event::TrackerAnomaly { kind, ts }),
FlowEvent::StateChange { .. } => None,
}
}