use std::collections::VecDeque;
use std::hash::Hash;
use std::marker::PhantomData;
#[cfg(feature = "pcap")]
use std::path::Path;
use crate::event::FlowEvent;
use crate::extractor::FlowExtractor;
#[cfg(feature = "pcap")]
use crate::pcap::PcapFlowSource;
use crate::session::{DatagramParser, SessionEvent, SessionParser};
use crate::tracker::FlowTrackerConfig;
#[cfg(feature = "session")]
use crate::datagram_driver::FlowDatagramDriver;
#[cfg(feature = "session")]
use crate::session_driver::FlowSessionDriver;
#[derive(Debug, Default, Clone, Copy)]
pub struct NoSessionParser;
impl SessionParser for NoSessionParser {
type Message = ();
fn feed_initiator(&mut self, _bytes: &[u8], _ts: crate::Timestamp) -> Vec<Self::Message> {
Vec::new()
}
fn feed_responder(&mut self, _bytes: &[u8], _ts: crate::Timestamp) -> Vec<Self::Message> {
Vec::new()
}
}
#[derive(Debug, Default, Clone, Copy)]
pub struct NoDatagramParser;
impl DatagramParser for NoDatagramParser {
type Message = ();
fn parse(
&mut self,
_payload: &[u8],
_side: crate::event::FlowSide,
_ts: crate::Timestamp,
) -> Vec<Self::Message> {
Vec::new()
}
}
pub struct PipelineBuilder<E, S = NoSessionParser, D = NoDatagramParser> {
extractor: E,
config: FlowTrackerConfig,
session_parser: Option<S>,
datagram_parser: Option<D>,
emit_anomalies: bool,
monotonic: bool,
}
impl<E> PipelineBuilder<E, NoSessionParser, NoDatagramParser>
where
E: FlowExtractor,
{
pub fn new(extractor: E) -> Self {
Self {
extractor,
config: FlowTrackerConfig::default(),
session_parser: None,
datagram_parser: None,
emit_anomalies: true,
monotonic: true,
}
}
}
impl<E, S, D> PipelineBuilder<E, S, D>
where
E: FlowExtractor,
{
pub fn config(mut self, c: FlowTrackerConfig) -> Self {
self.config = c;
self
}
pub fn emit_anomalies(mut self, on: bool) -> Self {
self.emit_anomalies = on;
self
}
pub fn monotonic(mut self, on: bool) -> Self {
self.monotonic = on;
self
}
pub fn session<P>(self, p: P) -> PipelineBuilder<E, P, D>
where
P: SessionParser + Clone + Send + 'static,
{
PipelineBuilder {
extractor: self.extractor,
config: self.config,
session_parser: Some(p),
datagram_parser: self.datagram_parser,
emit_anomalies: self.emit_anomalies,
monotonic: self.monotonic,
}
}
pub fn datagram<P>(self, p: P) -> PipelineBuilder<E, S, P>
where
P: DatagramParser + Clone + Send + 'static,
{
PipelineBuilder {
extractor: self.extractor,
config: self.config,
session_parser: self.session_parser,
datagram_parser: Some(p),
emit_anomalies: self.emit_anomalies,
monotonic: self.monotonic,
}
}
}
#[cfg(feature = "session")]
impl<E, S, D> PipelineBuilder<E, S, D>
where
E: FlowExtractor + Clone,
E::Key: Hash + Eq + Clone + Send + 'static,
S: SessionParser + Clone + Send + 'static,
D: DatagramParser + Clone + Send + 'static,
S: Default,
D: Default,
{
pub fn build(self) -> Pipeline<E, S, D> {
let session_parser = self.session_parser.unwrap_or_default();
let datagram_parser = self.datagram_parser.unwrap_or_default();
let session = FlowSessionDriver::with_config(
self.extractor.clone(),
session_parser.clone(),
self.config.clone(),
)
.with_emit_anomalies(self.emit_anomalies)
.with_monotonic_timestamps(self.monotonic);
let datagram = FlowDatagramDriver::with_config(
self.extractor.clone(),
datagram_parser.clone(),
self.config.clone(),
)
.with_emit_anomalies(self.emit_anomalies)
.with_monotonic_timestamps(self.monotonic);
Pipeline {
extractor: self.extractor,
config: self.config,
session_parser,
datagram_parser,
emit_anomalies: self.emit_anomalies,
monotonic: self.monotonic,
session: Some(session),
datagram: Some(datagram),
_markers: PhantomData,
}
}
}
#[derive(Debug, Clone)]
#[non_exhaustive]
pub enum Event<K, SM = (), DM = ()> {
Flow(FlowEvent<K>),
Tcp(SessionEvent<K, SM>),
Udp(SessionEvent<K, DM>),
}
impl<K, SM, DM> Event<K, SM, DM> {
pub fn kind(&self) -> EventKind {
match self {
Event::Flow(FlowEvent::Started { .. }) => EventKind::FlowStarted,
Event::Flow(FlowEvent::Packet { .. }) => EventKind::FlowPacket,
Event::Flow(FlowEvent::Established { .. }) => EventKind::FlowEstablished,
Event::Flow(FlowEvent::Ended { .. }) => EventKind::FlowEnded,
Event::Flow(FlowEvent::Tick { .. }) => EventKind::FlowTick,
Event::Flow(FlowEvent::FlowAnomaly { .. }) => EventKind::FlowAnomaly,
Event::Flow(FlowEvent::TrackerAnomaly { .. }) => EventKind::TrackerAnomaly,
Event::Flow(_) => EventKind::FlowOther,
Event::Tcp(SessionEvent::Application { .. }) => EventKind::TcpApplication,
Event::Tcp(SessionEvent::Closed { .. }) => EventKind::TcpClosed,
Event::Tcp(SessionEvent::Started { .. }) => EventKind::TcpFlowStarted,
Event::Tcp(SessionEvent::FlowTick { .. }) => EventKind::TcpTick,
Event::Tcp(_) => EventKind::TcpOther,
Event::Udp(SessionEvent::Application { .. }) => EventKind::UdpApplication,
Event::Udp(SessionEvent::Closed { .. }) => EventKind::UdpClosed,
Event::Udp(SessionEvent::Started { .. }) => EventKind::UdpFlowStarted,
Event::Udp(SessionEvent::FlowTick { .. }) => EventKind::UdpTick,
Event::Udp(_) => EventKind::UdpOther,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
#[non_exhaustive]
pub enum EventKind {
FlowStarted,
FlowPacket,
FlowEstablished,
FlowEnded,
FlowTick,
FlowAnomaly,
TrackerAnomaly,
FlowOther,
TcpApplication,
TcpClosed,
TcpFlowStarted,
TcpTick,
TcpOther,
UdpApplication,
UdpClosed,
UdpFlowStarted,
UdpTick,
UdpOther,
}
#[cfg(feature = "session")]
pub struct Pipeline<E, S = NoSessionParser, D = NoDatagramParser>
where
E: FlowExtractor,
S: SessionParser,
D: DatagramParser,
{
extractor: E,
config: FlowTrackerConfig,
session_parser: S,
datagram_parser: D,
emit_anomalies: bool,
monotonic: bool,
session: Option<FlowSessionDriver<E, S, ()>>,
datagram: Option<FlowDatagramDriver<E, D, ()>>,
_markers: PhantomData<()>,
}
impl<E> Pipeline<E, NoSessionParser, NoDatagramParser>
where
E: FlowExtractor,
{
pub fn builder(extractor: E) -> PipelineBuilder<E> {
PipelineBuilder::new(extractor)
}
}
#[cfg(feature = "session")]
impl<E, S, D> Pipeline<E, S, D>
where
E: FlowExtractor + Clone,
E::Key: Hash + Eq + Clone + Send + 'static,
S: SessionParser + Clone + Send + 'static,
D: DatagramParser + Clone + Send + 'static,
{
pub fn reset(&mut self) {
let session = FlowSessionDriver::with_config(
self.extractor.clone(),
self.session_parser.clone(),
self.config.clone(),
)
.with_emit_anomalies(self.emit_anomalies)
.with_monotonic_timestamps(self.monotonic);
let datagram = FlowDatagramDriver::with_config(
self.extractor.clone(),
self.datagram_parser.clone(),
self.config.clone(),
)
.with_emit_anomalies(self.emit_anomalies)
.with_monotonic_timestamps(self.monotonic);
self.session = Some(session);
self.datagram = Some(datagram);
}
}
#[cfg(all(feature = "pcap", feature = "session", feature = "reassembler"))]
impl<E, S, D> Pipeline<E, S, D>
where
E: FlowExtractor + Clone,
E::Key: Hash + Eq + Clone + Send + 'static,
S: SessionParser + Clone + Send + 'static,
D: DatagramParser + Clone + Send + 'static,
{
pub fn run_pcap(&mut self, path: impl AsRef<Path>) -> crate::Result<PipelineIter<'_, E, S, D>> {
let source = PcapFlowSource::open(path)?;
Ok(PipelineIter {
views: Box::new(source.views()),
pipeline: self,
pending: VecDeque::new(),
finished: false,
})
}
}
#[cfg(all(feature = "session", feature = "reassembler"))]
impl<E, S, D> Pipeline<E, S, D>
where
E: FlowExtractor + Clone,
E::Key: Hash + Eq + Clone + Send + 'static,
S: SessionParser + Clone + Send + 'static,
D: DatagramParser + Clone + Send + 'static,
{
pub fn run_iter<I>(&mut self, iter: I) -> PipelineIter<'_, E, S, D>
where
I: IntoIterator<Item = crate::OwnedPacketView> + 'static,
{
PipelineIter {
views: Box::new(iter.into_iter().map(Ok)),
pipeline: self,
pending: VecDeque::new(),
finished: false,
}
}
}
#[cfg(all(feature = "session", feature = "reassembler"))]
pub struct PipelineIter<'a, E, S, D>
where
E: FlowExtractor,
S: SessionParser,
D: DatagramParser,
{
views: Box<dyn Iterator<Item = crate::Result<crate::OwnedPacketView>> + 'a>,
pipeline: &'a mut Pipeline<E, S, D>,
pending: VecDeque<Event<E::Key, S::Message, D::Message>>,
finished: bool,
}
#[cfg(all(feature = "session", feature = "reassembler"))]
impl<'a, E, S, D> Iterator for PipelineIter<'a, E, S, D>
where
E: FlowExtractor + Clone,
E::Key: Hash + Eq + Clone + Send + 'static,
S: SessionParser + Clone + Send + 'static,
D: DatagramParser + Clone + Send + 'static,
{
type Item = crate::Result<Event<E::Key, S::Message, D::Message>>;
fn next(&mut self) -> Option<Self::Item> {
loop {
if let Some(ev) = self.pending.pop_front() {
return Some(Ok(ev));
}
if self.finished {
if let Some(session) = self.pipeline.session.as_mut() {
let drained = session.finish();
for ev in drained {
self.pending.push_back(Event::Tcp(ev));
}
}
if let Some(datagram) = self.pipeline.datagram.as_mut() {
let drained = datagram.finish();
for ev in drained {
self.pending.push_back(Event::Udp(ev));
}
}
return self.pending.pop_front().map(Ok);
}
match self.views.next() {
None => {
self.finished = true;
continue;
}
Some(Err(e)) => return Some(Err(e)),
Some(Ok(view)) => {
if let Some(session) = self.pipeline.session.as_mut() {
for ev in session.track(&view) {
self.pending.push_back(Event::Tcp(ev));
}
}
if let Some(datagram) = self.pipeline.datagram.as_mut() {
for ev in datagram.track(&view) {
self.pending.push_back(Event::Udp(ev));
}
}
}
}
}
}
}